docs: TRTLLM Example of Llama4+Eagle3 (Speculative Decoding)

[KrishnanPrash]

Overview:

• Speculative Decoding Example of Llama 4 + Eagle 3 utilizing Speculative Decoding (Multi-Token Prediction)

Summary by CodeRabbit

• New Features
  • Added configuration files for LLaMA4 Eagle model deployment, including settings for aggregation, decoding, and prefill engines.
  • Introduced YAML files for both aggregated and disaggregated deployment modes, supporting resource allocation, batch sizes, token limits, and speculative decoding.
  • Provided options for model serving endpoints, GPU allocation, and advanced features such as CUDA graph optimizations and FP8/FP4 precision support.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[KrishnanPrash]

TODO: Add commands to run example to examples/tensorrt_llm/README.md

[coderabbitai]

Walkthrough

Six new YAML configuration files have been added to the TensorRT LLM examples directory for the LLaMA4 Eagle model. These files define engine parameters, runtime settings, and deployment configurations for aggregation, decoding, prefill, and both aggregated and disaggregated model serving setups, including resource allocation and speculative decoding options.

Changes

File(s)	Change Summary
.../agg_config.yaml .../decode_config.yaml .../prefill_config.yaml	Added engine configuration files specifying backend, parallelism, batch sizes, token limits, key-value cache settings, speculative decoding, CUDA graph options, and logging for LLaMA4 Eagle.
.../mtp_agg.yaml	Added aggregated deployment configuration with Frontend and TensorRTLLMWorker sections, model serving parameters, and resource allocation.
.../mtp_disagg.yaml	Added disaggregated deployment configuration with Frontend, TensorRTLLMWorker, and TensorRTLLMPrefillWorker sections, supporting modular and distributed serving.

Poem

A warren of configs, fresh and new,
For LLaMA4 Eagle, the models grew.
Prefill, decode, and aggregate too,
With YAMLs to guide what the engines do.
Rabbits rejoice—deployment is neat!
🐇✨ Now LLM servings are hard to beat!

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Explain this complex logic.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai explain this code block.
  • @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read src/utils.ts and explain its main purpose.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
  • @coderabbitai help me debug CodeRabbit configuration file.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[coderabbitai]

Actionable comments posted: 3

🧹 Nitpick comments (3)

examples/tensorrt_llm/configs/llama4/eagle/engine_configs/agg_config.yaml (1)

50-50: kv_cache_dtype: fp8 is experimental – guard with capability check
FP8 cache is currently limited to Hopper/Blackwell class GPUs and recent TRT-LLM nightly builds. Add an inline comment or pre-run capability probe to avoid silent precision down-casting on older cards.

examples/tensorrt_llm/configs/llama4/eagle/engine_configs/prefill_config.yaml (1)

28-33: Inconsistent KV-cache memory policy across engines
Prefill leaves 75 % of memory free, whereas decode/agg leave 15–70 %. Such disparity may starve the prefill worker when sharing a node with decode workers. Align the fractions or document the reasoning to avoid debugging surprises.

examples/tensorrt_llm/configs/llama4/eagle/engine_configs/decode_config.yaml (1)

26-31: max_num_tokens: 512 sits exactly at the formula boundary – leave head-room
max_num_tokens >= max_batch(256) × (layers+1)(2) ⇒ 512. Any future bump of batch-sizes or num_nextn_predict_layers will break graph capture. Recommend 640 or 768 to stay safe.

📜 Review details

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 6835dd7 and 7b4c32a.

📒 Files selected for processing (5)

• examples/tensorrt_llm/configs/llama4/eagle/engine_configs/agg_config.yaml (1 hunks)
• examples/tensorrt_llm/configs/llama4/eagle/engine_configs/decode_config.yaml (1 hunks)
• examples/tensorrt_llm/configs/llama4/eagle/engine_configs/prefill_config.yaml (1 hunks)
• examples/tensorrt_llm/configs/llama4/eagle/mtp_agg.yaml (1 hunks)
• examples/tensorrt_llm/configs/llama4/eagle/mtp_disagg.yaml (1 hunks)

