kanshan/TensorRT-LLMs
1
0
Fork 0
mirror of https://github.com/NVIDIA/TensorRT-LLM.git synced 2026-02-11 05:23:38 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Merge branch 'main' into fix_spec_gate

Browse Source
This commit is contained in:
Zheyu Fu 2026-01-07 16:24:57 -08:00 committed by GitHub
commit 116853cdb4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 29 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

38
docs/source/features/feature-combination-matrix.md
View File

@ -1,19 +1,23 @@
# Feature Combination Matrix
| Feature | Overlap Scheduler | CUDA Graph | Attention Data Parallelism | Disaggregated Serving | Chunked Prefill | MTP | EAGLE-3(One Model Engine) | EAGLE-3(Two Model Engine) | Torch Sampler | TLLM C++ Sampler | KV Cache Reuse | Slide Window Attention | Logits Post Processor | Guided Decoding | LoRA |
| -------------------------- | ----------------- | ---------- | -------------------------- | --------------------- | --------------- | -------- | ------------------------- | ------------------------- | ------------- | ---------------- | -------------- | ---------------------- | --------------------- | --------------- | ---- |
| Overlap Scheduler | --- | | | | | | | | | | | | | | |
| CUDA Graph | Yes | --- | | | | | | | | | | | | | |
| Attention Data Parallelism | Yes | Yes | --- | | | | | | | | | | | | |
| Disaggregated Serving | Yes | Yes | Yes | --- | | | | | | | | | | | |
| Chunked Prefill | Yes | Yes | Yes | Yes | --- | | | | | | | | | | |
| MTP | Yes | Yes | Yes | Yes | Yes | --- | | | | | | | | | |
| EAGLE-3(One Model Engine) | Yes | Yes | Yes | Yes | Yes | No | --- | | | | | | | | |
| EAGLE-3(Two Model Engine) | Yes | Yes | Yes | Yes | Yes | No | No | --- | | | | | | | |
| Torch Sampler | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | --- | | | | | | |
| TLLM C++ Sampler | Yes | Yes | Yes | Yes | Yes | No | No | No | No | --- | | | | | |
| KV Cache Reuse | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | --- | | | | |
| Slide Window Attention | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | --- | | | |
| Logits Post Processor | Yes | Yes | Yes | No | Yes | No | No | No | Yes | Yes | Yes | Yes | --- | | |
| Guided Decoding | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | --- | |
| LoRA | Yes | No | Untested | Untested | Untested | Untested | Untested | Untested | Yes | Yes | Yes | Yes | Yes | Untested | --- |
| Feature | Overlap Scheduler | CUDA Graph | Tensor Parallelism | Pipeline Parallelism | Expert Parallelism | Helix Parallelism | Attention Data Parallelism | Disaggregated Serving | Chunked Prefill | MTP | EAGLE-3(One Model Engine) | EAGLE-3(Two Model Engine) | Torch Sampler | TLLM C++ Sampler | KV Cache Reuse | Slide Window Attention | Logits Post Processor | Guided Decoding | LoRA |
| -------------------------- | ----------------- | ---------- | ------------------ | -------------------- | ------------------ | ----------------- | -------------------------- | --------------------- | --------------- | -------- | ------------------------- | ------------------------- | ------------- | ---------------- | -------------- | ---------------------- | --------------------- | --------------- | -------- |
| Overlap Scheduler | --- | | | | | | | | | | | | | | | | | | |
| CUDA Graph | Yes | --- | | | | | | | | | | | | | | | | | |
| Tensor Parallelism | Yes | Yes | --- | | | | | | | | | | | | | | | | |
| Pipeline Parallelism | Yes | Yes | Yes | --- | | | | | | | | | | | | | | | |
| Expert Parallelism | Yes | Yes | Yes | Yes | --- | | | | | | | | | | | | | | |
| Helix Parallelism | Untested | Yes | Yes | Yes | Yes | --- | | | | | | | | | | | | | |
| Attention Data Parallelism | Yes | Yes | Yes | Yes | Yes | Known issues | --- | | | | | | | | | | | | |
| Disaggregated Serving | Yes | Yes | Yes | Yes | Yes | Yes | Yes | --- | | | | | | | | | | | |
| Chunked Prefill | Yes | Yes | Yes | Untested | Yes | Yes | Yes | Yes | --- | | | | | | | | | | |
| MTP | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | --- | | | | | | | | | |
| EAGLE-3(One Model Engine) | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | No | --- | | | | | | | | |
| EAGLE-3(Two Model Engine) | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | No | No | --- | | | | | | | |
| Torch Sampler | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | --- | | | | | | |
| TLLM C++ Sampler | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | No | No | --- | | | | | |
| KV Cache Reuse | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | --- | | | | |
| Slide Window Attention | Yes | Yes | Yes | Yes | Yes | Untested | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | --- | | | |
| Logits Post Processor | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | Yes | No | No | No | Yes | Yes | Yes | Yes | --- | | |
| Guided Decoding | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | --- | |
| LoRA | Yes | No | Yes | Yes | Untested | Untested | Untested | Untested | Yes | Untested | Untested | Untested | Yes | Yes | Yes | Yes | Yes | Untested | --- |

8
tensorrt_llm/_torch/models/modeling_speculative.py
View File

@ -953,6 +953,14 @@ class SpecDecOneEngineForCausalLM(DecoderModelForCausalLM[TModel, TConfig],
hidden_states = hidden_states[:attn_metadata.num_tokens]
if self.draft_model is not None:
# For one-model speculative decoding with PP, only the last PP rank
# has valid hidden_states from the target model. The spec_worker (which
# runs the draft model loop) should only run on the last PP rank.
# Non-last PP ranks return None and let the PP sync handle the results.
mapping = self.model.model_config.mapping
if mapping.has_pp() and not mapping.is_last_pp_rank():
return None
# get logits
logits = self.logits_processor.forward(
hidden_states[spec_metadata.gather_ids],