From d9f262b5d535d6cb54d721d280dab99ef7c4aa6d Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 8 Jul 2026 08:10:23 -0700 Subject: [PATCH 1/8] [TRTLLM-13321][feat] Extend rejection sampling to one-model speculative decoding Support rejection sampling for the one-model speculative-decoding family (vanilla MTP, DRAFT_TARGET, PARD, DFlash, MTP-Eagle; Eagle3 pre-existing). Unify draft-token production, sampling, and acceptance behind shared SpecWorkerBase entry points (sample_draft_tokens, sample_and_accept_draft_tokens), gate the d2t-expanded full_draft_probs buffer on the draft head's own vocab, detect vocab-sharded draft logits by the draft head's width, and add the config-level gating for unsupported combinations in llm_args. Signed-off-by: ZhaoyangWang --- tensorrt_llm/_torch/model_config.py | 5 + .../_torch/models/modeling_speculative.py | 18 +- tensorrt_llm/_torch/models/modeling_utils.py | 2 +- tensorrt_llm/_torch/speculative/dflash.py | 29 +- .../_torch/speculative/draft_target.py | 45 +- tensorrt_llm/_torch/speculative/eagle3.py | 166 +--- .../_torch/speculative/eagle3_dynamic_tree.py | 19 +- tensorrt_llm/_torch/speculative/interface.py | 716 +++++++++++++----- tensorrt_llm/_torch/speculative/mtp.py | 103 +-- tensorrt_llm/_torch/speculative/pard.py | 75 +- tensorrt_llm/_torch/speculative/utils.py | 48 ++ tensorrt_llm/llmapi/llm_args.py | 138 +++- .../test_lists/test-db/l0_h100.yml | 1 + .../speculative/hw_agnostic/test_dflash.py | 30 + .../hw_agnostic/test_draft_target.py | 46 ++ .../speculative/hw_agnostic/test_mtp.py | 96 ++- .../speculative/hw_agnostic/test_pard.py | 45 ++ .../test_rejection_buffers_guard.py | 275 +++++++ 18 files changed, 1305 insertions(+), 552 deletions(-) create mode 100644 tests/unittest/_torch/speculative/test_rejection_buffers_guard.py diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index fd15550fd6e1..7b445e991c21 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -141,6 +141,11 @@ class ModelConfig(Generic[TConfig]): skip_create_weights_in_init: bool = False spec_config: Optional["DecodingBaseConfig"] = None + # When False, the column-parallel LM head keeps its vocab-sharded output + # instead of all-gathering to full vocab. Used for one-model speculative + # draft models so greedy draft sampling can do a lighter TP gather. Defaults + # to True to preserve behavior for every non-draft model. + lm_head_gather_output: bool = True lora_config: Optional["LoraConfig"] = None sparse_attention_config: Optional["SparseAttentionConfig"] = None diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 6e17258067a1..df685c4f2a27 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -777,7 +777,9 @@ def __init__(self, draft_config): draft_config.pretrained_config) # Remove spec_config to prevent recursive spec-dec initialization - draft_config_no_spec = replace(draft_config, spec_config=None) + draft_config_no_spec = replace(draft_config, + spec_config=None, + lm_head_gather_output=False) # Weights will be loaded later by ModelLoader.load_draft_weights() self.draft_model_full = DraftModelClass(draft_config_no_spec) @@ -863,7 +865,9 @@ def __init__(self, draft_config): pretrained_cfg.architectures = original_archs # Remove spec_config to prevent recursive spec-dec initialization - draft_config_no_spec = replace(draft_config, spec_config=None) + draft_config_no_spec = replace(draft_config, + spec_config=None, + lm_head_gather_output=False) # Weights will be loaded later by ModelLoader.load_draft_weights() self.draft_model_full = DraftModelClass(draft_config_no_spec) @@ -1653,6 +1657,12 @@ def get_draft_model(model_config, draft_config, lm_head, model): elif spec_dec_mode.is_dflash(): return DFlashForCausalLM(draft_config) elif spec_dec_mode.is_draft_target_one_model(): + # Keep the draft LM head vocab-sharded so greedy draft sampling uses the + # lighter TP gather (see SpecWorkerBase.greedy_sample_draft_with_tp_gather). + was_frozen = draft_config._frozen + draft_config._frozen = False + draft_config.lm_head_gather_output = False + draft_config._frozen = was_frozen return AutoModelForCausalLM.from_config(draft_config) else: raise NotImplementedError( @@ -1741,6 +1751,10 @@ def __init__(self, model: TModel, model_config: ModelConfig[TConfig]): model_config.mapping, use_separate_draft_kv_cache=self.use_separate_draft_kv_cache) if self.spec_worker is not None: + # Cache the static draft->target vocab map now that the draft + # model is loaded, so workers read self._d2t instead of probing + # draft_model.model.d2t on every forward. + self.spec_worker.set_draft_model(self.draft_model) self.epilogue.append(self.spec_worker) self.layer_idx = -1 diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 5355bb7d3193..c45572f577f6 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -431,7 +431,7 @@ def __init__(self, model: TModel, *, config: ModelConfig[TConfig], dtype=config.pretrained_config.torch_dtype, mapping=config.mapping, tensor_parallel_mode=TensorParallelMode.COLUMN, - gather_output=True, + gather_output=config.lm_head_gather_output, reduce_output=False, use_custom_cublas_mm=getattr(model, 'use_custom_cublas_mm', False), diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index efa31fa2cc1e..4ef753b78530 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -407,17 +407,8 @@ def forward( self._execute_guided_decoder_if_present(logits) - # Target now emits K+1 logits per gen request and the previous step - # stored K draft tokens per gen request (no filler padding). - if num_gens > 0: - draft_tokens = spec_metadata.draft_tokens.reshape(num_gens, K) - else: - draft_tokens = spec_metadata.draft_tokens.reshape(0, K) - - logits_for_accept = logits - - accepted_tokens, num_accepted_tokens = self._sample_and_accept_draft_tokens_base( - logits_for_accept, draft_tokens, num_contexts, batch_size, spec_metadata + accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( + logits, attn_metadata, spec_metadata ) # Update GDN/Mamba recurrent states to the accepted token's state. @@ -471,6 +462,8 @@ def forward( draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + self.reset_draft_probs_valid_for_capture(spec_metadata) + if num_gens > 0: with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): hidden_states_out = draft_model.dflash_forward( @@ -498,13 +491,13 @@ def forward( vocab_size = gen_logits.shape[-1] gen_logits = gen_logits.reshape(num_gens, K, vocab_size) - d2t = getattr(draft_model.model, "d2t", None) - gen_draft_tokens = torch.argmax(gen_logits, dim=-1, keepdim=False).long() - - if d2t is not None: - gen_draft_tokens = d2t[gen_draft_tokens] + gen_draft_tokens - - gen_draft_tokens = gen_draft_tokens.type(torch.int32) + gen_draft_tokens = self.sample_draft_tokens( + gen_logits, + spec_metadata, + batch_size, + num_contexts=num_contexts, + is_last_draft_cycle=True, + ) else: gen_draft_tokens = torch.empty((0, K), dtype=torch.int32, device="cuda") diff --git a/tensorrt_llm/_torch/speculative/draft_target.py b/tensorrt_llm/_torch/speculative/draft_target.py index fe6dc7d1e4c7..47c51187e617 100644 --- a/tensorrt_llm/_torch/speculative/draft_target.py +++ b/tensorrt_llm/_torch/speculative/draft_target.py @@ -213,6 +213,8 @@ def forward( next_draft_tokens = [] original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens + self.reset_draft_probs_valid_for_capture(spec_metadata) + # Get the draft KV cache manager if using separate layouts draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) @@ -257,10 +259,15 @@ def forward( hidden_states[gather_ids], draft_model.lm_head, attn_metadata, True ) if self.guided_decoder is not None: - d2t = getattr(draft_model.model, "d2t", None) - self.guided_decoder.execute_draft_batch(logits, d2t, draft_step=i) - - new_draft_token = self.draft_decoder(logits, draft_model) + self.guided_decoder.execute_draft_batch(logits, self._d2t, draft_step=i) + + new_draft_token = self.sample_draft_tokens( + logits, + spec_metadata, + batch_size, + draft_step=i, + is_last_draft_cycle=(i == runtime_draft_len - 1), + ) next_draft_tokens.append(new_draft_token) # Update inputs and metadata for next draft step @@ -314,36 +321,6 @@ def forward( "next_new_tokens": next_new_tokens, } - def sample_and_accept_draft_tokens( - self, - logits: torch.Tensor, - attn_metadata: AttentionMetadata, - spec_metadata: DraftTargetOneModelSpecMetadata, - ): - batch_size = attn_metadata.num_seqs - num_contexts = attn_metadata.num_contexts - num_gens = batch_size - num_contexts - runtime_draft_len = spec_metadata.runtime_draft_len - - if spec_metadata.draft_tokens is None: - draft_tokens = torch.zeros( - (num_gens, runtime_draft_len), dtype=torch.int, device=logits.device - ) - else: - draft_tokens = spec_metadata.draft_tokens.reshape(num_gens, runtime_draft_len) - - return self._sample_and_accept_draft_tokens_base( - logits, draft_tokens, num_contexts, batch_size, spec_metadata - ) - - def draft_decoder( - self, - logits: torch.Tensor, - draft_model: nn.Module, - ): - d2t = getattr(draft_model.model, "d2t", None) - return self._draft_sampler_greedy(logits, d2t) - def prepare_1st_drafter_inputs( self, input_ids: torch.LongTensor, diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index 3978040724b1..f69147fe7c73 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -10,7 +10,6 @@ from tensorrt_llm.mapping import Mapping from ..attention_backend import AttentionMetadata -from ..distributed.ops import allgather from ..model_config import ModelConfig from ..pyexecutor.llm_request import LlmRequest from ..pyexecutor.mamba_cache_manager import MambaHybridCacheManager @@ -806,6 +805,7 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, runtime_draft_len = spec_metadata.runtime_draft_len num_gens = batch_size - num_contexts next_draft_tokens = [] + self.reset_draft_probs_valid_for_capture(spec_metadata) last_tokens_idx = torch.cumsum( attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1 position_ids = inputs["position_ids"] @@ -891,9 +891,8 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, self.guided_decoder.execute_draft_batch(logits, draft_step=i) else: - d2t = getattr(draft_model.model, "d2t", None) self.guided_decoder.execute_draft_batch(logits, - d2t, + self._d2t, draft_step=i) # When ADP+LM-head-TP pads logits to max_num_requests, the @@ -905,13 +904,21 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, # would otherwise fail to broadcast in apply_temperature. This # also keeps next_draft_tokens and the draft_probs buffer # token_count-sized without a post-hoc trim. + mapping_lm_head_tp = None if use_lm_head_tp_in_adp: logits = logits[:token_count] - new_draft_token = self.draft_decoder(logits, - draft_model, - spec_metadata, - batch_size, - draft_step=i) + # The MTP head built this per-forward mapping when producing + # the vocab-sharded logits; the sampler needs it to gather. + mapping_lm_head_tp = getattr( + draft_model.mtp_layers[0].shared_head, + "mapping_lm_head_tp", None) + new_draft_token = self.sample_draft_tokens( + logits, + spec_metadata, + batch_size, + draft_step=i, + is_last_draft_cycle=(i == runtime_draft_len - 1), + mapping_lm_head_tp=mapping_lm_head_tp) next_draft_tokens.append(new_draft_token) # Update hidden states for the next iteration. @@ -977,20 +984,6 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, gen_draft_tokens) next_draft_tokens[num_contexts:] = gen_draft_tokens - # Probs were already scattered into the slot-indexed buffer by - # _draft_sampler_advanced_for_rejection on each draft step (non-greedy - # batches only). All-greedy batches skip storage — rejection sampling - # will be bypassed by _can_use_rejection_sampling. Finalize the validity - # flag and d2t for next-iter target-side verification. - if spec_metadata.use_rejection_sampling: - if not spec_metadata.is_all_greedy_sample: - d2t_param = getattr(getattr(draft_model, 'model', None), "d2t", - None) - spec_metadata.d2t = d2t_param.data if d2t_param is not None else None - spec_metadata.draft_probs_valid = True - else: - spec_metadata.draft_probs_valid = False - return next_draft_tokens def _get_step_all_rank_num_tokens(self, spec_metadata, step_idx: int): @@ -1043,63 +1036,6 @@ def _prepare_flash_mla_generation_layout(self, attn_metadata, num_contexts, attn_metadata.block_ids_per_seq[:batch_size, :].copy_( reorder_block_ids_per_seq, non_blocking=True) - @torch.compile(options={"max-autotune": True}) - def _get_local_max_and_combined(self, logits, mapping_lm_tp=None): - local_max_values, local_argmax = torch.max(logits, dim=-1, keepdim=True) - vocab_per_rank = logits.shape[-1] - mapping_lm_tp = mapping_lm_tp if mapping_lm_tp is not None else \ - self.model_config.mapping - max_index_per_rank = local_argmax.type( - torch.int32) + (mapping_lm_tp.tp_rank * vocab_per_rank) - max_index_per_rank_float = max_index_per_rank.float() - local_max_values_float32 = local_max_values.float() - combined = torch.stack( - [max_index_per_rank_float, local_max_values_float32], - dim=-1).flatten(-2) - return combined - - @torch.compile(options={"max-autotune": True}) - def _get_draft_tokens_from_gathered(self, gathered): - gathered_indices_float = gathered[..., 0::2] - gathered_values_float = gathered[..., 1::2] - max_indices = torch.argmax(gathered_values_float, dim=-1, keepdim=True) - draft_tokens = torch.gather(gathered_indices_float, -1, - max_indices).squeeze(-1).type(torch.int32) - return draft_tokens - - def draft_sampler( - self, - logits: torch.Tensor, - mapping_lm_head_tp=None, - ): - """TP-aware greedy draft token sampler (MTP Eagle path). - - Falls back to simple argmax when no tensor parallelism is active or - when only attention DP is enabled without LM-head TP. - """ - if (self.model_config is not None - and hasattr(self.model_config, 'mapping') - and self.model_config.mapping.tp_size > 1 - and not self.model_config.mapping.enable_attention_dp): - combined = self._get_local_max_and_combined(logits) - gathered = allgather(combined, self.model_config.mapping, dim=-1) - return self._get_draft_tokens_from_gathered(gathered) - elif (self.model_config is not None - and hasattr(self.model_config, 'mapping') - and self.model_config.mapping.tp_size > 1 - and self.model_config.mapping.enable_lm_head_tp_in_adp): - combined = self._get_local_max_and_combined(logits, - mapping_lm_head_tp) - gathered = allgather(combined, mapping_lm_head_tp, dim=-1) - batch_size = logits.shape[0] - local_batch_size = batch_size // mapping_lm_head_tp.tp_size - gathered = gathered.view(mapping_lm_head_tp.tp_size, - local_batch_size, -1) - sliced_gathered = gathered[mapping_lm_head_tp.tp_rank] - return self._get_draft_tokens_from_gathered(sliced_gathered) - else: - return self._draft_sampler_greedy(logits) - @torch.compile(options={"max-autotune": True}) def _topk_kernel(self, gen_logprobs, num_gens, mtp_num_modules, spec_metadata): @@ -1210,78 +1146,6 @@ def sample_and_accept_draft_tokens( return self._accept_draft_tokens(logits, draft_tokens, num_contexts, batch_size, spec_metadata) - def draft_decoder( - self, - logits: torch.Tensor, - draft_model: nn.Module, - spec_metadata: Eagle3OneModelSpecMetadata, - batch_size: int, - draft_step: Optional[int] = None, - ): - ''' - Sample draft tokens using the target's per-request sampling params - (temperature/top_k/top_p). - - When rejection sampling is enabled and draft_step is provided, take the - single-pass path that also scatters the draft prob distribution into the - slot-indexed buffer (avoids a redundant softmax later). - - Args: - logits: [batch_size, vocab_size] - Draft model logits. - draft_model: The draft model. - spec_metadata: Carries per-request sampling param tensors. - batch_size: Active requests, used to slice per-request tensors. - draft_step: Current draft step index (0..max_draft_len-1). Required - for the rejection-sampling code path so probs are written to - the correct slice of spec_metadata.draft_probs. - ''' - - d2t = getattr(getattr(draft_model, 'model', None), "d2t", None) - # All-greedy fast path must stay TP-aware. When the draft LM head is - # tensor-parallel (tp_size>1 without attention DP, or LM-head-TP in - # ADP), the draft logits are sharded along the vocab dim. A plain - # per-rank argmax then picks a different token on each rank, which - # desyncs the speculative-decoding control flow across ranks and - # deadlocks the next collective (observed as a generation hang on - # MTP-Eagle + TP). draft_sampler() all-gathers the sharded logits - # before argmax (and falls back to a plain argmax when no TP gather is - # needed). Eagle3 (non-MTP) keeps its d2t-aware argmax. - if spec_metadata.is_all_greedy_sample: - # Only plain tensor parallelism (tp_size>1 without attention DP) - # shards the draft logits over the vocab dim and thus needs - # draft_sampler()'s all-gather argmax. The LM-head-TP-in-ADP case - # already produces full-vocab logits per rank (gathered upstream), - # and the no-TP / Eagle3 cases need nothing, so they take the plain - # d2t-aware argmax. (Routing ADP/LM-head-TP through draft_sampler - # without its mapping_lm_head_tp arg hits the None-mapping branch - # and crashes with 'NoneType has no attribute tp_group'.) - if (self.is_mtp_eagle and self.model_config is not None - and hasattr(self.model_config, 'mapping') - and self.model_config.mapping.tp_size > 1 - and not self.model_config.mapping.enable_attention_dp): - return self.draft_sampler(logits) - return self._draft_sampler_greedy(logits, d2t) - # Non-greedy (advanced) draft sampling has the same TP hazard as the - # greedy path: when the draft LM head is plain tensor-parallel - # (tp_size>1 without attention DP), each rank only holds a vocab shard - # of the draft logits. Random per-rank sampling then draws different - # tokens on different ranks, desyncing the spec-decode control flow and - # deadlocking the next collective. All-gather the shards into the full - # vocab first so every rank samples from the same distribution with the - # shared seed. (Greedy uses draft_sampler()'s lighter max+index gather; - # random sampling needs the full distribution. The LM-head-TP-in-ADP - # case is handled upstream and must not be gathered again here.) - if (self.is_mtp_eagle and self.model_config is not None - and hasattr(self.model_config, 'mapping') - and self.model_config.mapping.tp_size > 1 - and not self.model_config.mapping.enable_attention_dp): - logits = allgather(logits, self.model_config.mapping, dim=-1) - if spec_metadata.use_rejection_sampling and draft_step is not None: - return self._draft_sampler_advanced_for_rejection( - logits, spec_metadata, batch_size, d2t, draft_step) - return self._draft_sampler_advanced(logits, spec_metadata, batch_size, - d2t) - def prepare_1st_drafter_inputs( self, input_ids: torch.LongTensor, diff --git a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py index 7d9e836f5b9a..7edbcc5e6756 100644 --- a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py @@ -585,7 +585,6 @@ def _forward_draft_loop( ): """Dynamic tree draft loop with growing context.""" spec_tree_manager = self.spec_tree_manager - self._d2t = getattr(draft_model.model, "d2t", None) assert batch_size <= self._max_batch_size, ( f"batch_size {batch_size} exceeds pre-allocated max_batch_size {self._max_batch_size}" @@ -636,9 +635,7 @@ def _forward_draft_loop( hidden_states[gather_ids], draft_model.lm_head, attn_metadata, True ) - new_draft_tokens, new_draft_scores = self.sample( - logits, self.K, draft_model=draft_model - ) + new_draft_tokens, new_draft_scores = self.sample(logits, self.K) previous_draft_scores = self.update_draft_tokens_and_scores( cur_draft_idx=0, @@ -698,9 +695,7 @@ def _forward_draft_loop( selected_hs, draft_model.lm_head, attn_metadata, True ) - new_draft_tokens, new_draft_scores = self.sample( - logits, self.K, draft_model=draft_model - ) + new_draft_tokens, new_draft_scores = self.sample(logits, self.K) # Reshape for update: [batch_size, K, K] new_draft_tokens = new_draft_tokens.reshape(batch_size, self.K, self.K) @@ -1015,14 +1010,12 @@ def _finalize_dynamic_tree_verify_outputs( ).to(torch.int32) @nvtx_range("eagle3_dyn.sample") - def sample( - self, logits: torch.Tensor, max_top_k: int, draft_model=None - ) -> tuple[torch.Tensor, torch.Tensor]: + def sample(self, logits: torch.Tensor, max_top_k: int) -> tuple[torch.Tensor, torch.Tensor]: """TopK sampling with softmax for dynamic tree.""" topk_indices, topk_values = _sample_softmax_topk(logits, max_top_k) - # Apply draft-to-target vocab mapping if the draft model has it - if draft_model is not None and hasattr(draft_model.model, "d2t"): - d2t = draft_model.model.d2t.data + # Apply draft-to-target vocab mapping when draft/target vocabs differ. + if self._d2t is not None: + d2t = self._d2t.data topk_indices = topk_indices + d2t[topk_indices] return topk_indices, topk_values diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 6244873efc43..ed7a2b59d9fa 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -521,7 +521,6 @@ class SpecMetadata: skip_temperature: bool = False skip_top_k: bool = False skip_top_p: bool = False - has_greedy_requests: bool = False # Pre-computed top_k_max scalar (CPU-side) to avoid CUDA-graph-incompatible # dynamic boolean tensor indexing inside verify_dynamic_tree_rejection_from_logits_out. top_k_max: int = 0 @@ -539,6 +538,10 @@ class SpecMetadata: # Slot-indexed buffers (draft_probs) must span this full range. # 0 falls back to max_num_requests. num_seq_slots: int = 0 + # Draft-model vocab size. full_draft_probs is allocated only when it differs + # from vocab_size; 0 (unknown) or a value equal to vocab_size means shared + # vocab and skips the buffer. + draft_vocab_size: int = 0 # Draft probabilities buffer for rejection sampling, indexed by py_seq_slot # so per-request data is stable across iterations regardless of batch # composition shifts (chunking ctx, gen completion, new ctx joining). @@ -553,54 +556,61 @@ class SpecMetadata: # Used to scatter draft probs by slot at write time and gather them by slot # at the next iter's verify. Shape: [max_num_requests], dtype=long. batch_slot_ids: Optional[torch.Tensor] = None - # Draft-to-target vocab offset tensor. - d2t: Optional[torch.Tensor] = None - # Pre-allocated scratch for draft probs expanded to the target vocab size. - # Filled with zeros once at prepare(); each rejection iter only overwrites - # the positions selected by d2t (or [:draft_vocab] when there is no d2t), - # so the zeros outside those positions persist across iterations and we - # avoid a per-iter 64 MB zero-fill on the (max_num_requests, max_draft_len, - # vocab_size) tensor. Shape: [max_num_requests, max_draft_len, vocab_size]. + # Draft probs expanded to the target vocab size. Zero-filled once at + # prepare(); each rejection iter overwrites only the d2t-selected positions + # (or [:draft_vocab] when there is no d2t). + # Shape: [max_num_requests, max_draft_len, vocab_size]. full_draft_probs: Optional[torch.Tensor] = None - # Cached d2t-projected target vocab indices, computed once on first use - # (d2t is a model-static tensor). Replaces the per-iter - # arange + (source + d2t) % vocab_size kernel sequence inside the d2t - # padding step. Shape: [draft_vocab_size], dtype long. + # Cached d2t-projected target vocab indices, computed once on first use. + # Shape: [draft_vocab_size], dtype long. d2t_target_indices: Optional[torch.Tensor] = None def __post_init__(self): pass - def prepare(self): + def prepare_rejection_sampling_buffers(self): """ - Hook to be called before the forward step of the model. + Allocate the slot-indexed buffers used by one-model rejection sampling. + + Idempotent and gated on ``use_rejection_sampling``. """ - if (self.use_rejection_sampling and self.draft_probs is None - and self.vocab_size > 0): - # 3D [slot, draft_step, vocab] so we can scatter/gather by slot id - # and avoid the brittle "batch position == buffer position" mapping. - slot_capacity = self.num_seq_slots or self.max_num_requests + if not self.use_rejection_sampling: + return + + # Slot-indexed buffers span the full SeqSlotManager pool: py_seq_slot + # can range over [0, num_seq_slots), which under DeepSeek-V4 overlap + # exceeds max_num_requests. Fall back to max_num_requests when the pool + # size is unknown (0). One extra scratch row at index ``slot_capacity`` + # absorbs CUDA-graph dummy/padding requests (``py_seq_slot is None``). + slot_capacity = self.num_seq_slots or self.max_num_requests + num_slot_rows = slot_capacity + 1 + + if self.draft_probs is None and self.vocab_size > 0: + # [slot, draft_step, vocab]: scatter/gather by stable slot id. self.draft_probs = torch.empty( - (slot_capacity, self.max_draft_len, self.vocab_size), + (num_slot_rows, self.max_draft_len, self.vocab_size), dtype=torch.float32, device='cuda') self.draft_probs_vocab_size = self.vocab_size - if (self.use_rejection_sampling and self.batch_slot_ids is None - and self.max_num_requests > 0): + if self.batch_slot_ids is None and self.max_num_requests > 0: self.batch_slot_ids = torch.empty((self.max_num_requests, ), dtype=torch.long, device='cuda') - if (self.use_rejection_sampling and self.full_draft_probs is None - and self.vocab_size > 0): - # Zero-fill once. Subsequent iters only overwrite the d2t-mapped - # positions (constant across iters since d2t is model-static), so - # untouched positions stay 0 forever — saves the per-iter 64 MB - # zero-fill in _sample_and_accept_draft_tokens_rejection. + # full_draft_probs (d2t-expanded) is read only when draft and target + # vocabularies differ; skip it otherwise. Zero-filled once. + if (self.full_draft_probs is None and self.vocab_size > 0 + and self.draft_vocab_size not in (0, self.vocab_size)): self.full_draft_probs = torch.zeros( - (self.max_num_requests, self.max_draft_len, self.vocab_size), + (num_slot_rows, self.max_draft_len, self.vocab_size), dtype=torch.float32, device='cuda') + def prepare(self): + """ + Hook to be called before the forward step of the model. + """ + self.prepare_rejection_sampling_buffers() + def create_cuda_graph_metadata(self, max_batch_size: int): """ Creates metadata for CUDA graph execution. @@ -634,8 +644,8 @@ def _scan_one_model_sampling( ) -> tuple[list[tuple[float, int, float, int]], list[int]]: """Single source of truth for one-engine sampling-param detection. - Scans the batch's sampling configs and sets skip_*/has_greedy_requests/ - is_all_greedy_sample (honoring the warmup capture override). Returns + Scans the batch's sampling configs and sets skip_*/is_all_greedy_sample + (honoring the warmup capture override). Returns ``(per_request_normalized, per_request_slot_ids)`` for buffer population. Does NOT allocate or fill GPU buffers, so it is safe to call before the CUDA graph key is built. @@ -693,7 +703,7 @@ def _normalize_request_sampling_params( temperature_enabled = False top_k_enabled = False top_p_enabled = False - has_greedy_requests = False + has_non_greedy_requests = False per_request_slot_ids: list[int] = [] for request in requests: @@ -722,32 +732,32 @@ def _normalize_request_sampling_params( temperature_enabled |= use_temperature top_k_enabled |= use_top_k top_p_enabled |= use_top_p - has_greedy_requests |= is_greedy + has_non_greedy_requests |= not is_greedy per_request_normalized.append( (temp_val, tk_val, tp_val, num_tokens)) - # py_seq_slot is a stable per-request id used to scatter / gather - # draft probs across iterations. Dummies / unallocated slots fall - # back to 0 (any valid index is fine — the data at that slot will - # be overwritten on the next real iteration before being read). + # py_seq_slot is a stable per-request id used to scatter/gather draft + # probs across iterations. Dummy/padding requests (py_seq_slot is + # None) route to the scratch row at index ``slot_capacity`` (== + # num_seq_slots or max_num_requests), matching the buffer sizing in + # prepare_rejection_sampling_buffers so a real slot never collides. per_request_slot_ids.append( - request.py_seq_slot if request.py_seq_slot is not None else 0) + request.py_seq_slot if request.py_seq_slot is not None else + (self.num_seq_slots or self.max_num_requests)) self.skip_temperature = not temperature_enabled self.skip_top_k = not top_k_enabled self.skip_top_p = not top_p_enabled - self.has_greedy_requests = has_greedy_requests # Used in the CUDA graph key to pick the argmax / advanced variant. - self.is_all_greedy_sample = (self.skip_temperature and self.skip_top_k - and self.skip_top_p) - - # Warmup-time override (set via runtime attribute by the model engine): - # force the advanced-sampling code path so the CUDA graph for the - # (is_all_greedy_sample=False) key gets captured. Dummy warmup requests - # carry no sampling params, so the natural detection above always - # returns True; this branch substitutes synthetic non-greedy scalars - # into the per-request data and lets Phase 2 run normally to populate - # the GPU buffers used by the captured kernels. + # All-greedy iff EVERY request is greedy. Derived from per-request + # greediness, not from the skip_* filter flags (a non-greedy request may + # enable no filter, e.g. temperature=1.0 with top_k/top_p unset). + self.is_all_greedy_sample = not has_non_greedy_requests + + # Warmup-time override: force the advanced-sampling path so the CUDA + # graph for the (is_all_greedy_sample=False) key gets captured. Dummy + # warmup requests carry no sampling params, so substitute synthetic + # non-greedy scalars to populate the GPU buffers. if getattr(self, '_force_non_greedy_for_capture', False): self.skip_temperature = False self.skip_top_k = False @@ -765,13 +775,8 @@ def update_is_all_greedy_sample(self, requests: list["LlmRequest"]) -> None: Must be called BEFORE the CUDA graph key is built (the key includes ``is_all_greedy_sample`` to choose the argmax vs advanced-sampling graph - variant). ``populate_sampling_params_for_one_model`` runs later, inside - ``_prepare_inputs``, and re-derives the same flag while filling the GPU - sampling buffers. Computing the flag here first keeps the selected graph - consistent with the buffers ``populate`` fills; otherwise the key would - use the previous iteration's stale value and could replay the advanced - graph against unpopulated (greedy) buffers, which can hang/corrupt the - run (notably for MTP with num_nextn>=2). + variant), so the selected graph stays consistent with the buffers + ``populate_sampling_params_for_one_model`` fills later. """ if not self.spec_dec_mode.use_one_engine(): return @@ -789,6 +794,11 @@ def populate_sampling_params_for_one_model( if not self.spec_dec_mode.use_one_engine(): return + # Allocate the rejection buffers before copying py_seq_slot values into + # batch_slot_ids below; this runs earlier than prepare() in the + # model-engine flow. No-op unless use_rejection_sampling is set. + self.prepare_rejection_sampling_buffers() + if self.temperatures is None: # Ensures determinism across ranks. torch.manual_seed(0) @@ -917,6 +927,9 @@ def __init__(self, use_separate_draft_kv_cache: bool = False): self.seed: Optional[torch.Tensor] = None self.offset: Optional[torch.Tensor] = None self.use_separate_draft_kv_cache = use_separate_draft_kv_cache + # Static draft->target vocab offset map, cached once the draft model is + # loaded (see set_draft_model). None when draft and target share a vocab. + self._d2t: Optional[torch.Tensor] = None # Lazily-initialized state for the fractional synthetic acceptance # rate. The pool is a fixed-seed, rank-independent table of uniform # [0, 1) values; the counter is a device-side int64 advanced in-place @@ -1025,6 +1038,13 @@ def set_guided_decoder(self, self.guided_decoder = guided_decoder return True + def set_draft_model(self, draft_model) -> None: + """Cache the static draft->target vocab offset map (``d2t``) once the + draft model is loaded. ``d2t`` is a model-static parameter present only + when the draft and target vocabularies differ; stays None otherwise. + """ + self._d2t = getattr(getattr(draft_model, "model", None), "d2t", None) + def _prepare_attn_metadata_for_spec_dec(self, attn_metadata): """ Prepare attention metadata before speculative decoding draft token generation. @@ -1248,19 +1268,222 @@ def _accept_draft_tokens(self, logits, draft_tokens, num_contexts, stored_vocab = (spec_metadata.draft_probs_last_dim if spec_metadata.draft_probs_last_dim > 0 else spec_metadata.draft_probs_vocab_size) - # Gather the slot rows for the gen subset. The buffer was filled - # at the previous draft step indexed by py_seq_slot, so each gen - # request reads back exactly its own probs, regardless of batch - # composition changes since then. - gen_slot_ids = spec_metadata.batch_slot_ids[num_contexts:batch_size] - draft_probs = spec_metadata.draft_probs[ - gen_slot_ids, :draft_len, :stored_vocab] - return self._sample_and_accept_draft_tokens_rejection( - logits, draft_tokens, draft_probs, num_contexts, batch_size, - spec_metadata) + # Fail closed: run the rejection kernel only when every buffer is + # present and correctly shaped; otherwise fall back to strict + # acceptance. draft_probs_valid is reset so stale state is not reused. + if self._rejection_buffers_valid(draft_tokens, draft_len, + stored_vocab, num_contexts, + batch_size, logits, spec_metadata): + # Gather the gen subset's slot rows, filled at the previous draft + # step indexed by py_seq_slot. + gen_slot_ids = spec_metadata.batch_slot_ids[ + num_contexts:batch_size] + draft_probs = spec_metadata.draft_probs[ + gen_slot_ids, :draft_len, :stored_vocab] + return self._sample_and_accept_draft_tokens_rejection( + logits, draft_tokens, draft_probs, num_contexts, batch_size, + spec_metadata) + spec_metadata.draft_probs_valid = False return self._sample_and_accept_draft_tokens_base( logits, draft_tokens, num_contexts, batch_size, spec_metadata) + def _draft_logits_are_sharded(self, logits, spec_metadata): + """Whether the draft logits are vocab-sharded and need a TP gather. + + Sharded when tp_size>1 and the logits' last dim is narrower than the + DRAFT head's own full vocab -- either plain TP (no attention DP) or the + ADP + LM-head-TP mode, both of which produce vocab-sharded draft logits. + Replicated full-vocab logits (borrowed/gathered LM head, plain attention + DP, or a single rank) are not sharded. + + The reference width is the draft head's own full vocab (``draft_vocab_size``, + falling back to ``vocab_size`` when unknown/shared), NOT the target + ``vocab_size``: an Eagle3 reduced-vocab draft head produces full, + replicated ``[tokens, draft_vocab_size]`` logits that are narrower than the + target vocab; comparing against the target vocab would misclassify those as + sharded and gather identical copies (overflowing the d2t table). + """ + mapping = self.mapping + if mapping is None or getattr(mapping, "tp_size", 1) <= 1: + return False + # Plain attention DP (without LM-head TP) replicates full-vocab logits. + if (getattr(mapping, "enable_attention_dp", False) + and not getattr(mapping, "enable_lm_head_tp_in_adp", False)): + return False + draft_full_vocab = (getattr(spec_metadata, "draft_vocab_size", 0) + or getattr(spec_metadata, "vocab_size", 0) or 0) + return bool(draft_full_vocab) and logits.shape[-1] < draft_full_vocab + + def maybe_gather_sharded_draft_logits(self, + logits, + spec_metadata, + mapping_lm_head_tp=None): + """All-gather TP-sharded draft logits to full vocab before advanced sampling. + + Advanced (non-greedy) draft sampling needs the full-vocab distribution. + Gathers shards only for a non-greedy batch when the logits are sharded + (see ``_draft_logits_are_sharded``); replicated full-vocab logits are + returned unchanged. + + Plain TP gathers vocab shards over ``self.mapping``. Under ADP + LM-head + TP the worker has already trimmed the LM-head-TP padding rows so each rank + holds ``[token_count, vocab_shard]`` for its own tokens; a vocab-dim + all-gather over ``mapping_lm_head_tp`` restores full vocab (no token + re-slice is needed after the trim). + """ + if (spec_metadata is None or spec_metadata.is_all_greedy_sample + or not self._draft_logits_are_sharded(logits, spec_metadata)): + return logits + + from ..distributed.ops import allgather + mapping = self.mapping + if (getattr(mapping, "enable_attention_dp", False) + and getattr(mapping, "enable_lm_head_tp_in_adp", False)): + assert mapping_lm_head_tp is not None, ( + "mapping_lm_head_tp is required to gather ADP + LM-head-TP draft " + "logits") + return allgather(logits, mapping_lm_head_tp, dim=-1) + return allgather(logits, mapping, dim=-1) + + def advanced_sample_draft(self, + logits: torch.Tensor, + spec_metadata: "SpecMetadata", + batch_size: int, + draft_step: Optional[int] = None): + """Per-step advanced (non-greedy) draft sampler for step workers + (MTP, DraftTarget). + + With rejection enabled and a ``draft_step``, samples via + ``sampling_batch_spec_dec_one_model_for_rejection`` and scatters this + step's proposal distribution into the slot-indexed ``draft_probs`` + buffer; otherwise uses ``sampling_batch_spec_dec_one_model`` (tokens + only). Returns tokens in draft-vocab space (the caller applies d2t). + Expects 2D ``[batch_size, vocab]`` logits (one row per request). + """ + temperatures = spec_metadata.request_temperatures[:batch_size] + top_ks = spec_metadata.request_top_ks[:batch_size] + top_ps = spec_metadata.request_top_ps[:batch_size] + + self._update_advance_draft_sampling_seed(logits.device) + if spec_metadata.use_rejection_sampling and draft_step is not None: + draft_tokens, probs = ( + sampling_batch_spec_dec_one_model_for_rejection( + logits, + temperatures, + top_ks, + top_ps, + seed=self.seed, + offset=self.offset)) + # Scatter probs into the slot-indexed buffer so each request's data + # lands at its stable py_seq_slot row regardless of batch shifts. + assert spec_metadata.batch_slot_ids is not None, ( + "batch_slot_ids must be populated by " + "populate_sampling_params_for_one_model before draft probs " + "storage") + batch_slots = spec_metadata.batch_slot_ids[:batch_size] + vocab = probs.shape[-1] + spec_metadata.draft_probs[batch_slots, draft_step, :vocab] = probs + spec_metadata.draft_probs_last_dim = vocab + else: + draft_tokens = sampling_batch_spec_dec_one_model(logits, + temperatures, + top_ks, + top_ps, + seed=self.seed, + offset=self.offset) + + return draft_tokens.type(torch.int32) + + def _reshape_draft_tokens_for_accept(self, spec_metadata, num_gens, device): + """Reshape the stored draft tokens to ``[num_gens, runtime_draft_len]`` + for acceptance. Default assumes one draft token per step (DraftTarget, + DFlash); workers with a different buffer layout (e.g. PARD's 2K-1 + entries) override this. + """ + runtime_draft_len = spec_metadata.runtime_draft_len + if spec_metadata.draft_tokens is None: + return torch.zeros((num_gens, runtime_draft_len), + dtype=torch.int, + device=device) + return spec_metadata.draft_tokens.reshape(num_gens, runtime_draft_len) + + def _reshape_logits_for_accept(self, logits, num_contexts, num_gens, + spec_metadata): + """Reshape target logits to the ``[num_contexts + num_gens*(K+1), vocab]`` + layout expected by acceptance. Default is identity (target already emits + one logit per accepted position); workers that emit extra positions + (e.g. PARD's 2K per gen request) override this. + """ + return logits + + def sample_and_accept_draft_tokens(self, logits, attn_metadata, + spec_metadata): + """Sample the golden token and verify previously proposed draft tokens. + + Default implementation for one-model workers whose acceptance differs + only in how draft tokens / target logits are reshaped (DraftTarget, + PARD, DFlash): unpack batch sizes, reshape via the overridable hooks, + then route through ``_accept_draft_tokens`` (strict or rejection). Workers + with a materially different acceptance path (MTP, Eagle3: relaxed / + THOP / extra ``input_ids``) override this method entirely. + """ + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + draft_tokens = self._reshape_draft_tokens_for_accept( + spec_metadata, num_gens, logits.device) + logits = self._reshape_logits_for_accept(logits, num_contexts, num_gens, + spec_metadata) + return self._accept_draft_tokens(logits, draft_tokens, num_contexts, + batch_size, spec_metadata) + + def reset_draft_probs_valid_for_capture(self, spec_metadata): + """Clear the producer->consumer ``draft_probs_valid`` flag at the start + of every draft forward (when rejection is enabled), so it is trusted + only after the current forward finishes scattering its draft probs. This + fails closed across batch-composition changes: a pure-context or + partial/failed capture cannot leave a stale True for the next forward's + acceptance. The worker re-sets it to True only after a full capture. + """ + if getattr(spec_metadata, "use_rejection_sampling", False): + spec_metadata.draft_probs_valid = False + + def _rejection_buffers_valid(self, draft_tokens, draft_len, stored_vocab, + num_contexts, batch_size, logits, + spec_metadata) -> bool: + """Fail-closed guard: return True only when the slot-indexed draft-prob + buffers exist and every shape the rejection path dereferences is valid; + otherwise the caller falls back to strict acceptance. Inspects only + host-side tensor shapes -- no ``.item()`` / value read on CUDA tensors -- + so it stays CUDA-graph-capture safe. + """ + draft_probs = spec_metadata.draft_probs + batch_slot_ids = spec_metadata.batch_slot_ids + if draft_probs is None or batch_slot_ids is None: + return False + if stored_vocab <= 0: + return False + num_gens = batch_size - num_contexts + # draft_probs must cover the slice [:, :draft_len, :stored_vocab]. + if draft_probs.dim() != 3: + return False + if draft_probs.shape[1] < draft_len or draft_probs.shape[ + 2] < stored_vocab: + return False + if draft_tokens.dim() != 2 or draft_tokens.shape[0] != num_gens: + return False + # logits must cover context rows (1 each) + gen rows (draft_len + 1 each). + logits_rows = logits.shape[0] if logits.dim() > 1 else 1 + if logits_rows < num_contexts + num_gens * (draft_len + 1): + return False + # Slot ids for the gen subset must exist (range safety is guaranteed by + # construction). + if batch_slot_ids.shape[0] < batch_size: + return False + if batch_slot_ids[num_contexts:batch_size].shape[0] != num_gens: + return False + return True + def _can_use_rejection_sampling(self, spec_metadata: SpecMetadata) -> bool: # Skip rejection sampling when the whole batch is greedy: the accepted # result is identical to argmax and the base path is cheaper. Mixed @@ -1343,16 +1566,17 @@ def _sample_and_accept_draft_tokens_rejection( assert draft_probs.shape[1] == runtime_draft_len, ( f"draft_probs draft length mismatch: {draft_probs.shape[1]} != " f"{runtime_draft_len}") - d2t = getattr(spec_metadata, "d2t", None) + d2t = self._d2t.data if self._d2t is not None else None if draft_vocab_size != vocab_size: - # Use the pre-allocated buffer from spec_metadata.prepare() - # (zero-filled once at init; untouched positions stay 0). - # Falls back to per-iter allocation if the buffer is not - # configured, e.g. when use_rejection_sampling was off at - # prepare() time. if spec_metadata.full_draft_probs is not None: - full_draft_probs = spec_metadata.full_draft_probs[:num_gens] + # Slice to runtime_draft_len so the max_draft_len buffer + # never passes stale extra rows to the rejection kernel. + full_draft_probs = spec_metadata.full_draft_probs[: + num_gens, : + runtime_draft_len] else: + # Buffer not pre-allocated (e.g. rejection off at prepare()): + # fall back to a per-iter allocation. full_draft_probs = torch.zeros( (num_gens, runtime_draft_len, vocab_size), dtype=torch.float32, @@ -1419,142 +1643,250 @@ def _sample_and_accept_draft_tokens_rejection( spec_metadata=spec_metadata) return accepted_tokens, num_accepted_tokens - def _draft_sampler_greedy(self, logits: torch.Tensor, d2t=None): + def _update_advance_draft_sampling_seed(self, device): + """Increment the draft sampler's RNG seed for this draft-sampling call + (lazily initializing the seed/offset tensors on first use), so each call + samples with a fresh, deterministic seed.""" + if self.seed is None: + self.seed = torch.tensor([0], dtype=torch.int64, device=device) + self.offset = torch.tensor([0], dtype=torch.int64, device=device) + self.seed += 1 + self.seed %= (2**31) + + def _draft_sampler_greedy(self, logits: torch.Tensor): """ Simple greedy draft token sampling using argmax. Args: logits: [num_tokens, vocab_size] - Draft model logits - d2t: Optional dictionary offset tensor for vocab mapping Returns: draft_tokens: [num_tokens] - Sampled draft token ids (int32) """ draft_tokens = greedy(logits, return_probs=False)[0] - # Apply d2t (offsets between draft and target model dictionaries) - if d2t is not None: - draft_tokens = d2t[draft_tokens] + draft_tokens + # Apply the cached draft->target vocab offset map. + if self._d2t is not None: + draft_tokens = self._d2t[draft_tokens] + draft_tokens return draft_tokens.type(torch.int32) - def _draft_sampler_advanced( - self, - logits: torch.Tensor, - spec_metadata: "SpecMetadata", - batch_size: int, - d2t: Optional[torch.Tensor] = None, - ): + def _get_local_max_and_combined(self, logits, mapping_lm_tp=None): + """Pack each rank's local (global_argmax_index, max_value) for a + distributed argmax over a vocab-sharded draft LM head. """ - Draft token sampling using per-request sampling parameters from the - target's sampling config. Falls back to argmax when the batch is - all-greedy. - - Args: - logits: [batch_size, vocab_size] - Draft model logits (one row per - request, since each draft step emits one token per request). - spec_metadata: Source of per-request temperatures / top_k / top_p - tensors populated by populate_sampling_params_for_one_model. - batch_size: Number of active requests in the batch. - d2t: Optional dictionary offset tensor for vocab mapping. - - Returns: - draft_tokens: [batch_size] - Sampled draft token ids (int32) + local_max_values, local_argmax = torch.max(logits, dim=-1, keepdim=True) + vocab_per_rank = logits.shape[-1] + mapping_lm_tp = mapping_lm_tp if mapping_lm_tp is not None else self.mapping + max_index_per_rank = local_argmax.type( + torch.int32) + (mapping_lm_tp.tp_rank * vocab_per_rank) + max_index_per_rank_float = max_index_per_rank.float() + local_max_values_float32 = local_max_values.float() + # Interleaved layout: [idx0, val0, idx1, val1, ...] after all-gather. + combined = torch.stack( + [max_index_per_rank_float, local_max_values_float32], + dim=-1).flatten(-2) + return combined + + @torch.compile(options={"max-autotune": True}) + def _get_draft_tokens_from_gathered(self, gathered): + """Pick the global-argmax token id from the all-gathered per-rank + (index, value) pairs produced by ``_get_local_max_and_combined``. """ - if spec_metadata.is_all_greedy_sample: - return self._draft_sampler_greedy(logits, d2t) - - temperatures = spec_metadata.request_temperatures[:batch_size] - top_ks = spec_metadata.request_top_ks[:batch_size] - top_ps = spec_metadata.request_top_ps[:batch_size] - - if self.seed is None: - self.seed = torch.tensor([0], - dtype=torch.int64, - device=logits.device) - self.offset = torch.tensor([0], - dtype=torch.int64, - device=logits.device) - self.seed += 1 - self.seed %= (2**31) - - draft_tokens = sampling_batch_spec_dec_one_model(logits, - temperatures, - top_ks, - top_ps, - seed=self.seed, - offset=self.offset) - - if d2t is not None: - draft_tokens = d2t[draft_tokens] + draft_tokens - - return draft_tokens.type(torch.int32) - - def _draft_sampler_advanced_for_rejection( - self, - logits: torch.Tensor, - spec_metadata: "SpecMetadata", - batch_size: int, - d2t: Optional[torch.Tensor] = None, - draft_step: int = 0, - ): + gathered_indices_float = gathered[..., 0::2] + gathered_values_float = gathered[..., 1::2] + max_indices = torch.argmax(gathered_values_float, dim=-1, keepdim=True) + draft_tokens = torch.gather(gathered_indices_float, -1, + max_indices).squeeze(-1).type(torch.int32) + return draft_tokens + + def greedy_sample_draft_with_tp_gather(self, + logits: torch.Tensor, + spec_metadata=None, + mapping_lm_head_tp=None): + """Greedy draft-token sampling with a TP all-gather of the argmax. + + When the draft LM head is vocab-sharded under plain tensor parallelism, + a per-rank argmax disagrees across ranks and desyncs speculative + decoding. Gather only each rank's local (index, value) and pick the + global argmax. Falls back to plain argmax when the logits are not + vocab-sharded (see ``_draft_logits_are_sharded``) -- e.g. a borrowed or + gathered full-vocab draft head. Returns tokens in draft-vocab space (the + caller applies d2t). Expects 2D ``[num_tokens, vocab_shard]`` logits. """ - Rejection-sampling-aware variant of ``_draft_sampler_advanced``. - - Single-pass compute + sample + scatter: computes the per-request prob - distribution once via TRT-LLM's fused ``compute_probs_from_logits`` - (temp + top_k + top_p + softmax + greedy override in one CUDA kernel), - samples the draft token from that distribution, and scatters the same - probs into the slot-indexed ``spec_metadata.draft_probs`` buffer for - next-iter rejection verification. Replaces the previous two-stage path - (flashinfer fused sampling kernel + a redundant softmax pass to store - probs). - - All-greedy batches take the cheaper argmax path — - ``_can_use_rejection_sampling`` will bypass rejection for those anyway. + mapping = self.mapping + sharded = self._draft_logits_are_sharded(logits, spec_metadata) + if (sharded and mapping is not None + and getattr(mapping, "tp_size", 1) > 1 + and not mapping.enable_attention_dp): + from ..distributed.ops import allgather + combined = self._get_local_max_and_combined(logits) + gathered = allgather(combined, mapping, dim=-1) + return self._get_draft_tokens_from_gathered(gathered) + elif (sharded and mapping is not None + and getattr(mapping, "tp_size", 1) > 1 + and mapping.enable_lm_head_tp_in_adp): + # ADP + LM-head TP: the worker trimmed the LM-head-TP padding rows, so + # each rank holds [token_count, vocab_shard] for its own tokens. Gather + # the per-rank (index, value) over mapping_lm_head_tp and pick the + # global argmax per row (no token re-slice needed after the trim). + assert mapping_lm_head_tp is not None, ( + "mapping_lm_head_tp is required for ADP + LM-head-TP greedy " + "draft sampling") + from ..distributed.ops import allgather + combined = self._get_local_max_and_combined(logits, + mapping_lm_head_tp) + gathered = allgather(combined, mapping_lm_head_tp, dim=-1) + return self._get_draft_tokens_from_gathered(gathered) + # No TP gather needed: plain argmax (draft-vocab space; caller applies d2t). + return torch.argmax(logits, dim=-1).type(torch.int32) + + def advanced_sample_draft_block(self, gen_logits: torch.Tensor, + spec_metadata: "SpecMetadata", + num_contexts: int, batch_size: int): + """Block counterpart of ``advanced_sample_draft`` for gen-only workers + (PARD, DFLASH): produces all ``K`` draft positions per gen request in one + forward from ``[num_gens, K, vocab]`` logits. + + With rejection enabled, samples via + ``sampling_batch_spec_dec_one_model_for_rejection`` and scatters the K + proposal rows into ``draft_probs[gen_slot_ids, 0:K, :]``; otherwise uses + ``sampling_batch_spec_dec_one_model`` (tokens only). Only called for a + non-greedy batch (the all-greedy path is handled by the caller). Returns + ``[num_gens, K]`` int32 tokens in draft-vocab space (the caller applies + d2t); stored probs likewise stay in draft-vocab space. """ - if spec_metadata.is_all_greedy_sample: - return self._draft_sampler_greedy(logits, d2t) - - temperatures = spec_metadata.request_temperatures[:batch_size] - top_ks = spec_metadata.request_top_ks[:batch_size] - top_ps = spec_metadata.request_top_ps[:batch_size] - - if self.seed is None: - self.seed = torch.tensor([0], - dtype=torch.int64, - device=logits.device) - self.offset = torch.tensor([0], - dtype=torch.int64, - device=logits.device) - self.seed += 1 - self.seed %= (2**31) - - draft_tokens, probs = sampling_batch_spec_dec_one_model_for_rejection( - logits, - temperatures, - top_ks, - top_ps, - seed=self.seed, - offset=self.offset, - ) - - # Scatter probs into the slot-indexed buffer (shaped - # [max_num_requests, max_draft_len, vocab_size]). Each request's data - # always lands at its stable py_seq_slot row regardless of batch - # composition shifts across iterations. - assert spec_metadata.batch_slot_ids is not None, ( - "batch_slot_ids must be populated by " - "populate_sampling_params_for_one_model before draft probs storage") - batch_slots = spec_metadata.batch_slot_ids[:batch_size] - vocab = probs.shape[-1] - spec_metadata.draft_probs[batch_slots, draft_step, :vocab] = probs - spec_metadata.draft_probs_last_dim = vocab - - if d2t is not None: - draft_tokens = d2t[draft_tokens] + draft_tokens - - return draft_tokens.type(torch.int32) + num_gens, K, vocab = gen_logits.shape + if num_gens == 0: + return torch.empty((0, K), + dtype=torch.int32, + device=gen_logits.device) + + # Take the gen slice and repeat each request's value K times to line up + # with the flattened [num_gens*K, vocab] logits (K rows per request). + temps = spec_metadata.request_temperatures[ + num_contexts:batch_size].repeat_interleave(K) + top_ks = spec_metadata.request_top_ks[ + num_contexts:batch_size].repeat_interleave(K) + top_ps = spec_metadata.request_top_ps[ + num_contexts:batch_size].repeat_interleave(K) + + self._update_advance_draft_sampling_seed(gen_logits.device) + flat_logits = gen_logits.reshape(num_gens * K, vocab) + + if getattr(spec_metadata, "use_rejection_sampling", False): + flat_tokens, flat_probs = ( + sampling_batch_spec_dec_one_model_for_rejection( + flat_logits, + temps, + top_ks, + top_ps, + seed=self.seed, + offset=self.offset)) + # Scatter the K prob rows per gen request into its stable slot row. + if spec_metadata.draft_probs is not None: + assert spec_metadata.batch_slot_ids is not None, ( + "batch_slot_ids must be populated before block draft prob " + "storage") + gen_slot_ids = spec_metadata.batch_slot_ids[ + num_contexts:batch_size] + probs = flat_probs.reshape(num_gens, K, vocab) + spec_metadata.draft_probs[gen_slot_ids, :K, :vocab] = probs + spec_metadata.draft_probs_last_dim = vocab + else: + flat_tokens = sampling_batch_spec_dec_one_model(flat_logits, + temps, + top_ks, + top_ps, + seed=self.seed, + offset=self.offset) + + return flat_tokens.reshape(num_gens, K).type(torch.int32) + + def sample_draft_tokens(self, + logits, + spec_metadata, + batch_size, + *, + num_contexts=0, + draft_step=None, + is_last_draft_cycle=False, + mapping_lm_head_tp=None): + """Unified draft-token production entry for all one-model workers. + + Branches by logits rank: 3D ``[num_gens, K, vocab]`` is the block form + (gen-only workers emitting all K positions in one forward, e.g. + PARD/DFlash); 2D ``[num_tokens, vocab]`` is the per-step form + (autoregressive workers called once per draft step, e.g. MTP/ + DraftTarget). ``draft_step`` applies only to the step form and + ``num_contexts`` only to the block form. d2t is read from ``self._d2t``. + + ``is_last_draft_cycle`` marks the final production call of the draft + forward. On that call ``draft_probs_valid`` is set True only when + rejection is enabled, the batch is not all-greedy, and (block form) the + batch is pure generation; otherwise a partial/greedy/mixed capture fails + closed to strict acceptance next iteration. + + In a mixed (context + generation) batch the block form receives + gen-only logits (context requests draft from accepted tokens they do not + have yet), so ``num_contexts`` slices the full-batch per-request metadata + down to the gen segment. The step form receives full-batch logits + (context requests draft from target hidden states already available), so + no slicing is needed. + """ + is_block = logits.dim() == 3 + use_rejection = getattr(spec_metadata, "use_rejection_sampling", False) + + # Draft tokens use argmax unless rejection sampling is engaged for a + # non-greedy batch. Rejection sampling is the only path that needs the + # draft's stochastic proposal distribution (stored in draft_probs); every + # other path (all-greedy, or non-greedy with strict/exact-match + # acceptance) accepts a draft token only when it equals the target's + # choice, and there argmax maximizes the acceptance rate (E[accept] = + # max_i p_i >= sum_i p_i^2 = E[accept] for a stochastic draft). This + # matches sglang/vLLM, which draft with argmax/top-k by default and apply + # sampling params only on the target/acceptance side. + advanced = use_rejection and not spec_metadata.is_all_greedy_sample + + # All samplers below return tokens in draft-vocab space; d2t is applied + # once after the branch. + if not advanced: + # greedy_sample_draft_with_tp_gather expects 2D [tokens, vocab]; + # flatten the block form and restore its shape (step form is 2D). + batch_shape = logits.shape[:-1] + tokens = self.greedy_sample_draft_with_tp_gather( + logits.reshape(-1, logits.shape[-1]), spec_metadata, + mapping_lm_head_tp) + tokens = tokens.reshape(batch_shape) + else: + # Advanced sampling gathers the vocab-sharded draft logits to full + # vocab, then samples (scattering this step's proposal distribution + # into draft_probs). + logits = self.maybe_gather_sharded_draft_logits( + logits, spec_metadata, mapping_lm_head_tp) + if is_block: + tokens = self.advanced_sample_draft_block( + logits, spec_metadata, num_contexts, batch_size).long() + else: + tokens = self.advanced_sample_draft(logits, + spec_metadata, + batch_size, + draft_step=draft_step) + + # Map draft-vocab token ids to target vocab (no-op for shared vocab). + if self._d2t is not None: + tokens = self._d2t[tokens] + tokens + + if is_last_draft_cycle and use_rejection: + # Trust the captured draft_probs only for a non-all-greedy batch + # (and, for the block form, a pure-gen batch: a mixed batch's block + # capture is partial). + valid = not spec_metadata.is_all_greedy_sample + if is_block: + valid = valid and num_contexts == 0 + spec_metadata.draft_probs_valid = valid + + return tokens.type(torch.int32) def _execute_guided_decoder_if_present(self, logits): """Execute guided decoder on target model logits if available.""" diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 6afed0fc02e9..f24d53bb912e 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -5,10 +5,8 @@ import torch from tensorrt_llm._utils import prefer_pinned -from tensorrt_llm.mapping import Mapping from ..attention_backend import AttentionMetadata -from ..distributed.ops import allgather from ..pyexecutor.llm_request import LlmRequest from ..pyexecutor.resource_manager import BaseResourceManager, SlotManager from ..pyexecutor.sampler import TorchSampler @@ -282,6 +280,7 @@ def __init__(self, super().__init__(use_separate_draft_kv_cache) self.spec_config = spec_config self.model_config = model_config + self.mapping = getattr(model_config, "mapping", None) self.is_thop = False self.sa_enhancer: Optional[SADraftEnhancer] = None if spec_config.sa_config is not None: @@ -447,6 +446,12 @@ def forward( last_tokens_idx = torch.cumsum( attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1 + # Rejection sampling captures each draft step's sampling distribution + # into the slot-indexed spec_metadata.draft_probs buffer (engaged only + # for non-greedy batches; all-greedy uses argmax). draft_probs_valid is + # finalized inside sample_draft_tokens on the last draft cycle. + self.reset_draft_probs_valid_for_capture(spec_metadata) + draft_kv_cache_manager = self.get_draft_kv_cache_manager( resource_manager) @@ -469,7 +474,12 @@ def forward( self.guided_decoder.execute_draft_batch(logits, draft_step=i) - new_draft_token = self.draft_sampler(logits) + new_draft_token = self.sample_draft_tokens( + logits, + spec_metadata, + batch_size, + draft_step=i, + is_last_draft_cycle=(i == runtime_draft_len - 1)) next_draft_tokens.append(new_draft_token) # shift input_ids and hidden_states input_ids = draft_inputs["input_ids"] @@ -841,6 +851,16 @@ def sample_and_accept_draft_tokens( runtime_draft_len, spec_metadata=spec_metadata) + # Rejection sampling acceptance. _can_use_rejection_sampling() requires + # use_rejection_sampling, draft_probs_valid, and a non-all-greedy batch; + # otherwise falls through to strict acceptance below. Context rows take + # the target's first sampled token; gen rows run the rejection kernel. + elif self._can_use_rejection_sampling(spec_metadata): + draft_tokens = spec_metadata.draft_tokens.reshape( + num_gens, runtime_draft_len) + accepted_tokens, num_accepted_tokens = self._accept_draft_tokens( + logits, draft_tokens, num_contexts, batch_size, spec_metadata) + # Strict acceptance else: if self.is_thop: @@ -1092,80 +1112,3 @@ def prepare_drafter_inputs( "hidden_states": return_hidden_states, "attn_metadata": attn_metadata, } - - @torch.compile(options={"max-autotune": True}) - def get_local_max_and_combined(self, logits, mapping_lm_tp=None): - local_max_values, local_argmax = torch.max(logits, dim=-1, keepdim=True) - # Adjust indices based on TP rank and size - vocab_per_rank = logits.shape[-1] - mapping_lm_tp = mapping_lm_tp if mapping_lm_tp is not None else self.model_config.mapping - max_index_per_rank = local_argmax.type( - torch.int32) + (mapping_lm_tp.tp_rank * vocab_per_rank) - # Use torch.stack and flatten instead of view+cat to avoid torch.compile issues - # Convert both to float32 to ensure consistent dtype - max_index_per_rank_float = max_index_per_rank.float() - local_max_values_float32 = local_max_values.float() - - # Stack and flatten to get interleaved layout: [idx0, val0, idx1, val1, ...] - combined = torch.stack( - [max_index_per_rank_float, local_max_values_float32], - dim=-1).flatten(-2) - return combined - - @torch.compile(options={"max-autotune": True}) - def get_draft_tokens_from_gathered(self, gathered): - gathered_indices_float = gathered[..., 0::2] # Even positions: indices - gathered_values_float = gathered[..., 1::2] # Odd positions: values - - # Find the rank with maximum value - max_indices = torch.argmax(gathered_values_float, dim=-1, keepdim=True) - - # Get the corresponding token indices and convert back to int32 - draft_tokens = torch.gather(gathered_indices_float, -1, - max_indices).squeeze(-1).type(torch.int32) - return draft_tokens - - def draft_sampler( - self, - logits: torch.Tensor, - mapping_lm_head_tp: Mapping = None, - ): - ''' - Sampling draft tokens. - - Args: - logits: torch.Tensor - [num_tokens, vocab_size] - Logits produced by the draft model. - - Returns: - draft_tokens: torch.Tensor - [batch_size * max_draft_len] - Draft token ids. Flattened. - ''' - if (self.model_config is not None - and hasattr(self.model_config, 'mapping') - and self.model_config.mapping.tp_size - > 1) and not (self.model_config.mapping.enable_attention_dp): - combined = self.get_local_max_and_combined(logits) - gathered = allgather(combined, self.model_config.mapping, dim=-1) - draft_tokens = self.get_draft_tokens_from_gathered(gathered) - elif (self.model_config is not None - and hasattr(self.model_config, 'mapping') - and self.model_config.mapping.tp_size - > 1) and self.model_config.mapping.enable_lm_head_tp_in_adp: - # For ADP + LM head TP mode, we need to find the global argmax across all TP ranks - combined = self.get_local_max_and_combined(logits, - mapping_lm_head_tp) - gathered = allgather(combined, mapping_lm_head_tp, dim=-1) - batch_size = logits.shape[0] - local_batch_size = batch_size // mapping_lm_head_tp.tp_size - gathered = gathered.view(mapping_lm_head_tp.tp_size, - local_batch_size, -1) - sliced_gathered = gathered[mapping_lm_head_tp.tp_rank] - draft_tokens = self.get_draft_tokens_from_gathered(sliced_gathered) - else: - # Simple argmax if no TP or no model config - draft_tokens = self._draft_sampler_greedy(logits) - - return draft_tokens diff --git a/tensorrt_llm/_torch/speculative/pard.py b/tensorrt_llm/_torch/speculative/pard.py index 34a4509314cd..86cb48ac62b2 100644 --- a/tensorrt_llm/_torch/speculative/pard.py +++ b/tensorrt_llm/_torch/speculative/pard.py @@ -196,25 +196,8 @@ def forward( self._execute_guided_decoder_if_present(logits) - # draft_tokens buffer has (2K-1) entries per gen request; extract the K real drafts - if num_gens > 0: - draft_tokens = spec_metadata.draft_tokens[: num_gens * (2 * K - 1)] - draft_tokens = draft_tokens.reshape(num_gens, 2 * K - 1)[:, :K] - else: - draft_tokens = spec_metadata.draft_tokens.reshape(0, K) - - # logits have 2K entries per gen request; extract K+1 for acceptance - if num_gens > 0: - ctx_logits = logits[:num_contexts] - vocab_size = logits.shape[-1] - gen_logits_2k = logits[num_contexts:].reshape(num_gens, 2 * K, vocab_size) - gen_logits_kp1 = gen_logits_2k[:, : K + 1, :].reshape(-1, vocab_size) - logits_for_accept = torch.cat([ctx_logits, gen_logits_kp1], dim=0) - else: - logits_for_accept = logits - - accepted_tokens, num_accepted_tokens = self._sample_and_accept_draft_tokens_base( - logits_for_accept, draft_tokens, num_contexts, batch_size, spec_metadata + accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( + logits, attn_metadata, spec_metadata ) # Pad accepted_tokens from (batch, K+1) to (batch, 2K) to match sampler buffer @@ -255,6 +238,8 @@ def forward( draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + self.reset_draft_probs_valid_for_capture(spec_metadata) + if num_gens > 0: with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): hidden_states_out = draft_model.model(**inputs) @@ -285,14 +270,13 @@ def forward( vocab_size = gen_logits.shape[-1] gen_logits = gen_logits.reshape(num_gens, K, vocab_size) - # Use torch.argmax directly to avoid cute_argmax stride issues - d2t = getattr(draft_model.model, "d2t", None) - gen_draft_tokens = torch.argmax(gen_logits, dim=-1, keepdim=False).long() - - if d2t is not None: - gen_draft_tokens = d2t[gen_draft_tokens] + gen_draft_tokens - - gen_draft_tokens = gen_draft_tokens.type(torch.int32) + gen_draft_tokens = self.sample_draft_tokens( + gen_logits, + spec_metadata, + batch_size, + num_contexts=num_contexts, + is_last_draft_cycle=True, + ) if self.sa_enhancer is not None and sa_manager is not None: gen_draft_tokens = self.sa_enhancer.maybe_override_all_draft_tokens( @@ -347,23 +331,26 @@ def forward( "next_new_tokens": next_new_tokens, } - def draft_decoder( - self, - logits: torch.Tensor, - draft_model: nn.Module, - ): - """ - Sample draft tokens using greedy decoding. - - Args: - logits: [num_tokens, vocab_size] from the draft model. - draft_model: The draft model (used to read the d2t mapping). - - Returns: - draft_tokens: [batch_size * max_draft_len] flattened token ids. - """ - d2t = getattr(draft_model.model, "d2t", None) - return self._draft_sampler_greedy(logits, d2t) + def _reshape_draft_tokens_for_accept(self, spec_metadata, num_gens, device): + # draft_tokens buffer has (2K-1) entries per gen request; extract the K real drafts. + K = spec_metadata.runtime_draft_len + if num_gens > 0: + draft_tokens = spec_metadata.draft_tokens[: num_gens * (2 * K - 1)] + return draft_tokens.reshape(num_gens, 2 * K - 1)[:, :K] + # Context-only batch: return an empty view without reshaping the + # (possibly preallocated, non-empty) draft_tokens buffer. + return torch.empty((0, K), dtype=torch.int32, device=device) + + def _reshape_logits_for_accept(self, logits, num_contexts, num_gens, spec_metadata): + # logits have 2K entries per gen request; extract K+1 for acceptance. + if num_gens == 0: + return logits + K = spec_metadata.runtime_draft_len + ctx_logits = logits[:num_contexts] + vocab_size = logits.shape[-1] + gen_logits_2k = logits[num_contexts:].reshape(num_gens, 2 * K, vocab_size) + gen_logits_kp1 = gen_logits_2k[:, : K + 1, :].reshape(-1, vocab_size) + return torch.cat([ctx_logits, gen_logits_kp1], dim=0) def prepare_1st_drafter_inputs( self, diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index d7a0e8185da5..7738ccf962fb 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -42,6 +42,37 @@ def _is_effective_dynamic_tree(spec_config) -> bool: and getattr(spec_config, 'dynamic_tree_max_topK', 0) > 1) +def _get_draft_vocab_size(spec_config, target_vocab_size: int) -> int: + """Draft-model vocab size, used to decide whether rejection sampling needs + the d2t-expanded ``full_draft_probs`` buffer (only when it differs from the + target vocab). + + Reads the draft model's ``config.json`` from ``spec_config.speculative_model``. + Eagle3 configs store the target vocab in ``vocab_size`` and the reduced head + width in ``draft_vocab_size``, so ``draft_vocab_size`` is read first, falling + back to ``vocab_size`` (or a nested ``text_config.vocab_size``). Returns + ``target_vocab_size`` (shared vocab, no buffer needed) when there is no + separate draft model or the config cannot be read. + """ + draft_dir = getattr(spec_config, "speculative_model", None) + if not draft_dir: + return target_vocab_size + try: + import json + import os + cfg_path = os.path.join(str(draft_dir), "config.json") + if not os.path.isfile(cfg_path): + return target_vocab_size + with open(cfg_path) as f: + cfg = json.load(f) + vs = cfg.get("draft_vocab_size") or cfg.get("vocab_size") + if vs is None: + vs = (cfg.get("text_config") or {}).get("vocab_size") + return int(vs) if vs else target_vocab_size + except (OSError, ValueError, TypeError, AttributeError): + return target_vocab_size + + def get_spec_metadata(spec_config, model_config, max_num_requests, @@ -57,6 +88,9 @@ def get_spec_metadata(spec_config, num_seq_slots = (num_seq_slots if num_seq_slots is not None else max_num_requests) vocab_size = getattr(model_config, "vocab_size", 0) + # Draft-model vocab size, used to gate the d2t-expanded full_draft_probs + # buffer allocation (see SpecMetadata.prepare_rejection_sampling_buffers). + draft_vocab_size = _get_draft_vocab_size(spec_config, vocab_size) if spec_config.spec_dec_mode.is_mtp_eagle_one_model(): # MTP Eagle one-model reuses Eagle3 one-model metadata for the # unified worker/sampler/slot_ids plumbing, but skips per-layer @@ -76,6 +110,7 @@ def get_spec_metadata(spec_config, use_rejection_sampling=use_rejection_sampling, vocab_size=vocab_size, num_seq_slots=num_seq_slots, + draft_vocab_size=draft_vocab_size, spec_resource_manager=spec_resource_manager, ) if spec_config.spec_dec_mode.is_mtp_vanilla(): @@ -86,6 +121,9 @@ def get_spec_metadata(spec_config, mtp_num_modules=spec_config.max_draft_len, max_num_requests=max_num_requests, mtp_hidden_states_manager=spec_resource_manager, + use_rejection_sampling=use_rejection_sampling, + vocab_size=vocab_size, + draft_vocab_size=draft_vocab_size, ) if spec_config.spec_dec_mode.is_mtp_eagle(): return Eagle3SpecMetadata( @@ -134,6 +172,7 @@ def get_spec_metadata(spec_config, layers_to_capture=spec_config.eagle3_layers_to_capture, use_rejection_sampling=use_rejection_sampling, vocab_size=vocab_size, + draft_vocab_size=draft_vocab_size, spec_resource_manager=spec_resource_manager, use_dynamic_tree=_is_effective_dynamic_tree(spec_config), eagle_choices=spec_config.eagle_choices, @@ -145,6 +184,9 @@ def get_spec_metadata(spec_config, spec_dec_mode=spec_config.spec_dec_mode, max_num_requests=max_num_requests, spec_resource_manager=spec_resource_manager, + use_rejection_sampling=use_rejection_sampling, + vocab_size=vocab_size, + draft_vocab_size=draft_vocab_size, ) if spec_config.spec_dec_mode.is_dflash(): target_layer_ids = getattr(spec_config, 'target_layer_ids', None) @@ -157,6 +199,9 @@ def get_spec_metadata(spec_config, hidden_size=model_config.hidden_size, max_num_tokens=max_num_tokens, dtype=model_config.torch_dtype, + use_rejection_sampling=use_rejection_sampling, + vocab_size=vocab_size, + draft_vocab_size=draft_vocab_size, ) if spec_config.spec_dec_mode.is_draft_target_one_model(): return DraftTargetOneModelSpecMetadata( @@ -165,6 +210,9 @@ def get_spec_metadata(spec_config, spec_dec_mode=spec_config.spec_dec_mode, max_num_requests=max_num_requests, max_num_tokens=max_num_tokens, + use_rejection_sampling=use_rejection_sampling, + vocab_size=vocab_size, + draft_vocab_size=draft_vocab_size, ) if spec_config.spec_dec_mode.is_save_hidden_states(): return SaveHiddenStatesSpecMetadata( diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index cc6a66e53a08..115ca8ab1138 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1690,14 +1690,17 @@ def validate_draft_len_schedule_and_sort(cls, v, info): @model_validator(mode='after') def validate_rejection_sampling_config(self): - """Disable rejection sampling when SA-enhanced configurations are - active, since SA may override the proposed draft tokens. This is a - silent fallback so the new default (True) does not break sa_config - users. + """Disable rejection sampling when SA-enhanced configurations are active. + + Only silently disable a default-inherited value; an explicit + ``use_rejection_sampling=True`` is preserved so + ``TorchLlmArgs.validate_speculative_config`` can raise for the + unsupported SA combination. """ if self.use_rejection_sampling and getattr(self, 'sa_config', None) is not None: - self.use_rejection_sampling = False + if "use_rejection_sampling" not in self.model_fields_set: + self.use_rejection_sampling = False return self @model_validator(mode='before') @@ -4963,17 +4966,120 @@ def validate_speculative_config(self): exclude={"decoding_type"}) self.speculative_config = Eagle3DecodingConfig(**eagle_data) - if self.speculative_config.use_rejection_sampling and not isinstance( - self.speculative_config, Eagle3DecodingConfig): - # Rejection sampling is only wired up for Eagle3 one-model paths. - # Silently fall back for other spec types so the new default - # (True) does not break them. - # TODO: extend rejection sampling to the remaining speculative - # decoding paths (MTP / DraftTarget / PARD / DFlash / - # SaveHiddenStates / SA) and unify the dispatch in SpecMetadata - # so new spec algorithms get rejection sampling for free; once - # all paths are covered this whitelist guard can be removed. - self.speculative_config.use_rejection_sampling = False + if self.speculative_config.use_rejection_sampling: + # Supported paths: Eagle3 one-model, MTP-Eagle one-model, + # vanilla MTP, PARD, DFlash, DraftTarget one-model. Classify by + # spec-dec mode; MTP-Eagle one-model shares the Eagle3 one-model + # worker/metadata/sampler, so it rides the same unified + # production/acceptance path. + from tensorrt_llm._torch.speculative.interface import \ + SpeculativeDecodingMode as TorchSpeculativeDecodingMode + + spec_mode = self.speculative_config.spec_dec_mode + is_supported = (isinstance(self.speculative_config, + Eagle3DecodingConfig) + or spec_mode == TorchSpeculativeDecodingMode.MTP + or spec_mode.is_mtp_eagle_one_model() + or spec_mode.is_pard() or spec_mode.is_dflash() + or spec_mode.is_draft_target_one_model()) + + # Combinations that break the proposal-distribution invariant. + # SA is gated for every method; relaxed / parallel / guided gates + # apply only to the newly wired methods (vanilla MTP, PARD, + # DFlash, DraftTarget one-model). + is_new_rejection_method = ( + spec_mode == TorchSpeculativeDecodingMode.MTP + or spec_mode.is_pard() or spec_mode.is_dflash() + or spec_mode.is_draft_target_one_model()) + # Plain tensor parallelism is supported (the draft path + # all-gathers vocab-sharded draft logits before rejection, see + # SpecWorkerBase.maybe_gather_sharded_draft_logits). + # attention-DP and context parallelism remain gated. + rs_parallel_active = (self.context_parallel_size > 1 + or self.enable_attention_dp) + rs_guided_active = self.guided_decoding_backend is not None + rs_sa_active = getattr(self.speculative_config, "sa_config", + None) is not None + rs_relaxed_active = getattr( + self.speculative_config, + "use_relaxed_acceptance_for_thinking", False) + # Eagle3 dynamic-tree rejection records at most kMaxTriedPerLevel + # (=32, cpp/.../speculativeDecoding/dynamicTreeKernels.cu) tried + # siblings per level; dynamic_tree_max_topK > 32 yields a wrong + # correction distribution, so cap supported topK at 32. + _DYNAMIC_TREE_REJECTION_MAX_TOPK = 32 + rs_dynamic_tree_active = getattr(self.speculative_config, + "use_dynamic_tree", False) + rs_dynamic_tree_topk_unsupported = ( + rs_dynamic_tree_active + and (self.speculative_config.dynamic_tree_max_topK + or 0) > _DYNAMIC_TREE_REJECTION_MAX_TOPK) + mtp_unsupported_combo = ( + rs_sa_active or rs_dynamic_tree_topk_unsupported + or (is_new_rejection_method and + (rs_relaxed_active or rs_parallel_active + or rs_guided_active))) + + if not is_supported or mtp_unsupported_combo: + # Explicit opt-in raises; a default-inherited value is + # silently disabled for backward compatibility. + explicitly_set = "use_rejection_sampling" in \ + self.speculative_config.model_fields_set + if explicitly_set: + reasons = [] + if not is_supported: + reasons.append( + f"spec mode {spec_mode.name} is not a supported " + "rejection-sampling path (supported: the Eagle3 " + "one-model path, vanilla MTP, PARD, DFlash, and " + "DraftTarget one-model)") + if rs_sa_active: + reasons.append("SA (sa_config) is active") + if rs_dynamic_tree_topk_unsupported: + reasons.append( + "dynamic_tree_max_topK=" + f"{self.speculative_config.dynamic_tree_max_topK}" + " exceeds the dynamic-tree rejection kernel limit " + f"of {_DYNAMIC_TREE_REJECTION_MAX_TOPK} tried " + "siblings per level") + if is_new_rejection_method: + if rs_relaxed_active: + reasons.append( + "relaxed-thinking acceptance is enabled") + if rs_parallel_active: + reasons.append( + "tensor/context parallelism or attention-DP " + "is active (the draft path resolves only " + "the global argmax, not full distributions)" + ) + if rs_guided_active: + reasons.append("guided decoding is enabled") + raise ValueError( + "use_rejection_sampling=True is not supported for " + f"this configuration: {'; '.join(reasons)}.") + self.speculative_config.use_rejection_sampling = False + elif self.speculative_config.use_rejection_sampling: + # The non-dynamic-tree one-model rejection path depends on + # FlashInfer (sampling_from_probs / + # chain_speculative_sampling); fail fast if it is missing. + # The dynamic-tree path uses + # verify_dynamic_tree_rejection_out / + # compute_probs_from_logits (CUDA op with a PyTorch + # fallback) and does not require FlashInfer. + rs_dynamic_tree_active = getattr(self.speculative_config, + "use_dynamic_tree", False) + if not rs_dynamic_tree_active: + from tensorrt_llm._torch.flashinfer_utils import \ + IS_FLASHINFER_AVAILABLE + if not IS_FLASHINFER_AVAILABLE: + raise ValueError( + "use_rejection_sampling=True requires FlashInfer, " + "which is not available in this environment. " + "Install flashinfer-python or set " + "use_rejection_sampling=False.") + # The enabled paths (Eagle3, vanilla MTP, PARD, DFlash, + # DraftTarget one-model) share the SpecMetadata slot-indexed + # dispatch. if isinstance(self.speculative_config, PARDDecodingConfig): assert self.speculative_config.max_draft_len > 0, "PARD max_draft_len must be > 0" diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 3fc1f846b163..df098b45454b 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -49,6 +49,7 @@ l0_h100: # function's leaves on Hopper, keep the rest of the sampler dir. - unittest/_torch/sampler -k "not test_speculative_d2h_parity_real_predictor" - unittest/_torch/speculative/test_eagle3.py + - unittest/_torch/speculative/test_rejection_buffers_guard.py - unittest/_torch/speculative/hw_agnostic - unittest/_torch/thop/parallel - unittest/_torch/thop/parallel_hw_agnostic diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash.py index 2a3def69a639..bb920e27bfae 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash.py @@ -104,5 +104,35 @@ def test_dflash_qwen3_5_4b(disable_overlap_scheduler: bool): _run_and_check(llm_config, min_avg_accepted=1.0) +@pytest.mark.high_cuda_memory +def test_dflash_qwen3_8b_rejection(): + """DFlash with rejection sampling on: the block-capture rejection path + (draft-prob scatter -> fail-closed guard -> rejection acceptance) runs + end-to-end with non-greedy sampling and produces coherent output.""" + total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 + if total_mem_gb < 35: + pytest.skip("Not enough memory to load target + draft model") + + models_path = llm_models_root() + llm_config = _make_llm_config( + target_model_dir=f"{models_path}/Qwen3/Qwen3-8B", + dflash_model_dir=f"{models_path}/Qwen3-8B-DFlash-b16", + disable_overlap_scheduler=True, + ) + llm_config["speculative_config"].use_rejection_sampling = True + + llm = LLM(**llm_config) + # Non-greedy so rejection sampling actually engages (all-greedy bypasses it). + outputs = llm.generate( + PROMPTS, SamplingParams(max_tokens=32, temperature=0.8, top_p=0.95, top_k=50, seed=1234) + ) + llm.shutdown() + + assert len(outputs) == len(PROMPTS) + for out in outputs: + assert len(out.outputs[0].token_ids) > 0 + assert out.outputs[0].text.strip() + + if __name__ == "__main__": unittest.main() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_target.py b/tests/unittest/_torch/speculative/hw_agnostic/test_draft_target.py index 4bda8318c08c..dd2b3e5cd7a0 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_target.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_draft_target.py @@ -66,5 +66,51 @@ def test_llama_draft_target(use_cuda_graph: bool, attn_backend: str): assert similar(text_spec, text_ref) +@pytest.mark.high_cuda_memory +def test_llama_draft_target_rejection(): + """DraftTarget one-model with rejection sampling on: the rejection path + (draft-prob capture -> fail-closed guard -> rejection acceptance) runs + end-to-end with non-greedy sampling and produces coherent output.""" + total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 + if total_mem_gb < 60: + pytest.skip("Not enough memory to load target model") + + models_path = llm_models_root() + target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" + draft_model_dir = f"{models_path}/llama-3.2-models/Llama-3.2-1B-Instruct" + + spec_config = DraftTargetDecodingConfig( + max_draft_len=4, + speculative_model=draft_model_dir, + use_rejection_sampling=True, + ) + + llm = LLM( + model=target_model_dir, + backend="pytorch", + attn_backend="TRTLLM", + disable_overlap_scheduler=True, + max_batch_size=2, + kv_cache_config=KvCacheConfig(enable_block_reuse=False, max_tokens=8192), + max_num_tokens=2048, + speculative_config=spec_config, + ) + prompts = [ + "The capital of France is", + "The president of the United States is", + ] + # Non-greedy so rejection sampling actually engages (all-greedy bypasses it). + sampling_params = SamplingParams( + max_tokens=32, temperature=0.8, top_p=0.95, top_k=50, seed=1234 + ) + outputs = llm.generate(prompts, sampling_params) + llm.shutdown() + + assert len(outputs) == len(prompts) + for out in outputs: + assert len(out.outputs[0].token_ids) > 0 + assert out.outputs[0].text.strip() + + if __name__ == "__main__": unittest.main() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py index d987ed442b5d..a381dbd94907 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py @@ -1,13 +1,20 @@ +import os +import sys import unittest +import pytest import torch from parameterized import parameterized import tensorrt_llm +from tensorrt_llm import LLM, SamplingParams from tensorrt_llm._torch.attention_backend import TrtllmAttentionMetadata from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.speculative.mtp import MTPHiddenStatesManager, MTPSpecMetadata, MTPWorker -from tensorrt_llm.llmapi import MTPDecodingConfig +from tensorrt_llm.llmapi import KvCacheConfig, MTPDecodingConfig + +sys.path.append(os.path.join(os.path.dirname(__file__), "..")) +from utils.llm_data import llm_models_root def unittest_name_func(testcase_func, param_num, param): @@ -1724,3 +1731,90 @@ def test_prepare_drafter_inputs( torch.testing.assert_close(draft_inputs["input_ids"], ref_input_ids) torch.testing.assert_close(draft_inputs["hidden_states"], ref_previous_hidden_states) + + +@pytest.mark.high_cuda_memory +def test_vanilla_mtp_rejection(): + """Vanilla MTP with rejection sampling on: the per-step rejection path + (draft-prob scatter -> fail-closed guard -> rejection acceptance) runs + end-to-end with non-greedy sampling and produces coherent output.""" + total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 + if total_mem_gb < 60: + pytest.skip("Not enough memory to load DeepSeek-V3-Lite") + + model_dir = f"{llm_models_root()}/DeepSeek-V3-Lite/bf16" + spec_config = MTPDecodingConfig( + max_draft_len=1, + use_mtp_vanilla=True, + speculative_model=model_dir, + use_rejection_sampling=True, + ) + llm = LLM( + model=model_dir, + backend="pytorch", + disable_overlap_scheduler=True, + max_batch_size=2, + kv_cache_config=KvCacheConfig(enable_block_reuse=False, max_tokens=8192), + max_num_tokens=2048, + speculative_config=spec_config, + trust_remote_code=True, + ) + prompts = [ + "The capital of France is", + "The president of the United States is", + ] + # Non-greedy so rejection sampling actually engages (all-greedy bypasses it). + sampling_params = SamplingParams( + max_tokens=32, temperature=0.8, top_p=0.95, top_k=50, seed=1234 + ) + outputs = llm.generate(prompts, sampling_params) + llm.shutdown() + + assert len(outputs) == len(prompts) + for out in outputs: + assert len(out.outputs[0].token_ids) > 0 + assert out.outputs[0].text.strip() + + +@pytest.mark.high_cuda_memory +def test_mtp_eagle_one_model_rejection(): + """MTP-Eagle one-model with rejection sampling on: it shares the Eagle3 + one-model worker/sampler, so the same per-step rejection path runs + end-to-end with non-greedy sampling and produces coherent output.""" + total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 + if total_mem_gb < 60: + pytest.skip("Not enough memory to load DeepSeek-V3-Lite") + + model_dir = f"{llm_models_root()}/DeepSeek-V3-Lite/bf16" + spec_config = MTPDecodingConfig( + max_draft_len=1, + use_mtp_vanilla=False, + mtp_eagle_one_model=True, + speculative_model=model_dir, + use_rejection_sampling=True, + ) + llm = LLM( + model=model_dir, + backend="pytorch", + disable_overlap_scheduler=True, + max_batch_size=2, + kv_cache_config=KvCacheConfig(enable_block_reuse=False, max_tokens=8192), + max_num_tokens=2048, + speculative_config=spec_config, + trust_remote_code=True, + ) + prompts = [ + "The capital of France is", + "The president of the United States is", + ] + # Non-greedy so rejection sampling actually engages (all-greedy bypasses it). + sampling_params = SamplingParams( + max_tokens=32, temperature=0.8, top_p=0.95, top_k=50, seed=1234 + ) + outputs = llm.generate(prompts, sampling_params) + llm.shutdown() + + assert len(outputs) == len(prompts) + for out in outputs: + assert len(out.outputs[0].token_ids) > 0 + assert out.outputs[0].text.strip() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_pard.py b/tests/unittest/_torch/speculative/hw_agnostic/test_pard.py index 543f269cec2e..8f2b5afb5557 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_pard.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_pard.py @@ -76,5 +76,50 @@ def test_pard(disable_overlap_scheduler: bool): llm_spec.shutdown() +@pytest.mark.high_cuda_memory +def test_pard_rejection(): + """PARD with rejection sampling on: the block-capture rejection path + (draft-prob scatter -> fail-closed guard -> rejection acceptance) runs + end-to-end with non-greedy sampling and produces coherent output.""" + total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 + if total_mem_gb < 35: + pytest.skip("Not enough memory to load target + draft model") + + models_path = llm_models_root() + pard_model_dir = f"{models_path}/PARD-Llama-3.2-1B" + target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" + + spec_config = PARDDecodingConfig( + max_draft_len=4, + speculative_model=pard_model_dir, + use_rejection_sampling=True, + ) + llm_spec = LLM( + model=target_model_dir, + attn_backend="TRTLLM", + disable_overlap_scheduler=True, + cuda_graph_config=CudaGraphConfig(batch_sizes=[1, 2], enable_padding=True), + max_batch_size=2, + kv_cache_config=KvCacheConfig(enable_block_reuse=False, max_tokens=2048), + max_seq_len=2048, + speculative_config=spec_config, + ) + prompts = [ + "The capital of France is", + "The president of the United States is", + ] + # Non-greedy so rejection sampling actually engages (all-greedy bypasses it). + sampling_params = SamplingParams( + max_tokens=32, temperature=0.8, top_p=0.95, top_k=50, seed=1234 + ) + outputs = llm_spec.generate(prompts, sampling_params) + llm_spec.shutdown() + + assert len(outputs) == len(prompts) + for out in outputs: + assert len(out.outputs[0].token_ids) > 0 + assert out.outputs[0].text.strip() + + if __name__ == "__main__": unittest.main() diff --git a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py new file mode 100644 index 000000000000..c9e5d301448c --- /dev/null +++ b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py @@ -0,0 +1,275 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU unit tests for the one-model rejection slot-buffer allocation and the +fail-closed acceptance guard. + +These exercise the pure buffer/guard logic with synthetic tensors and lightweight +attribute stand-ins (``types.SimpleNamespace``), so they do not construct a full +``SpecMetadata`` or run any model forward: + +- ``SpecMetadata.prepare_rejection_sampling_buffers`` only reads/writes ``self`` + attributes (``use_rejection_sampling``, ``vocab_size``, ``max_num_requests``, + ``max_draft_len``, ``draft_probs``, ``batch_slot_ids``, ``full_draft_probs``). +- ``SpecWorkerBase._rejection_buffers_valid`` only reads its arguments and + ``spec_metadata`` attributes (no ``self`` use), so it can be called unbound. + +Requires CUDA because the buffers are allocated on ``device='cuda'``. +""" + +import types + +import pytest +import torch + +from tensorrt_llm._torch.speculative.interface import SpecMetadata, SpecWorkerBase + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="rejection buffers are CUDA tensors" +) + +R, K, V = 8, 4, 32 # max_num_requests, max_draft_len, vocab_size + + +def _alloc_meta(**over): + base = dict( + use_rejection_sampling=True, + draft_probs=None, + vocab_size=V, + draft_vocab_size=V, + max_num_requests=R, + max_draft_len=K, + num_seq_slots=0, + batch_slot_ids=None, + full_draft_probs=None, + draft_probs_vocab_size=0, + ) + base.update(over) + return types.SimpleNamespace(**base) + + +def test_prepare_buffers_allocates_when_enabled(): + # Shared draft/target vocab: draft_probs + batch_slot_ids are allocated; + # full_draft_probs is skipped (only needed when the vocabs differ). + m = _alloc_meta() + SpecMetadata.prepare_rejection_sampling_buffers(m) + assert m.draft_probs is not None and tuple(m.draft_probs.shape) == (R + 1, K, V) + assert m.batch_slot_ids is not None and m.batch_slot_ids.shape[0] == R + assert m.batch_slot_ids.dtype == torch.long + assert m.full_draft_probs is None + assert m.draft_probs_vocab_size == V + + +def test_prepare_buffers_allocates_full_draft_probs_on_vocab_mismatch(): + # Distinct draft vocab: full_draft_probs (d2t-expanded) is allocated. + m = _alloc_meta(draft_vocab_size=V - 1) + SpecMetadata.prepare_rejection_sampling_buffers(m) + assert m.full_draft_probs is not None and tuple(m.full_draft_probs.shape) == (R + 1, K, V) + + +def test_prepare_buffers_span_seq_slot_pool(): + # Under DeepSeek-V4 overlap scheduling the SeqSlotManager pool + # (num_seq_slots) can exceed max_num_requests; py_seq_slot then indexes past + # max_num_requests. Slot-indexed buffers must span the full pool plus the + # dummy scratch row, and dummy_slot_row must land on that last row so a real + # slot never aliases the dummy row. + pool = 2 * R + m = _alloc_meta(num_seq_slots=pool) + SpecMetadata.prepare_rejection_sampling_buffers(m) + assert tuple(m.draft_probs.shape) == (pool + 1, K, V) + assert m.dummy_slot_row == pool + + +def test_prepare_buffers_noop_when_disabled(): + m = _alloc_meta(use_rejection_sampling=False) + SpecMetadata.prepare_rejection_sampling_buffers(m) + assert m.draft_probs is None + assert m.batch_slot_ids is None + assert m.full_draft_probs is None + + +def _valid_state(): + return types.SimpleNamespace( + draft_probs=torch.empty((R, K, V), device="cuda"), + batch_slot_ids=torch.arange(R, device="cuda", dtype=torch.long), + ) + + +# num_contexts=0, num_gens=4 +_NUM_CTX = 0 +_BATCH = 4 +_NUM_GENS = _BATCH - _NUM_CTX + + +def _good_args(): + draft_tokens = torch.zeros((_NUM_GENS, K), dtype=torch.int, device="cuda") + logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 1), V), device="cuda") + return draft_tokens, K, V, _NUM_CTX, _BATCH, logits + + +_DUMMY = object() # _rejection_buffers_valid does not use self + + +def _call(meta, draft_tokens, draft_len, stored_vocab, num_contexts, batch_size, logits): + return SpecWorkerBase._rejection_buffers_valid( + _DUMMY, draft_tokens, draft_len, stored_vocab, num_contexts, batch_size, logits, meta + ) + + +def test_guard_true_on_valid_state(): + assert _call(_valid_state(), *_good_args()) is True + + +# Each builder returns (meta, *_call args) for one malformed state that the +# fail-closed guard must reject. Kept as a single parametrized test so the guard's +# rejection branches read as one table. +def _case_draft_probs_missing(): + m = _valid_state() + m.draft_probs = None + return (m, *_good_args()) + + +def _case_batch_slot_ids_missing(): + m = _valid_state() + m.batch_slot_ids = None + return (m, *_good_args()) + + +def _case_stored_vocab_nonpositive(): + dt, dl, _sv, nc, bs, lg = _good_args() + return (_valid_state(), dt, dl, 0, nc, bs, lg) + + +def _case_stored_vocab_exceeds_buffer(): + dt, dl, _sv, nc, bs, lg = _good_args() + return (_valid_state(), dt, dl, V + 1, nc, bs, lg) + + +def _case_draft_len_exceeds_buffer(): + # draft_probs has K steps; ask for K+1. + draft_tokens = torch.zeros((_NUM_GENS, K + 1), dtype=torch.int, device="cuda") + logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 2), V), device="cuda") + return (_valid_state(), draft_tokens, K + 1, V, _NUM_CTX, _BATCH, logits) + + +def _case_draft_tokens_wrong_rows(): + dt = torch.zeros((_NUM_GENS + 1, K), dtype=torch.int, device="cuda") + _, dl, sv, nc, bs, lg = _good_args() + return (_valid_state(), dt, dl, sv, nc, bs, lg) + + +def _case_too_few_logits_rows(): + dt, dl, sv, nc, bs, _lg = _good_args() + too_few = torch.zeros((nc + _NUM_GENS, V), device="cuda") # missing +1 each + return (_valid_state(), dt, dl, sv, nc, bs, too_few) + + +def _case_batch_slot_ids_too_short(): + m = _valid_state() + m.batch_slot_ids = torch.arange(_BATCH - 1, device="cuda", dtype=torch.long) + return (m, *_good_args()) + + +@pytest.mark.parametrize( + "build_case", + [ + _case_draft_probs_missing, + _case_batch_slot_ids_missing, + _case_stored_vocab_nonpositive, + _case_stored_vocab_exceeds_buffer, + _case_draft_len_exceeds_buffer, + _case_draft_tokens_wrong_rows, + _case_too_few_logits_rows, + _case_batch_slot_ids_too_short, + ], + ids=lambda f: f.__name__.removeprefix("_case_"), +) +def test_guard_false_on_malformed_state(build_case): + assert _call(*build_case()) is False + + +# -------------------------------------------------------------------------- +# Fail-closed acceptance dispatch: prove _accept_draft_tokens() routes to +# strict/base acceptance (and clears draft_probs_valid) when the buffers are +# malformed even though draft_probs_valid was left True, and routes to the +# rejection method when the state is valid. The two acceptance methods are +# stubbed to record which path ran; the real _can_use_rejection_sampling and +# _rejection_buffers_valid are exercised. +# -------------------------------------------------------------------------- + + +class _Worker(SpecWorkerBase): + @property + def max_draft_len(self) -> int: + return K + + +def _dispatch_meta(**over): + base = dict( + use_rejection_sampling=True, + draft_probs_valid=True, + is_all_greedy_sample=False, + draft_probs_vocab_size=V, + draft_probs_last_dim=V, + batch_slot_ids=torch.arange(R, device="cuda", dtype=torch.long), + draft_probs=torch.zeros((R, K, V), device="cuda"), + ) + base.update(over) + return types.SimpleNamespace(**base) + + +def _make_worker(): + w = _Worker() + calls = {"base": 0, "rejection": 0} + + def _base(logits, draft_tokens, num_contexts, batch_size, spec_metadata): + calls["base"] += 1 + return ("base", None) + + def _rej(logits, draft_tokens, draft_probs, num_contexts, batch_size, spec_metadata): + calls["rejection"] += 1 + return ("rejection", None) + + w._sample_and_accept_draft_tokens_base = _base + w._sample_and_accept_draft_tokens_rejection = _rej + return w, calls + + +def test_accept_dispatch_fails_closed_on_malformed_buffers(): + w, calls = _make_worker() + # draft_probs_valid is True, but the buffers are malformed (draft_probs None) + # -> _rejection_buffers_valid must return False, so acceptance falls back to + # base, the rejection kernel is skipped, and draft_probs_valid is cleared. + meta = _dispatch_meta(draft_probs=None) + draft_tokens = torch.zeros((_NUM_GENS, K), dtype=torch.int, device="cuda") + logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 1), V), device="cuda") + out = w._accept_draft_tokens(logits, draft_tokens, _NUM_CTX, _BATCH, meta) + assert out == ("base", None) + assert calls["base"] == 1 and calls["rejection"] == 0 + assert meta.draft_probs_valid is False # stale flag cleared on fallback + + +def test_accept_dispatch_routes_to_rejection_on_valid_state(): + w, calls = _make_worker() + meta = _dispatch_meta() # valid buffers + draft_tokens = torch.zeros((_NUM_GENS, K), dtype=torch.int, device="cuda") + logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 1), V), device="cuda") + out = w._accept_draft_tokens(logits, draft_tokens, _NUM_CTX, _BATCH, meta) + assert out == ("rejection", None) + assert calls["rejection"] == 1 and calls["base"] == 0 + + +def test_accept_dispatch_base_when_all_greedy(): + w, calls = _make_worker() + # All-greedy batch: rejection is bypassed regardless of buffers. + meta = _dispatch_meta(is_all_greedy_sample=True) + draft_tokens = torch.zeros((_NUM_GENS, K), dtype=torch.int, device="cuda") + logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 1), V), device="cuda") + out = w._accept_draft_tokens(logits, draft_tokens, _NUM_CTX, _BATCH, meta) + assert out == ("base", None) + assert calls["base"] == 1 and calls["rejection"] == 0 + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) From 8fd477d0115939ffec0926895063361f6284fbce Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Fri, 10 Jul 2026 02:01:01 -0700 Subject: [PATCH 2/8] [TRTLLM-13321][fix] Fix CUDA-graph rejection producer-state hazards Two CUDA-graph correctness fixes for one-model rejection sampling: - Clear draft_probs_valid after graph capture. Warmup runs dummy requests through the full draft forward and sets the producer flag while only the dummy scratch slot is populated; the first real replay's acceptance would otherwise read stale per-slot probabilities. capture() now invalidates it so acceptance fails closed to strict until the real draft forward repopulates the buffers. - Route dummy requests to a scratch row captured at allocation time. create_cuda_graph_metadata shrinks max_num_requests to the graph bucket size while sharing the full-size draft_probs buffer, so deriving the scratch row from max_num_requests at use time aliased a real request's slot and overwrote its draft probabilities. Capture the scratch row against the real allocation size instead. Adds regression tests for both. Signed-off-by: ZhaoyangWang --- .../_torch/pyexecutor/cuda_graph_runner.py | 18 ++++++++ tensorrt_llm/_torch/speculative/interface.py | 22 ++++++--- .../speculative/hw_agnostic/test_pard.py | 45 +++++++++++++++++++ 3 files changed, 80 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index e0bd91ac52a8..7efa6de22c18 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -444,6 +444,24 @@ def _setup_spec_decoding_and_forward(key: KeyType, forward_fn: Callable, self.graph_outputs[key] = make_weak_ref(output) self.memory_pool = graph.pool() + self._invalidate_draft_probs_after_capture( + self.graph_metadata[key].get("spec_metadata")) + + @staticmethod + def _invalidate_draft_probs_after_capture(spec_metadata): + """Clear the one-model rejection producer flag on the captured metadata. + + Warmup ran dummy requests through the full draft forward, which sets + ``draft_probs_valid=True`` while only the dummy scratch slot was + populated. Clearing it here makes the first real replay's acceptance fail + closed to strict acceptance until the real draft forward repopulates the + per-slot draft-prob buffers. No-op when spec metadata is absent or does + not carry the flag (non-rejection paths). + """ + if spec_metadata is not None and hasattr(spec_metadata, + "draft_probs_valid"): + spec_metadata.draft_probs_valid = False + def replay(self, key: KeyType, current_inputs: Dict[str, Any]) -> Optional[torch.Tensor]: """Replays a previously captured graph.""" diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index ed7a2b59d9fa..7d040a7b9d64 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -548,6 +548,13 @@ class SpecMetadata: # Shape: [num_seq_slots, max_draft_len, vocab_size]. draft_probs: Optional[torch.Tensor] = None draft_probs_vocab_size: int = 0 + # Scratch row index that dummy/padding requests (py_seq_slot is None) route + # to, captured when draft_probs is allocated so it tracks the buffer's real + # size. It must NOT be re-derived from max_num_requests at use time: + # create_cuda_graph_metadata shrinks max_num_requests to the graph bucket + # size while sharing the full-size buffer, so a bucket-relative index would + # collide with a real request's slot row and overwrite its draft probs. + dummy_slot_row: int = 0 # Whether draft_probs contains valid data. draft_probs_valid: bool = False # Last dimension size of the draft logits/probs stored in draft_probs. @@ -592,6 +599,11 @@ def prepare_rejection_sampling_buffers(self): dtype=torch.float32, device='cuda') self.draft_probs_vocab_size = self.vocab_size + # Dummy requests route to the extra scratch row (the buffer's last + # row at index ``slot_capacity``). Capture it here against the real + # allocation size so it stays correct on graph copies that shrink + # max_num_requests and under overlap where slot_capacity exceeds it. + self.dummy_slot_row = slot_capacity if self.batch_slot_ids is None and self.max_num_requests > 0: self.batch_slot_ids = torch.empty((self.max_num_requests, ), dtype=torch.long, @@ -738,12 +750,12 @@ def _normalize_request_sampling_params( (temp_val, tk_val, tp_val, num_tokens)) # py_seq_slot is a stable per-request id used to scatter/gather draft # probs across iterations. Dummy/padding requests (py_seq_slot is - # None) route to the scratch row at index ``slot_capacity`` (== - # num_seq_slots or max_num_requests), matching the buffer sizing in - # prepare_rejection_sampling_buffers so a real slot never collides. + # None) route to the scratch row captured at allocation time (the + # buffer's real last row), not max_num_requests, which a graph copy + # shrinks to the bucket size and would alias a real request's slot. per_request_slot_ids.append( - request.py_seq_slot if request.py_seq_slot is not None else - (self.num_seq_slots or self.max_num_requests)) + request.py_seq_slot if request. + py_seq_slot is not None else self.dummy_slot_row) self.skip_temperature = not temperature_enabled self.skip_top_k = not top_k_enabled diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_pard.py b/tests/unittest/_torch/speculative/hw_agnostic/test_pard.py index 8f2b5afb5557..4fee8da04a8c 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_pard.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_pard.py @@ -121,5 +121,50 @@ def test_pard_rejection(): assert out.outputs[0].text.strip() +@pytest.mark.high_cuda_memory +def test_pard_rejection_draft_probs_valid_cleared_after_warmup(enforce_single_worker): + """CUDA-graph warmup runs dummy requests through the draft forward, which sets + the rejection producer flag draft_probs_valid=True on each captured graph + metadata copy while only the dummy scratch slot is populated. capture() must + clear it so the first real replay's acceptance fails closed. Assert every + captured graph's flag is False after warmup (which runs during LLM init, + before any generate()).""" + total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 + if total_mem_gb < 35: + pytest.skip("Not enough memory to load target + draft model") + + models_path = llm_models_root() + pard_model_dir = f"{models_path}/PARD-Llama-3.2-1B" + target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" + + spec_config = PARDDecodingConfig( + max_draft_len=4, + speculative_model=pard_model_dir, + use_rejection_sampling=True, + ) + llm_spec = LLM( + model=target_model_dir, + attn_backend="TRTLLM", + disable_overlap_scheduler=True, + cuda_graph_config=CudaGraphConfig(batch_sizes=[1, 2], enable_padding=True), + max_batch_size=2, + kv_cache_config=KvCacheConfig(enable_block_reuse=False, max_tokens=2048), + max_seq_len=2048, + speculative_config=spec_config, + ) + try: + runner = llm_spec._executor.engine.model_engine.cuda_graph_runner + checked = 0 + for meta in runner.graph_metadata.values(): + spec_metadata = meta.get("spec_metadata") + if spec_metadata is not None and hasattr(spec_metadata, "draft_probs_valid"): + assert spec_metadata.draft_probs_valid is False + checked += 1 + # At least one rejection-carrying graph must have been captured and checked. + assert checked > 0 + finally: + llm_spec.shutdown() + + if __name__ == "__main__": unittest.main() From dcb3998df8135944a7829502ec5b3fdc935ad267 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Sun, 12 Jul 2026 23:32:09 -0700 Subject: [PATCH 3/8] [TRTLLM-13321][feat] Write one-hot draft probs for block-worker context slots Block workers (PARD/DFlash) do not draft context requests (they get a zero placeholder token), leaving their slot-indexed draft_probs rows unwritten. When such a request became a generation request the next iteration, rejection acceptance would read a stale row, which was previously guarded by forcing draft_probs_valid=False for block mixed batches. Fill each context request's slot row with a one-hot placeholder distribution (prob 1.0 at draft-vocab token id 0) at the worker's placeholder-assembly point, so the row is a legal distribution: rejection rejects the placeholder (target_prob(0) ~ 0) and resamples from the target residual, equivalent to strict acceptance. Extracted into SpecWorkerBase.write_context_onehot_draft_probs and shared by PARD and DFlash. This lets block mixed batches keep draft_probs_valid=True (drop the num_contexts==0 clause). Signed-off-by: ZhaoyangWang --- tensorrt_llm/_torch/speculative/dflash.py | 6 +++ tensorrt_llm/_torch/speculative/interface.py | 45 +++++++++++++++++--- tensorrt_llm/_torch/speculative/pard.py | 6 +++ 3 files changed, 52 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 4ef753b78530..fbce2b24c7d6 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -502,6 +502,12 @@ def forward( else: gen_draft_tokens = torch.empty((0, K), dtype=torch.int32, device="cuda") + # Context requests are not drafted by the block worker (zero placeholder + # token); fill their draft-prob slot rows with a one-hot placeholder so + # they are a legal distribution when they become gen requests next iter. + gen_vocab = vocab_size if num_gens > 0 else None + self.write_context_onehot_draft_probs(spec_metadata, num_contexts, num_gens, K, gen_vocab) + if num_contexts > 0 and num_gens > 0: ctx_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") next_draft_tokens = torch.cat([ctx_draft_tokens, gen_draft_tokens], dim=0) diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 7d040a7b9d64..5047ae60f760 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -1815,6 +1815,39 @@ def advanced_sample_draft_block(self, gen_logits: torch.Tensor, return flat_tokens.reshape(num_gens, K).type(torch.int32) + def write_context_onehot_draft_probs(self, + spec_metadata, + num_contexts, + num_gens, + draft_len, + gen_vocab=None): + """Write a one-hot draft-prob distribution (prob 1.0 at draft-vocab token + id 0, the placeholder draft token) into each context request's stable + slot row, so the row is a legal distribution when the context request + becomes a generation request next iteration and rejection acceptance + reads ``draft_probs[slot, :draft_len, :stored_vocab]``. + + Block workers (PARD/DFlash) do not draft context requests (they get a + zero placeholder token), leaving their slot rows unwritten. Mixed + iterations reuse the vocab width the gen scatter just set (``gen_vocab``); + pure-context iterations have no gen logits, so use the buffer's full + draft-vocab width and publish it via ``draft_probs_last_dim`` (acceptance + next iter reads this scalar). No-op unless rejection is enabled and the + slot-indexed buffers exist. Static shapes -> CUDA-graph safe. + """ + if (num_contexts <= 0 + or not getattr(spec_metadata, "use_rejection_sampling", False) + or spec_metadata.draft_probs is None): + return + ctx_slot_ids = spec_metadata.batch_slot_ids[:num_contexts] + if num_gens > 0: + onehot_vocab = gen_vocab # matches the gen scatter's width + else: + onehot_vocab = spec_metadata.draft_probs_vocab_size + spec_metadata.draft_probs_last_dim = onehot_vocab + spec_metadata.draft_probs[ctx_slot_ids, :draft_len, :onehot_vocab] = 0.0 + spec_metadata.draft_probs[ctx_slot_ids, :draft_len, 0] = 1.0 + def sample_draft_tokens(self, logits, spec_metadata, @@ -1890,12 +1923,14 @@ def sample_draft_tokens(self, tokens = self._d2t[tokens] + tokens if is_last_draft_cycle and use_rejection: - # Trust the captured draft_probs only for a non-all-greedy batch - # (and, for the block form, a pure-gen batch: a mixed batch's block - # capture is partial). + # Trust the captured draft_probs for any non-all-greedy batch. For + # the block form, the gen subset is scattered here and context slots + # are filled with a one-hot placeholder distribution at the worker's + # placeholder-assembly point, so every slot the next iteration's gen + # acceptance reads is a legal distribution. valid = not spec_metadata.is_all_greedy_sample - if is_block: - valid = valid and num_contexts == 0 + if is_block and spec_metadata.draft_probs is None: + valid = False spec_metadata.draft_probs_valid = valid return tokens.type(torch.int32) diff --git a/tensorrt_llm/_torch/speculative/pard.py b/tensorrt_llm/_torch/speculative/pard.py index 86cb48ac62b2..54d802319170 100644 --- a/tensorrt_llm/_torch/speculative/pard.py +++ b/tensorrt_llm/_torch/speculative/pard.py @@ -298,6 +298,12 @@ def forward( else: gen_draft_tokens = torch.empty((0, 2 * K - 1), dtype=torch.int32, device="cuda") + # Context requests are not drafted by the block worker (zero placeholder + # token); fill their draft-prob slot rows with a one-hot placeholder so + # they are a legal distribution when they become gen requests next iter. + gen_vocab = vocab_size if num_gens > 0 else None + self.write_context_onehot_draft_probs(spec_metadata, num_contexts, num_gens, K, gen_vocab) + if num_contexts > 0 and num_gens > 0: ctx_draft_tokens = torch.zeros( (num_contexts, 2 * K - 1), dtype=torch.int32, device="cuda" From 5a467be2dcc6afbb0ad87377f8f2ddf5185e41fe Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Mon, 13 Jul 2026 00:33:50 -0700 Subject: [PATCH 4/8] [TRTLLM-13321][fix] One-hot draft probs on dynamic-draft-len zero toggle With one-model rejection sampling and dynamic draft length, runtime_draft_len is batch-global and can toggle K->0->K as batch size fluctuates. On the 0 iteration every gen request goes through skip_drafting, producing no draft tokens and never scattering draft_probs. The next K iteration then runs rejection acceptance (its branch frozen in the captured CUDA graph, so a Python flag reset cannot redirect it) reading the stale draft_probs rows paired with padded-zero draft tokens, silently corrupting output. Mirror the draft-token zero-padding: mark such gen requests host-side in _handle_dynamic_draft_len (pre-pad, current_num_draft_tokens == 0), collect their slots in _prepare_tp_inputs, and eagerly write a one-hot placeholder distribution (draft-vocab token id 0) into their draft_probs rows after spec_metadata.prepare(). The eager write into the stable draft_probs buffer is read by the replayed rejection kernel, which then rejects the placeholder and resamples from the target -- equivalent to strict acceptance. Idempotent with the context-slot one-hot; no-op when rejection is off. Signed-off-by: ZhaoyangWang --- tensorrt_llm/_torch/pyexecutor/llm_request.py | 4 +++ .../_torch/pyexecutor/model_engine.py | 13 ++++++++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 9 ++++++ tensorrt_llm/_torch/speculative/interface.py | 31 +++++++++++++++++++ 4 files changed, 57 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 63bfb318e284..6d7522ec79a8 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -736,6 +736,10 @@ def __init__( self.py_target_probs = None self.py_last_draft_tokens = None self.py_num_accepted_draft_tokens = 0 + # One-model rejection: set per-iteration by _handle_dynamic_draft_len when + # this gen request produced 0 real draft tokens, so _prepare_tp_inputs + # one-hots its stale draft_probs slot. Consumed (and cleared) there. + self.py_needs_onehot_draft_probs = False self.py_num_accepted_draft_tokens_indices = [] self.py_rewind_draft_token_separate_adjustment = 0 self.py_per_pos_drafted = [0] * MAX_SPEC_DECODE_POSITIONS diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 2ec62d5d3cae..00aed5522ae0 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3296,6 +3296,10 @@ def _prepare_tp_inputs( draft_tokens = [] draft_lens = [] gen_request_seq_slots = [] # per generation request + # One-model rejection: slots of gen requests that produced 0 real draft + # tokens this step (marked in _handle_dynamic_draft_len); their stale + # draft_probs rows are one-hot'd after spec_metadata.prepare(). + padding_gen_slots = [] multimodal_params_list = [] mrope_position_ids = [ ] # (start_idx, end_idx, (3,1,L) mrope_pos_ids) per multimodal request @@ -3541,6 +3545,10 @@ def append_cross_attention_state(request: LlmRequest, self.runtime_draft_len) runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1 for request in extend_requests: + if getattr(request, "py_needs_onehot_draft_probs", False): + if request.py_seq_slot is not None: + padding_gen_slots.append(request.py_seq_slot) + request.py_needs_onehot_draft_probs = False # consume once request_ids.append(request.py_request_id) request_accepted_path[ request. @@ -4278,6 +4286,11 @@ def previous_seq_slots_device(): spec_metadata.populate_sampling_params_for_one_model( scheduled_requests.all_requests()) spec_metadata.prepare() + # One-model rejection: one-hot the stale draft_probs rows of gen + # requests that produced no draft tokens this step, so the (possibly + # captured) rejection kernel reads a legal placeholder distribution. + spec_metadata.write_padding_onehot_draft_probs( + padding_gen_slots, self.runtime_draft_len) inputs['spec_metadata'] = spec_metadata if self.enable_attention_dp: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index bf20ed9e495e..ed45c421f290 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3194,8 +3194,17 @@ def _handle_dynamic_draft_len(self, scheduled_batch.batch_size, self.model_engine.max_draft_len) # 2. Pad or truncate draft tokens to the resolved length DRAFT_BUFFER_PAD = 0 # Buffer sentinel, not PARD mask_token_id. + rejection_on = getattr(self.model_engine.spec_config, + "use_rejection_sampling", False) for request in scheduled_batch.generation_requests: current_num_draft_tokens = len(request.py_draft_tokens) + # One-model rejection: a gen request entering with 0 real draft + # tokens produced no draft-prob scatter for its slot last iter, + # so next iter's rejection kernel would read a stale draft_probs + # row. Mark it (pre-pad signal) so _prepare_tp_inputs writes a + # one-hot placeholder row after spec_metadata.prepare(). + request.py_needs_onehot_draft_probs = ( + rejection_on and current_num_draft_tokens == 0) if spec_dec_mode.is_pard(): # special case: PARD carries 2K-1 draft tokens per request runtime_draft_token_buffer_width = ( diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 5047ae60f760..b14b75696680 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -617,6 +617,37 @@ def prepare_rejection_sampling_buffers(self): dtype=torch.float32, device='cuda') + def write_padding_onehot_draft_probs(self, padding_slot_ids, draft_len): + """Write a one-hot draft-prob row (prob 1.0 at draft-vocab token id 0, + the placeholder token) into each padding gen request's stable slot row. + + Padding requests are gen requests that entered this iteration with 0 real + draft tokens (e.g. a runtime_draft_len K->0->K dynamic-draft-len toggle). + Their slot's draft_probs row was never scattered by the draft sampler, so + the next iteration's (possibly CUDA-graph-captured) rejection kernel would + read a stale/uninitialized distribution. Writing a legal one-hot row makes + acceptance reject the placeholder and resample from the target (equivalent + to strict acceptance) for those rows. Written eagerly before graph replay + into the stable draft_probs buffer, so the replayed kernel reads it. + + Idempotent w.r.t. context->gen transitions whose row was already one-hot'd + by write_context_onehot_draft_probs. The width matches the value already + published in draft_probs_last_dim (what acceptance reads), so it is NOT + overwritten here. No-op unless rejection is enabled and slots exist. + Static shapes -> CUDA-graph safe. + """ + if (not padding_slot_ids + or not getattr(self, "use_rejection_sampling", False) + or self.draft_probs is None): + return + onehot_vocab = (self.draft_probs_last_dim if self.draft_probs_last_dim + > 0 else self.draft_probs_vocab_size) + slots = torch.tensor(padding_slot_ids, + dtype=torch.long, + device=self.draft_probs.device) + self.draft_probs[slots, :draft_len, :onehot_vocab] = 0.0 + self.draft_probs[slots, :draft_len, 0] = 1.0 + def prepare(self): """ Hook to be called before the forward step of the model. From 0b59abc2a1f53f28a9fa8ed696b8eba91f9f93da Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Mon, 13 Jul 2026 00:57:27 -0700 Subject: [PATCH 5/8] [TRTLLM-13321][refactor] Remove draft_probs_valid producer flag The draft_probs_valid producer->consumer flag is now redundant: every slot the rejection acceptance reads is guaranteed to hold a legal distribution written by the request's immediately-preceding step. Context slots in block workers are one-hot'd by write_context_onehot_draft_probs; dynamic-draft-len zero-toggle padding slots by write_padding_onehot_draft_probs; all-greedy and pure-context batches are already excluded by is_all_greedy_sample and the num_gens>0 guard; CUDA-graph warmup only touches the dummy scratch row (real slots are written by each request's own context/gen forward before any rejection reads them). Delete the field, its term in _can_use_rejection_sampling, the fail-closed assignment, reset_draft_probs_valid_for_capture and its five callers, the set in sample_draft_tokens (and the now-dead is_last_draft_cycle parameter), and _invalidate_draft_probs_after_capture. _can_use_rejection_sampling now reduces to `use_rejection_sampling and not is_all_greedy_sample`, matching the dynamic-tree worker's own gate. The _rejection_buffers_valid shape guard is retained. Drop the obsolete warmup-flag test; update the dispatch tests. Signed-off-by: ZhaoyangWang --- .../_torch/pyexecutor/cuda_graph_runner.py | 18 ------- tensorrt_llm/_torch/speculative/dflash.py | 3 -- .../_torch/speculative/draft_target.py | 3 -- tensorrt_llm/_torch/speculative/eagle3.py | 2 - tensorrt_llm/_torch/speculative/interface.py | 35 +------------ tensorrt_llm/_torch/speculative/mtp.py | 22 +++----- tensorrt_llm/_torch/speculative/pard.py | 3 -- .../speculative/hw_agnostic/test_pard.py | 51 ------------------- .../test_rejection_buffers_guard.py | 11 ++-- 9 files changed, 12 insertions(+), 136 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 7efa6de22c18..e0bd91ac52a8 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -444,24 +444,6 @@ def _setup_spec_decoding_and_forward(key: KeyType, forward_fn: Callable, self.graph_outputs[key] = make_weak_ref(output) self.memory_pool = graph.pool() - self._invalidate_draft_probs_after_capture( - self.graph_metadata[key].get("spec_metadata")) - - @staticmethod - def _invalidate_draft_probs_after_capture(spec_metadata): - """Clear the one-model rejection producer flag on the captured metadata. - - Warmup ran dummy requests through the full draft forward, which sets - ``draft_probs_valid=True`` while only the dummy scratch slot was - populated. Clearing it here makes the first real replay's acceptance fail - closed to strict acceptance until the real draft forward repopulates the - per-slot draft-prob buffers. No-op when spec metadata is absent or does - not carry the flag (non-rejection paths). - """ - if spec_metadata is not None and hasattr(spec_metadata, - "draft_probs_valid"): - spec_metadata.draft_probs_valid = False - def replay(self, key: KeyType, current_inputs: Dict[str, Any]) -> Optional[torch.Tensor]: """Replays a previously captured graph.""" diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index fbce2b24c7d6..9c90628d8708 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -462,8 +462,6 @@ def forward( draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) - self.reset_draft_probs_valid_for_capture(spec_metadata) - if num_gens > 0: with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): hidden_states_out = draft_model.dflash_forward( @@ -496,7 +494,6 @@ def forward( spec_metadata, batch_size, num_contexts=num_contexts, - is_last_draft_cycle=True, ) else: diff --git a/tensorrt_llm/_torch/speculative/draft_target.py b/tensorrt_llm/_torch/speculative/draft_target.py index 47c51187e617..6df31b2bd223 100644 --- a/tensorrt_llm/_torch/speculative/draft_target.py +++ b/tensorrt_llm/_torch/speculative/draft_target.py @@ -213,8 +213,6 @@ def forward( next_draft_tokens = [] original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens - self.reset_draft_probs_valid_for_capture(spec_metadata) - # Get the draft KV cache manager if using separate layouts draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) @@ -266,7 +264,6 @@ def forward( spec_metadata, batch_size, draft_step=i, - is_last_draft_cycle=(i == runtime_draft_len - 1), ) next_draft_tokens.append(new_draft_token) diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index f69147fe7c73..e9612e5d3d19 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -805,7 +805,6 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, runtime_draft_len = spec_metadata.runtime_draft_len num_gens = batch_size - num_contexts next_draft_tokens = [] - self.reset_draft_probs_valid_for_capture(spec_metadata) last_tokens_idx = torch.cumsum( attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1 position_ids = inputs["position_ids"] @@ -917,7 +916,6 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, spec_metadata, batch_size, draft_step=i, - is_last_draft_cycle=(i == runtime_draft_len - 1), mapping_lm_head_tp=mapping_lm_head_tp) next_draft_tokens.append(new_draft_token) diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index b14b75696680..2669b3fb5481 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -555,8 +555,6 @@ class SpecMetadata: # size while sharing the full-size buffer, so a bucket-relative index would # collide with a real request's slot row and overwrite its draft probs. dummy_slot_row: int = 0 - # Whether draft_probs contains valid data. - draft_probs_valid: bool = False # Last dimension size of the draft logits/probs stored in draft_probs. draft_probs_last_dim: int = 0 # Per-request slot ids (py_seq_slot) for the current batch, in batch order. @@ -1313,7 +1311,7 @@ def _accept_draft_tokens(self, logits, draft_tokens, num_contexts, spec_metadata.draft_probs_vocab_size) # Fail closed: run the rejection kernel only when every buffer is # present and correctly shaped; otherwise fall back to strict - # acceptance. draft_probs_valid is reset so stale state is not reused. + # acceptance. if self._rejection_buffers_valid(draft_tokens, draft_len, stored_vocab, num_contexts, batch_size, logits, spec_metadata): @@ -1326,7 +1324,6 @@ def _accept_draft_tokens(self, logits, draft_tokens, num_contexts, return self._sample_and_accept_draft_tokens_rejection( logits, draft_tokens, draft_probs, num_contexts, batch_size, spec_metadata) - spec_metadata.draft_probs_valid = False return self._sample_and_accept_draft_tokens_base( logits, draft_tokens, num_contexts, batch_size, spec_metadata) @@ -1480,17 +1477,6 @@ def sample_and_accept_draft_tokens(self, logits, attn_metadata, return self._accept_draft_tokens(logits, draft_tokens, num_contexts, batch_size, spec_metadata) - def reset_draft_probs_valid_for_capture(self, spec_metadata): - """Clear the producer->consumer ``draft_probs_valid`` flag at the start - of every draft forward (when rejection is enabled), so it is trusted - only after the current forward finishes scattering its draft probs. This - fails closed across batch-composition changes: a pure-context or - partial/failed capture cannot leave a stale True for the next forward's - acceptance. The worker re-sets it to True only after a full capture. - """ - if getattr(spec_metadata, "use_rejection_sampling", False): - spec_metadata.draft_probs_valid = False - def _rejection_buffers_valid(self, draft_tokens, draft_len, stored_vocab, num_contexts, batch_size, logits, spec_metadata) -> bool: @@ -1533,7 +1519,6 @@ def _can_use_rejection_sampling(self, spec_metadata: SpecMetadata) -> bool: # batches (context + gen) are handled via slot-indexed draft probs and # are split inside _sample_and_accept_draft_tokens_rejection. return (spec_metadata.use_rejection_sampling - and spec_metadata.draft_probs_valid and not spec_metadata.is_all_greedy_sample) def _sample_and_accept_draft_tokens_rejection( @@ -1886,7 +1871,6 @@ def sample_draft_tokens(self, *, num_contexts=0, draft_step=None, - is_last_draft_cycle=False, mapping_lm_head_tp=None): """Unified draft-token production entry for all one-model workers. @@ -1897,12 +1881,6 @@ def sample_draft_tokens(self, DraftTarget). ``draft_step`` applies only to the step form and ``num_contexts`` only to the block form. d2t is read from ``self._d2t``. - ``is_last_draft_cycle`` marks the final production call of the draft - forward. On that call ``draft_probs_valid`` is set True only when - rejection is enabled, the batch is not all-greedy, and (block form) the - batch is pure generation; otherwise a partial/greedy/mixed capture fails - closed to strict acceptance next iteration. - In a mixed (context + generation) batch the block form receives gen-only logits (context requests draft from accepted tokens they do not have yet), so ``num_contexts`` slices the full-batch per-request metadata @@ -1953,17 +1931,6 @@ def sample_draft_tokens(self, if self._d2t is not None: tokens = self._d2t[tokens] + tokens - if is_last_draft_cycle and use_rejection: - # Trust the captured draft_probs for any non-all-greedy batch. For - # the block form, the gen subset is scattered here and context slots - # are filled with a one-hot placeholder distribution at the worker's - # placeholder-assembly point, so every slot the next iteration's gen - # acceptance reads is a legal distribution. - valid = not spec_metadata.is_all_greedy_sample - if is_block and spec_metadata.draft_probs is None: - valid = False - spec_metadata.draft_probs_valid = valid - return tokens.type(torch.int32) def _execute_guided_decoder_if_present(self, logits): diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index f24d53bb912e..f3086c33e32f 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -446,12 +446,6 @@ def forward( last_tokens_idx = torch.cumsum( attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1 - # Rejection sampling captures each draft step's sampling distribution - # into the slot-indexed spec_metadata.draft_probs buffer (engaged only - # for non-greedy batches; all-greedy uses argmax). draft_probs_valid is - # finalized inside sample_draft_tokens on the last draft cycle. - self.reset_draft_probs_valid_for_capture(spec_metadata) - draft_kv_cache_manager = self.get_draft_kv_cache_manager( resource_manager) @@ -474,12 +468,10 @@ def forward( self.guided_decoder.execute_draft_batch(logits, draft_step=i) - new_draft_token = self.sample_draft_tokens( - logits, - spec_metadata, - batch_size, - draft_step=i, - is_last_draft_cycle=(i == runtime_draft_len - 1)) + new_draft_token = self.sample_draft_tokens(logits, + spec_metadata, + batch_size, + draft_step=i) next_draft_tokens.append(new_draft_token) # shift input_ids and hidden_states input_ids = draft_inputs["input_ids"] @@ -852,9 +844,9 @@ def sample_and_accept_draft_tokens( spec_metadata=spec_metadata) # Rejection sampling acceptance. _can_use_rejection_sampling() requires - # use_rejection_sampling, draft_probs_valid, and a non-all-greedy batch; - # otherwise falls through to strict acceptance below. Context rows take - # the target's first sampled token; gen rows run the rejection kernel. + # use_rejection_sampling and a non-all-greedy batch; otherwise falls + # through to strict acceptance below. Context rows take the target's + # first sampled token; gen rows run the rejection kernel. elif self._can_use_rejection_sampling(spec_metadata): draft_tokens = spec_metadata.draft_tokens.reshape( num_gens, runtime_draft_len) diff --git a/tensorrt_llm/_torch/speculative/pard.py b/tensorrt_llm/_torch/speculative/pard.py index 54d802319170..969add54136e 100644 --- a/tensorrt_llm/_torch/speculative/pard.py +++ b/tensorrt_llm/_torch/speculative/pard.py @@ -238,8 +238,6 @@ def forward( draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) - self.reset_draft_probs_valid_for_capture(spec_metadata) - if num_gens > 0: with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): hidden_states_out = draft_model.model(**inputs) @@ -275,7 +273,6 @@ def forward( spec_metadata, batch_size, num_contexts=num_contexts, - is_last_draft_cycle=True, ) if self.sa_enhancer is not None and sa_manager is not None: diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_pard.py b/tests/unittest/_torch/speculative/hw_agnostic/test_pard.py index 4fee8da04a8c..6e2c805b54f9 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_pard.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_pard.py @@ -12,12 +12,6 @@ sys.path.append(os.path.join(os.path.dirname(__file__), "..")) -@pytest.fixture(scope="function") -def enforce_single_worker(monkeypatch): - monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") - yield - - @pytest.mark.parametrize("disable_overlap_scheduler", [True, False]) def test_pard(disable_overlap_scheduler: bool): """Test PARD speculative decoding with CUDA graph support. @@ -121,50 +115,5 @@ def test_pard_rejection(): assert out.outputs[0].text.strip() -@pytest.mark.high_cuda_memory -def test_pard_rejection_draft_probs_valid_cleared_after_warmup(enforce_single_worker): - """CUDA-graph warmup runs dummy requests through the draft forward, which sets - the rejection producer flag draft_probs_valid=True on each captured graph - metadata copy while only the dummy scratch slot is populated. capture() must - clear it so the first real replay's acceptance fails closed. Assert every - captured graph's flag is False after warmup (which runs during LLM init, - before any generate()).""" - total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 - if total_mem_gb < 35: - pytest.skip("Not enough memory to load target + draft model") - - models_path = llm_models_root() - pard_model_dir = f"{models_path}/PARD-Llama-3.2-1B" - target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" - - spec_config = PARDDecodingConfig( - max_draft_len=4, - speculative_model=pard_model_dir, - use_rejection_sampling=True, - ) - llm_spec = LLM( - model=target_model_dir, - attn_backend="TRTLLM", - disable_overlap_scheduler=True, - cuda_graph_config=CudaGraphConfig(batch_sizes=[1, 2], enable_padding=True), - max_batch_size=2, - kv_cache_config=KvCacheConfig(enable_block_reuse=False, max_tokens=2048), - max_seq_len=2048, - speculative_config=spec_config, - ) - try: - runner = llm_spec._executor.engine.model_engine.cuda_graph_runner - checked = 0 - for meta in runner.graph_metadata.values(): - spec_metadata = meta.get("spec_metadata") - if spec_metadata is not None and hasattr(spec_metadata, "draft_probs_valid"): - assert spec_metadata.draft_probs_valid is False - checked += 1 - # At least one rejection-carrying graph must have been captured and checked. - assert checked > 0 - finally: - llm_spec.shutdown() - - if __name__ == "__main__": unittest.main() diff --git a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py index c9e5d301448c..e2cbe8752d7a 100644 --- a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py +++ b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py @@ -189,8 +189,7 @@ def test_guard_false_on_malformed_state(build_case): # -------------------------------------------------------------------------- # Fail-closed acceptance dispatch: prove _accept_draft_tokens() routes to -# strict/base acceptance (and clears draft_probs_valid) when the buffers are -# malformed even though draft_probs_valid was left True, and routes to the +# strict/base acceptance when the buffers are malformed, and routes to the # rejection method when the state is valid. The two acceptance methods are # stubbed to record which path ran; the real _can_use_rejection_sampling and # _rejection_buffers_valid are exercised. @@ -206,7 +205,6 @@ def max_draft_len(self) -> int: def _dispatch_meta(**over): base = dict( use_rejection_sampling=True, - draft_probs_valid=True, is_all_greedy_sample=False, draft_probs_vocab_size=V, draft_probs_last_dim=V, @@ -236,16 +234,15 @@ def _rej(logits, draft_tokens, draft_probs, num_contexts, batch_size, spec_metad def test_accept_dispatch_fails_closed_on_malformed_buffers(): w, calls = _make_worker() - # draft_probs_valid is True, but the buffers are malformed (draft_probs None) - # -> _rejection_buffers_valid must return False, so acceptance falls back to - # base, the rejection kernel is skipped, and draft_probs_valid is cleared. + # Buffers are malformed (draft_probs None) -> _rejection_buffers_valid must + # return False, so acceptance falls back to base and the rejection kernel is + # skipped. meta = _dispatch_meta(draft_probs=None) draft_tokens = torch.zeros((_NUM_GENS, K), dtype=torch.int, device="cuda") logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 1), V), device="cuda") out = w._accept_draft_tokens(logits, draft_tokens, _NUM_CTX, _BATCH, meta) assert out == ("base", None) assert calls["base"] == 1 and calls["rejection"] == 0 - assert meta.draft_probs_valid is False # stale flag cleared on fallback def test_accept_dispatch_routes_to_rejection_on_valid_state(): From a004f30322803c5183cb1b7f7c2063421a1c27de Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Tue, 14 Jul 2026 19:56:22 -0700 Subject: [PATCH 6/8] [TRTLLM-13321][doc] Explain why NGram/SA drafting cannot use rejection sampling NGram and SA are retrieval-based drafters that emit only draft token ids with no per-token proposal distribution q(x). Rejection sampling needs q to form min(1, p/q) and the (p - q)+ residual, so both are gated. Document the reason on SuffixAutomatonManager and in the validate_speculative_config support whitelist. Signed-off-by: ZhaoyangWang --- tensorrt_llm/_torch/speculative/suffix_automaton.py | 11 +++++++++++ tensorrt_llm/llmapi/llm_args.py | 7 +++++++ 2 files changed, 18 insertions(+) diff --git a/tensorrt_llm/_torch/speculative/suffix_automaton.py b/tensorrt_llm/_torch/speculative/suffix_automaton.py index 1eb7be0bcea5..12bb77208fe4 100644 --- a/tensorrt_llm/_torch/speculative/suffix_automaton.py +++ b/tensorrt_llm/_torch/speculative/suffix_automaton.py @@ -87,6 +87,17 @@ class SuffixAutomatonManager(BaseResourceManager): and are CUDA graph compatible. Used as the resource manager for both NGram and MTP+SA speculative decoding. + + Rejection sampling is not supported for NGram or SA drafting. Both are + retrieval-based drafters: they propose draft tokens by matching against + previously seen context (n-gram lookup / suffix-automaton traversal) instead + of sampling from a model head, so they emit only token ids with no per-token + proposal distribution ``q(x)`` over the vocabulary. Rejection sampling's + acceptance test needs ``q`` to form the ratio ``min(1, p(x)/q(x))`` and the + residual ``(p - q)+`` correction on rejection; without ``q`` the target + distribution cannot be preserved. The gate lives in + ``TorchLlmArgs.validate_speculative_config`` (NGram falls outside the + supported-mode whitelist; SA is rejected via ``rs_sa_active``). """ def __init__( diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 115ca8ab1138..96bb64a02f3e 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -4972,6 +4972,13 @@ def validate_speculative_config(self): # spec-dec mode; MTP-Eagle one-model shares the Eagle3 one-model # worker/metadata/sampler, so it rides the same unified # production/acceptance path. + # + # Every supported path runs a neural draft head that yields a + # per-token proposal distribution q(x). Rejection sampling needs + # q to form min(1, p/q) and the (p - q)+ residual correction, so + # retrieval-based drafters that emit only token ids with no q are + # excluded: NGram is simply absent from this whitelist, and SA + # (sa_config) is rejected below via rs_sa_active. from tensorrt_llm._torch.speculative.interface import \ SpeculativeDecodingMode as TorchSpeculativeDecodingMode From 96b73b6a5634f08e55402ae714ffcc23bec5ef64 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 15 Jul 2026 04:25:55 -0700 Subject: [PATCH 7/8] [TRTLLM-13321][fix] Wire TP Mapping into MTP one-model workers get_spec_worker holds the correct Mapping but never passed it to the MTP workers: MTPEagleWorker hardcoded super().__init__(mapping=None) and MTPWorker read the unreliable model_config.mapping (None for the MTP worker). With the draft LM head kept vocab-sharded (lm_head_gather_output =False), a None mapping made greedy_sample_draft_with_tp_gather skip the TP argmax gather, so each rank argmaxed its own vocab shard, drafting divergent tokens and desyncing the ranks into a PyExecutor hang under plain TP + overlap scheduler (repro: DeepSeek-V3-Lite fp8 tp4 mtp_nextn=2 cuda_graph+overlap+torch_compile). Thread mapping from get_spec_worker into MTPWorker and MTPEagleWorker, matching every other one-model worker. Verified tp4 output is now byte-identical to the origin/main baseline. Signed-off-by: ZhaoyangWang --- tensorrt_llm/_torch/speculative/eagle3.py | 6 ++++-- tensorrt_llm/_torch/speculative/mtp.py | 12 ++++++++++-- tensorrt_llm/_torch/speculative/utils.py | 11 ++++++++--- 3 files changed, 22 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index e9612e5d3d19..d9b5be1ecc1e 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -1206,10 +1206,12 @@ class MTPEagleWorker(Eagle3OneModelWorker): def __init__(self, spec_config, model_config: Optional[ModelConfig] = None, - use_separate_draft_kv_cache: bool = False): + use_separate_draft_kv_cache: bool = False, + *, + mapping: Optional[Mapping] = None): super().__init__( spec_config, - mapping=None, + mapping=mapping, model_config=model_config, use_separate_draft_kv_cache=use_separate_draft_kv_cache) # Preserved for callers/tests that still expect this attribute. diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index f3086c33e32f..fcf0603a04c1 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -276,11 +276,19 @@ class MTPWorker(SpecWorkerBase): def __init__(self, spec_config: "MTPDecodingConfig", model_config=None, - use_separate_draft_kv_cache: bool = False): + use_separate_draft_kv_cache: bool = False, + *, + mapping=None): super().__init__(use_separate_draft_kv_cache) self.spec_config = spec_config self.model_config = model_config - self.mapping = getattr(model_config, "mapping", None) + # Use the mapping passed by get_spec_worker (the same Mapping that shards + # the draft LM head). model_config.mapping is unreliable here -- for the + # MTP worker it can be None -- so reading it caused greedy draft sampling + # to skip the TP argmax gather and desync the ranks into a hang under + # plain TP + overlap scheduler. + self.mapping = mapping if mapping is not None else getattr( + model_config, "mapping", None) self.is_thop = False self.sa_enhancer: Optional[SADraftEnhancer] = None if spec_config.sa_config is not None: diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 7738ccf962fb..4f9c5af8846a 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -444,10 +444,15 @@ def get_spec_worker(spec_config, use_separate_draft_kv_cache: bool = False): spec_dec_mode = spec_config.spec_dec_mode if spec_dec_mode.is_mtp_vanilla(): - return MTPWorker(spec_config, model_config, use_separate_draft_kv_cache) + return MTPWorker(spec_config, + model_config, + use_separate_draft_kv_cache, + mapping=mapping) if spec_dec_mode.is_mtp_eagle_one_model(): - return MTPEagleWorker(spec_config, model_config, - use_separate_draft_kv_cache) + return MTPEagleWorker(spec_config, + model_config, + use_separate_draft_kv_cache, + mapping=mapping) if spec_dec_mode.is_eagle3_one_model(): if _is_effective_dynamic_tree(spec_config): return Eagle3OneModelDynamicTreeWorker(spec_config, mapping, From f4334c9f757a6a8600173b516a251a91b6a0e6ab Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 15 Jul 2026 19:50:16 -0700 Subject: [PATCH 8/8] [TRTLLM-13321][fix] Do not TP-gather draft logits under attention DP Under attention DP each rank is data-parallel and owns a distinct set of requests, so the draft logits must not be gathered across ranks -- a per-rank argmax on the rank own logits is the correct proposal. The ADP + LM-head-TP path wrongly classified its vocab-sharded draft logits as needing the TP gather, splicing in an allgather over ranks that hold different token counts. That desynced the ranks and crashed in the DeepEP MoE dispatch (intranode.cu unspecified launch failure) / hung the cuda-graph batch-padding allgather (repro: DeepSeek-V3-Lite tp4 ep4 mtp_nextn=3 attention_dp + enable_lm_head_tp_in_adp, matching CI TestDeepSeekR1 test_nvfp4_multi_gpus[latency_adp_lmtp_tp4]). Make _draft_logits_are_sharded return False for any attention-DP mode so greedy takes a plain argmax and the advanced path skips the gather, matching the pre-rework behavior. Remove the now-unreachable ADP+LM-head-TP gather branches. Verified: ADP+LMTP tp4 ep4 and plain-TP tp4 both produce correct output byte-matching the baseline. Signed-off-by: ZhaoyangWang --- tensorrt_llm/_torch/speculative/interface.py | 45 ++++++++------------ 1 file changed, 18 insertions(+), 27 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 2669b3fb5481..275298a200c6 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -1346,9 +1346,16 @@ def _draft_logits_are_sharded(self, logits, spec_metadata): mapping = self.mapping if mapping is None or getattr(mapping, "tp_size", 1) <= 1: return False - # Plain attention DP (without LM-head TP) replicates full-vocab logits. - if (getattr(mapping, "enable_attention_dp", False) - and not getattr(mapping, "enable_lm_head_tp_in_adp", False)): + # Under attention DP every rank is data-parallel -- it owns a distinct + # set of requests -- so the draft logits must NOT be cross-rank gathered, + # whether they are replicated full-vocab (plain ADP) or vocab-sharded + # (ADP + LM-head TP). A per-rank argmax on the rank's own logits is the + # correct proposal (verified later); a gather would splice in a + # mismatched collective across ranks that hold different token counts, + # desyncing them into a hang / DeepEP launch failure. Only plain TP + # (tp>1 without ADP), where all ranks share the same tokens sharded over + # the vocab dim, needs the gather. + if getattr(mapping, "enable_attention_dp", False): return False draft_full_vocab = (getattr(spec_metadata, "draft_vocab_size", 0) or getattr(spec_metadata, "vocab_size", 0) or 0) @@ -1375,15 +1382,11 @@ def maybe_gather_sharded_draft_logits(self, or not self._draft_logits_are_sharded(logits, spec_metadata)): return logits + # Only plain TP (no attention DP) reaches here -- ADP variants return + # early via _draft_logits_are_sharded, since their ranks are + # data-parallel and must not gather draft logits across ranks. from ..distributed.ops import allgather - mapping = self.mapping - if (getattr(mapping, "enable_attention_dp", False) - and getattr(mapping, "enable_lm_head_tp_in_adp", False)): - assert mapping_lm_head_tp is not None, ( - "mapping_lm_head_tp is required to gather ADP + LM-head-TP draft " - "logits") - return allgather(logits, mapping_lm_head_tp, dim=-1) - return allgather(logits, mapping, dim=-1) + return allgather(logits, self.mapping, dim=-1) def advanced_sample_draft(self, logits: torch.Tensor, @@ -1751,22 +1754,10 @@ def greedy_sample_draft_with_tp_gather(self, combined = self._get_local_max_and_combined(logits) gathered = allgather(combined, mapping, dim=-1) return self._get_draft_tokens_from_gathered(gathered) - elif (sharded and mapping is not None - and getattr(mapping, "tp_size", 1) > 1 - and mapping.enable_lm_head_tp_in_adp): - # ADP + LM-head TP: the worker trimmed the LM-head-TP padding rows, so - # each rank holds [token_count, vocab_shard] for its own tokens. Gather - # the per-rank (index, value) over mapping_lm_head_tp and pick the - # global argmax per row (no token re-slice needed after the trim). - assert mapping_lm_head_tp is not None, ( - "mapping_lm_head_tp is required for ADP + LM-head-TP greedy " - "draft sampling") - from ..distributed.ops import allgather - combined = self._get_local_max_and_combined(logits, - mapping_lm_head_tp) - gathered = allgather(combined, mapping_lm_head_tp, dim=-1) - return self._get_draft_tokens_from_gathered(gathered) - # No TP gather needed: plain argmax (draft-vocab space; caller applies d2t). + # No TP gather for attention-DP (incl. ADP + LM-head TP): each rank owns + # its own requests, so a per-rank argmax is the correct proposal and a + # cross-rank gather here would desync the ranks (see + # _draft_logits_are_sharded). Plain argmax; caller applies d2t. return torch.argmax(logits, dim=-1).type(torch.int32) def advanced_sample_draft_block(self, gen_logits: torch.Tensor,