🧰 Additional context used

🧠 Learnings (3)

📓 Common learnings

Learnt from: tanmayv25
PR: ai-dynamo/dynamo#1391
File: examples/tensorrt_llm/common/base_engine.py:171-176
Timestamp: 2025-06-05T01:10:51.865Z
Learning: In examples/tensorrt_llm/common/base_engine.py, the _init_engine method is called only once during initialization, so direct mutation of the _default_sampling_params object during setup is safe and appropriate.

examples/tensorrt_llm/configs/llama4/eagle/mtp_agg.yaml (1)

Learnt from: nnshah1
PR: ai-dynamo/dynamo#1444
File: tests/fault_tolerance/configs/agg_tp_1_dp_8.yaml:31-38
Timestamp: 2025-07-01T15:33:53.262Z
Learning: In fault tolerance test configurations, the `resources` section under `ServiceArgs` specifies resources per individual worker, not total resources for all workers. So `workers: 8` with `gpu: '1'` means 8 workers × 1 GPU each = 8 GPUs total.

examples/tensorrt_llm/configs/llama4/eagle/mtp_disagg.yaml (1)

Learnt from: nnshah1
PR: ai-dynamo/dynamo#1444
File: tests/fault_tolerance/configs/agg_tp_1_dp_8.yaml:31-38
Timestamp: 2025-07-01T15:33:53.262Z
Learning: In fault tolerance test configurations, the `resources` section under `ServiceArgs` specifies resources per individual worker, not total resources for all workers. So `workers: 8` with `gpu: '1'` means 8 workers × 1 GPU each = 8 GPUs total.

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)

• GitHub Check: Build and Test - vllm

🔇 Additional comments (5)

examples/tensorrt_llm/configs/llama4/eagle/engine_configs/agg_config.yaml (2)

30-30: free_gpu_memory_fraction looks aggressive for 8 K-token, 256-batch workloads
With max_num_tokens: 8448, batch size 256, and kv_cache_dtype: fp8, leaving 30 % of the GPU memory unused may still push large sequences into OOM on 24 GB cards. Please double-check with the target HW and consider lowering the fraction (e.g. 0.15–0.20) or tightening max_batch_size.

---

32-36: Speculative MTP enabled on the aggregate engine – validate downstream decode logic
decoding_type: MTP + num_nextn_predict_layers: 1 assumes the aggregated engine will internally handle MTP. In most TRT-LLM deployments, MTP is applied only on the decode stage, while prefill stays vanilla. Verify that the runtime really supports MTP in aggregated mode; otherwise move this block to the decode-only config.

examples/tensorrt_llm/configs/llama4/eagle/engine_configs/prefill_config.yaml (1)

34-38: Confirm that MTP on the prefill stage is intended
Prefill traditionally computes only ISL; speculative tokens are generated during decode. Enabling MTP here will allocate extra scratch buffers without benefit. Consider removing the block unless you have a custom prefill-MTP kernel.

examples/tensorrt_llm/configs/llama4/eagle/engine_configs/decode_config.yaml (1)

32-34: Extremely high free_gpu_memory_fraction (0.85)
Only 15 % of memory is kept for the KV cache, while decode is usually the most memory-hungry phase. Please double-check; values >0.5 commonly lead to runtime allocation failures.

examples/tensorrt_llm/configs/llama4/eagle/mtp_agg.yaml (1)

29-31: GPU allocation doesn’t match TP×EP world size
tensor_parallel_size:4 and moe_expert_parallel_size:4 imply 16 ranks, yet only 8 GPUs are allocated (workers:1, gpu: 8). Either halve EP/TP or request 16 GPUs to avoid NCCL topology errors.

[Tabrizian]

Please update the TRT-LLM commit as well.

accidentally clicked approve, please resolve comments.

[richardhuo-nv]

We should add how to run these configs in the README.

You can reference the guide for MTP.