dgx-spark-playbooks/nvidia/speculative-decoding/README.md

Speculative Decoding

Learn how to setup speculative decoding for fast inference on Spark

Table of Contents

• Overview
• How to run inference with speculative decoding
  • Step 1. Run Eagle3 with GPT-OSS 120B
  • Step 2. Test the Eagle3 setup
  • Step 1. Run Draft-Target Speculative Decoding
  • Step 2. Test the Draft-Target setup
  • Troubleshooting
  • Cleanup
  • Next Steps

---

Overview

Basic idea

Speculative decoding speeds up text generation by using a small, fast model to draft several tokens ahead, then having the larger model quickly verify or adjust them. This way, the big model doesn’t need to predict every token step-by-step, reducing latency while keeping output quality.

What you'll accomplish

You'll explore two different speculative decoding approaches using TensorRT-LLM on NVIDIA Spark:

1. Eagle3 with GPT-OSS 120B - Advanced speculative decoding using Eagle3 draft models
2. Traditional Draft-Target - Classic speculative decoding with smaller model pairs (coming soon)

These examples demonstrate how to accelerate large language model inference while maintaining output quality.

What to know before starting

• Experience with Docker and containerized applications
• Understanding of speculative decoding concepts (Eagle3 vs traditional draft-target)
• Familiarity with TensorRT-LLM serving and API endpoints
• Knowledge of GPU memory management for large language models

Prerequisites

• NVIDIA Spark device with sufficient GPU memory available (80GB+ recommended for GPT-OSS 120B)
• Docker with GPU support enabled

  docker run --gpus all nvcr.io/nvidia/tensorrt-llm/release:spark-single-gpu-dev nvidia-smi
  

• Access to NVIDIA's internal container registry (for Eagle3 example)
• HuggingFace authentication configured (if needed for model downloads)

  huggingface-cli login
  

• Network connectivity for model downloads

Time & risk

Duration: 10-20 minutes for Eagle3 setup, additional time for model downloads (varies by network speed)

Risks: GPU memory exhaustion with large models, container registry access issues, network timeouts during downloads

Rollback: Stop Docker containers and optionally clean up downloaded model cache

How to run inference with speculative decoding

Example 1: Eagle3 Speculative Decoding with GPT-OSS 120B

Eagle3 is an advanced speculative decoding technique that uses a specialized draft model to accelerate inference of large language models.

Step 1. Run Eagle3 with GPT-OSS 120B

Execute the following command to download models and run Eagle3 speculative decoding:

docker run \
  -v $HOME/.cache/huggingface/:/root/.cache/huggingface/ \
  --rm -it --ulimit memlock=-1 --ulimit stack=67108864 \
  --gpus=all --ipc=host --network host nvcr.io/nvidia/tensorrt-llm/release:spark-single-gpu-dev \
  bash -c '
    hf download openai/gpt-oss-120b && \
    hf download nvidia/gpt-oss-120b-Eagle3 \
        --local-dir /opt/gpt-oss-120b-Eagle3/ && \
    cat > /tmp/extra-llm-api-config.yml <<EOF
enable_attention_dp: false
disable_overlap_scheduler: true
enable_autotuner: false
cuda_graph_config:
    max_batch_size: 1
speculative_config:
    decoding_type: Eagle
    max_draft_len: 4
    speculative_model_dir: /opt/gpt-oss-120b-Eagle3/

kv_cache_config:
    enable_block_reuse: false
EOF
    trtllm-serve openai/gpt-oss-120b \
      --backend pytorch --tp_size 1 \
      --max_batch_size 1 \
      --kv_cache_free_gpu_memory_fraction 0.95 \
      --extra_llm_api_options /tmp/extra-llm-api-config.yml'

Step 2. Test the Eagle3 setup

Once the server is running, you can test it with curl commands:

## Test completion endpoint
curl -X POST http://localhost:8000/v1/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "openai/gpt-oss-120b",
    "prompt": "The future of AI is",
    "max_tokens": 100,
    "temperature": 0.7
  }'

Example 2: Traditional Draft-Target Speculative Decoding

This example demonstrates traditional speculative decoding using a smaller draft model to accelerate a larger target model.

Step 1. Run Draft-Target Speculative Decoding

Execute the following command to set up and run traditional speculative decoding:

docker run \
  -v $HOME/.cache/huggingface/:/root/.cache/huggingface/ \
  --rm -it --ulimit memlock=-1 --ulimit stack=67108864 \
  --gpus=all --ipc=host --network host nvcr.io/nvidia/tensorrt-llm/release:spark-single-gpu-dev \
  bash -c "
#    # Download models
    hf download nvidia/Llama-3.3-70B-Instruct-FP4 && \
    hf download nvidia/Llama-3.1-8B-Instruct-FP4 \
    --local-dir /opt/Llama-3.1-8B-Instruct-FP4/ && \

#    # Create configuration file
    cat <<EOF > extra-llm-api-config.yml
print_iter_log: false
disable_overlap_scheduler: true
speculative_config:
  decoding_type: DraftTarget
  max_draft_len: 4
  speculative_model_dir: /opt/Llama-3.1-8B-Instruct-FP4/
kv_cache_config:
  enable_block_reuse: false
EOF

#    # Start TensorRT-LLM server
    trtllm-serve nvidia/Llama-3.3-70B-Instruct-FP4 \
      --backend pytorch --tp_size 1 \
      --max_batch_size 1 \
      --kv_cache_free_gpu_memory_fraction 0.9 \
      --extra_llm_api_options ./extra-llm-api-config.yml
  "

Step 2. Test the Draft-Target setup

Once the server is running, test it with API calls:

## Test completion endpoint
curl -X POST http://localhost:8000/v1/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "nvidia/Llama-3.3-70B-Instruct-FP4",
    "prompt": "Explain the benefits of speculative decoding:",
    "max_tokens": 150,
    "temperature": 0.7
  }'

Key Features of Draft-Target:

• Efficient resource usage: 8B draft model accelerates 70B target model
• Flexible configuration: Adjustable draft token length for optimization
• Memory efficient: Uses FP4 quantized models for reduced memory footprint
• Compatible models: Uses Llama family models with consistent tokenization

Troubleshooting

Common issues and solutions:

Symptom	Cause	Fix
"CUDA out of memory" error	Insufficient GPU memory	Reduce kv_cache_free_gpu_memory_fraction to 0.9 or use a device with more VRAM
Container fails to start	Docker GPU support issues	Verify nvidia-docker is installed and --gpus=all flag is supported
Model download fails	Network or authentication issues	Check HuggingFace authentication and network connectivity
Server doesn't respond	Port conflicts or firewall	Check if port 8000 is available and not blocked

Cleanup

Stop the Docker container when finished:

## Find and stop the container
docker ps
docker stop <container_id>

## Optional: Clean up downloaded models from cache
## rm -rf $HOME/.cache/huggingface/hub/models--*gpt-oss*

Next Steps

• Compare both Eagle3 and Draft-Target performance with baseline inference
• Experiment with different max_draft_len values (1, 2, 3, 4, 8) for both approaches
• Monitor token acceptance rates and throughput improvements across different model pairs
• Test with different prompt lengths and generation parameters
• Compare Eagle3 vs Draft-Target approaches for your specific use case
• Benchmark memory usage differences between the two methods