From 051edebb761177165b477835167762b73cdd3114 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Mon, 20 Apr 2026 20:26:29 -0700 Subject: [PATCH 01/17] [None][spec] add rejection sampling support for one-model speculative decoding Signed-off-by: ZhaoyangWang --- .../_torch/speculative/draft_target.py | 3 +- tensorrt_llm/_torch/speculative/eagle3.py | 16 +- tensorrt_llm/_torch/speculative/interface.py | 200 +++++++++++++++++- tensorrt_llm/_torch/speculative/mtp.py | 14 +- .../_torch/speculative/one_model_sampler.py | 64 +++++- tensorrt_llm/_torch/speculative/pard.py | 9 +- tensorrt_llm/_torch/speculative/utils.py | 11 + tensorrt_llm/llmapi/llm_args.py | 8 + 8 files changed, 316 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/draft_target.py b/tensorrt_llm/_torch/speculative/draft_target.py index c026b6d5b290..d14394452095 100644 --- a/tensorrt_llm/_torch/speculative/draft_target.py +++ b/tensorrt_llm/_torch/speculative/draft_target.py @@ -63,6 +63,7 @@ def __post_init__(self): def prepare(self): """Prepare the metadata before model forward.""" + super().prepare() assert self.request_ids is not None # Update batch indices num_seqs = len(self.request_ids) @@ -314,7 +315,7 @@ def sample_and_accept_draft_tokens( else: draft_tokens = spec_metadata.draft_tokens.reshape(num_gens, self.max_draft_len) - return self._sample_and_accept_draft_tokens_base( + return self._accept_draft_tokens( logits, draft_tokens, num_contexts, batch_size, spec_metadata ) diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index b7f95df29c81..828bb4b868cc 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -221,6 +221,7 @@ def __post_init__(self): self.is_spec_dec_dynamic_tree = False def prepare(self): + super().prepare() is_first_draft = self.eagle3_resource_manager.is_first_draft spec_tree_manager = self.eagle3_resource_manager.spec_tree_manager # Update start indices @@ -408,6 +409,7 @@ def is_layer_capture(self, layer_id: int): return layer_id in self.layers_to_capture def prepare(self): + super().prepare() assert self.request_ids is not None # update batch indices num_seqs = len(self.request_ids) @@ -652,6 +654,7 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, """Original linear draft loop (1 token per layer).""" runtime_draft_len = spec_metadata.runtime_draft_len next_draft_tokens = [] + draft_logits_list = [] position_ids = inputs["position_ids"] with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): @@ -704,6 +707,9 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, d2t, draft_step=i) + if spec_metadata.use_rejection_sampling: + draft_logits_list.append(logits.clone()) + new_draft_token = self.draft_decoder(logits, draft_model) next_draft_tokens.append(new_draft_token) # update inputs @@ -745,6 +751,12 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, gen_draft_tokens) next_draft_tokens[num_contexts:] = gen_draft_tokens + if spec_metadata.use_rejection_sampling and draft_logits_list: + d2t_param = getattr(draft_model.model, "d2t", None) + spec_metadata.d2t = d2t_param.data if d2t_param is not None else None + self._compute_and_store_draft_probs(draft_logits_list, + spec_metadata, batch_size) + return next_draft_tokens def sample_and_accept_draft_tokens( @@ -765,8 +777,8 @@ def sample_and_accept_draft_tokens( spec_metadata.runtime_draft_len, dtype=torch.int, device=logits.device) - return self._sample_and_accept_draft_tokens_base( - logits, draft_tokens, num_contexts, batch_size, spec_metadata) + return self._accept_draft_tokens(logits, draft_tokens, num_contexts, + batch_size, spec_metadata) def draft_decoder( self, diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index d39ace8988cd..e4586b2142eb 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -25,6 +25,10 @@ if IS_FLASHINFER_AVAILABLE: import flashinfer +from .one_model_sampler import (compute_probs_from_logits, + rejection_sampling_one_model, + sampling_batch_spec_dec_one_model) + # Environment variable name for forcing the number of accepted tokens in speculative decoding FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR = "TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS" @@ -434,10 +438,29 @@ class SpecMetadata: # For non-greedy sampling on 1-model. allow_advanced_sampling: bool = False + # Whether to use rejection sampling for one-model speculative decoding. + use_rejection_sampling: bool = False # Sampling parameters for non-greedy sampling (per-request) temperatures: Optional[torch.Tensor] = None top_ks: Optional[torch.Tensor] = None top_ps: Optional[torch.Tensor] = None + # Sampling parameters indexed per request. + request_temperatures: Optional[torch.Tensor] = None + request_top_ks: Optional[torch.Tensor] = None + request_top_ps: Optional[torch.Tensor] = None + # Whether to use sampling parameters when sampling draft tokens. + use_sampling_params_for_draft_tokens: bool = False + # Vocab size used for draft_probs buffer allocation. + vocab_size: int = 0 + # Draft probabilities buffer for rejection sampling, stored flat. + draft_probs: Optional[torch.Tensor] = None + draft_probs_vocab_size: 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 + # Draft-to-target vocab offset tensor. + d2t: Optional[torch.Tensor] = None def __post_init__(self): pass @@ -446,6 +469,14 @@ def prepare(self): """ Hook to be called before the forward step of the model. """ + if (self.use_rejection_sampling and self.draft_probs is None + and self.vocab_size > 0): + buffer_size = (self.max_num_requests * self.max_draft_len * + self.vocab_size) + self.draft_probs = torch.empty(buffer_size, + dtype=torch.float32, + device='cuda') + self.draft_probs_vocab_size = self.vocab_size def create_cuda_graph_metadata(self, max_batch_size: int): """ @@ -494,6 +525,9 @@ def populate_sampling_params_for_one_model( temperatures = [] top_ks = [] top_ps = [] + request_temperatures = [] + request_top_ks = [] + request_top_ps = [] # Need to use a very small value for temperature when disabled to avoid division by 0 DISABLE_TEMP_VAL = 1e-5 @@ -525,6 +559,9 @@ def populate_sampling_params_for_one_model( tk_val = DISABLE_TOPK_VAL if is_greedy or tk_val is None or tk_val <= 0 else tk_val tp_val = DISABLE_TOPP_VAL if is_greedy or tp_val is None else tp_val + request_temperatures.append(temp_val) + request_top_ks.append(tk_val) + request_top_ps.append(tp_val) temperatures.extend(temp_val for _ in range(num_tokens)) top_ks.extend(tk_val for _ in range(num_tokens)) top_ps.extend(tp_val for _ in range(num_tokens)) @@ -542,6 +579,15 @@ def populate_sampling_params_for_one_model( (self.max_draft_len + 1) * self.max_num_requests, dtype=torch.float32, device='cuda') + self.request_temperatures = torch.ones(self.max_num_requests, + dtype=torch.float32, + device='cuda') + self.request_top_ks = torch.zeros(self.max_num_requests, + dtype=torch.int32, + device='cuda') + self.request_top_ps = torch.ones(self.max_num_requests, + dtype=torch.float32, + device='cuda') self.temperatures[:len(temperatures)].copy_(torch.tensor( temperatures, dtype=torch.float32, pin_memory=prefer_pinned()), @@ -552,6 +598,17 @@ def populate_sampling_params_for_one_model( self.top_ps[:len(top_ps)].copy_(torch.tensor( top_ps, dtype=torch.float32, pin_memory=prefer_pinned()), non_blocking=True) + self.request_temperatures[:len(request_temperatures)].copy_( + torch.tensor(request_temperatures, + dtype=torch.float32, + pin_memory=prefer_pinned()), + non_blocking=True) + self.request_top_ks[:len(request_top_ks)].copy_(torch.tensor( + request_top_ks, dtype=torch.int32, pin_memory=prefer_pinned()), + non_blocking=True) + self.request_top_ps[:len(request_top_ps)].copy_(torch.tensor( + request_top_ps, dtype=torch.float32, pin_memory=prefer_pinned()), + non_blocking=True) class SpecWorkerBase(nn.Module, ABC): @@ -862,6 +919,107 @@ def _sample_and_accept_draft_tokens_base( return accepted_tokens, num_accepted_tokens + def _accept_draft_tokens(self, logits, draft_tokens, num_contexts, + batch_size, spec_metadata): + """ + Accept draft tokens with optional rejection sampling support. + """ + if self._can_use_rejection_sampling(spec_metadata, num_contexts): + draft_len = draft_tokens.shape[1] + stored_vocab = (spec_metadata.draft_probs_last_dim + if spec_metadata.draft_probs_last_dim > 0 else + spec_metadata.draft_probs_vocab_size) + draft_probs = spec_metadata.draft_probs[:batch_size * draft_len * + stored_vocab].reshape( + batch_size, draft_len, + stored_vocab) + return self._sample_and_accept_draft_tokens_rejection( + logits, draft_tokens, draft_probs, batch_size, spec_metadata) + return self._sample_and_accept_draft_tokens_base( + logits, draft_tokens, num_contexts, batch_size, spec_metadata) + + def _can_use_rejection_sampling(self, spec_metadata: SpecMetadata, + num_contexts: int) -> bool: + return (spec_metadata.use_rejection_sampling + and spec_metadata.draft_probs_valid and num_contexts == 0) + + def _sample_and_accept_draft_tokens_rejection( + self, + logits: torch.Tensor, + draft_tokens: torch.Tensor, + draft_probs: torch.Tensor, + batch_size: int, + spec_metadata, + ): + """ + Rejection-sampling acceptance for one-model speculative decoding. + """ + device = logits.device + vocab_size = logits.shape[-1] + + if logits.dim() == 1: + logits = logits.unsqueeze(0) + + draft_vocab_size = draft_probs.shape[-1] + num_target_tokens = batch_size * (self.max_draft_len + 1) + + temperatures = spec_metadata.temperatures[:num_target_tokens] + top_ks = spec_metadata.top_ks[:num_target_tokens] + top_ps = spec_metadata.top_ps[:num_target_tokens] + + target_probs_flat = compute_probs_from_logits(logits.clone(), + temperatures, top_ks, + top_ps) + target_probs = target_probs_flat.reshape(batch_size, + self.max_draft_len + 1, + vocab_size) + + runtime_draft_len = draft_probs.shape[1] + d2t = getattr(spec_metadata, "d2t", None) + if draft_vocab_size != vocab_size: + full_draft_probs = torch.zeros( + (batch_size, self.max_draft_len, vocab_size), + dtype=torch.float32, + device=device) + if d2t is not None: + assert d2t.numel() == draft_vocab_size, ( + f"d2t size mismatch: {d2t.numel()} != {draft_vocab_size}") + d2t = d2t.to(device=device) + source_indices = torch.arange(draft_vocab_size, + device=device, + dtype=torch.long) + target_indices = (source_indices + d2t) % vocab_size + full_draft_probs[:, :runtime_draft_len, + target_indices] = draft_probs + else: + assert draft_vocab_size < vocab_size + full_draft_probs[:, :runtime_draft_len, :draft_vocab_size] = ( + draft_probs) + else: + full_draft_probs = draft_probs + + full_draft_tokens = draft_tokens.to(torch.int32).contiguous() + + if self.seed is None: + self.seed = torch.tensor([0], dtype=torch.int64, device=device) + if self.offset is None: + self.offset = torch.tensor([0], dtype=torch.int64, device=device) + self.seed += 1 + self.seed %= 2**31 + + accepted_tokens, num_accepted_tokens = rejection_sampling_one_model( + draft_probs=full_draft_probs, + draft_token_ids=full_draft_tokens, + target_probs=target_probs, + deterministic=True, + seed=self.seed, + offset=self.offset, + ) + + num_accepted_tokens = self._apply_force_accepted_tokens( + num_accepted_tokens, 0, draft_tokens.shape[1]) + return accepted_tokens, num_accepted_tokens + def _draft_sampler_greedy(self, logits: torch.Tensor, d2t=None): """ Simple greedy draft token sampling using argmax. @@ -881,6 +1039,46 @@ def _draft_sampler_greedy(self, logits: torch.Tensor, d2t=None): return draft_tokens.type(torch.int32) + def _compute_and_store_draft_probs( + self, + draft_logits_list: List[torch.Tensor], + spec_metadata: SpecMetadata, + batch_size: int, + ): + """ + Compute draft probabilities and store them for next-step rejection sampling. + """ + draft_tokens_per_request = len(draft_logits_list) + vocab_size = draft_logits_list[0].shape[-1] + device = draft_logits_list[0].device + + draft_logits = torch.stack(draft_logits_list, dim=0) + draft_logits_flat = draft_logits.transpose(0, 1).reshape(-1, vocab_size) + + num_draft_tokens = batch_size * draft_tokens_per_request + if spec_metadata.temperatures is not None: + draft_temps = spec_metadata.temperatures[:batch_size].repeat_interleave( + draft_tokens_per_request) + draft_top_ks = spec_metadata.top_ks[:batch_size].repeat_interleave( + draft_tokens_per_request + ) if spec_metadata.top_ks is not None else None + draft_top_ps = spec_metadata.top_ps[:batch_size].repeat_interleave( + draft_tokens_per_request + ) if spec_metadata.top_ps is not None else None + else: + draft_temps = torch.ones(num_draft_tokens, device=device) + draft_top_ks = None + draft_top_ps = None + + draft_probs_flat = compute_probs_from_logits(draft_logits_flat, + draft_temps, draft_top_ks, + draft_top_ps) + num_elements = batch_size * draft_tokens_per_request * vocab_size + spec_metadata.draft_probs[:num_elements].copy_( + draft_probs_flat.flatten()) + spec_metadata.draft_probs_last_dim = vocab_size + spec_metadata.draft_probs_valid = True + def _execute_guided_decoder_if_present(self, logits): """Execute guided decoder on target model logits if available.""" if self.guided_decoder is not None: @@ -1011,8 +1209,6 @@ def _sample_tokens_for_batch( sampled_tokens: [num_tokens] - Sampled token ids """ if spec_metadata.allow_advanced_sampling: - from .one_model_sampler import sampling_batch_spec_dec_one_model - num_gens = batch_size - num_contexts num_tokens = num_contexts + num_gens * ( spec_metadata.runtime_draft_len + 1) diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index c310cf57816c..cf1814cdbf6b 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -172,6 +172,7 @@ def all_rank_num_seqs(self, value: List[int]): self.subseq_all_rank_num_tokens = value def prepare(self): + super().prepare() assert self.request_ids is not None num_seqs = len(self.request_ids) # update batch indices @@ -841,7 +842,7 @@ def sample_and_accept_draft_tokens( num_gens, mtp_num_modules) # Use base implementation for strict acceptance - accepted_tokens, num_accepted_tokens = self._sample_and_accept_draft_tokens_base( + accepted_tokens, num_accepted_tokens = self._accept_draft_tokens( logits, draft_tokens, num_contexts, batch_size, spec_metadata) @@ -1191,6 +1192,7 @@ def forward( # Predict draft tokens next_draft_tokens = [] + draft_logits_list = [] with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): for i in range(self.mtp_num_modules): if i == 0: @@ -1254,6 +1256,9 @@ def forward( self.guided_decoder.execute_draft_batch(logits, draft_step=i) + if spec_metadata.use_rejection_sampling: + draft_logits_list.append(logits[last_tokens_idx].clone()) + if self.model_config.mapping.enable_attention_dp and \ getattr(self.model_config.mapping, 'enable_lm_head_tp_in_adp', False): mapping_lm_head_tp = draft_model.mtp_layers[ @@ -1329,6 +1334,13 @@ def forward( stacked[num_contexts:] = gen_draft_tokens next_draft_tokens = [stacked[:, i] for i in range(stacked.shape[1])] + if spec_metadata.use_rejection_sampling and draft_logits_list: + d2t_inner = getattr(draft_model, "model", draft_model) + d2t_param = getattr(d2t_inner, "d2t", None) + spec_metadata.d2t = d2t_param.data if d2t_param is not None else None + self._compute_and_store_draft_probs(draft_logits_list, + spec_metadata, batch_size) + next_draft_tokens, next_new_tokens = self._prepare_next_tokens( next_draft_tokens, accepted_tokens, spec_metadata, batch_size, num_accepted_tokens) diff --git a/tensorrt_llm/_torch/speculative/one_model_sampler.py b/tensorrt_llm/_torch/speculative/one_model_sampler.py index 7d49aa85dd17..9949a6f588ee 100644 --- a/tensorrt_llm/_torch/speculative/one_model_sampler.py +++ b/tensorrt_llm/_torch/speculative/one_model_sampler.py @@ -1,7 +1,14 @@ -from typing import Optional +from typing import Optional, Tuple import torch -from flashinfer.sampling import top_k_top_p_sampling_from_logits + +from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE + +if IS_FLASHINFER_AVAILABLE: + from flashinfer.sampling import chain_speculative_sampling, top_k_top_p_sampling_from_logits +else: + chain_speculative_sampling = None + top_k_top_p_sampling_from_logits = None def forward_native( @@ -92,6 +99,59 @@ def sampling_batch_spec_dec_one_model( """ logits = apply_temperature(logits, temperatures) if use_flashinfer: + if top_k_top_p_sampling_from_logits is None: + raise RuntimeError("FlashInfer sampling requested but flashinfer is unavailable") return top_k_top_p_sampling_from_logits(logits, top_k, top_p, seed=seed, offset=offset) random_sampled = forward_native(logits, top_k, top_p) return random_sampled + + +def compute_probs_from_logits( + logits: torch.Tensor, + temperatures: torch.Tensor, + top_k: Optional[torch.Tensor], + top_p: Optional[torch.Tensor], +) -> torch.Tensor: + """ + Compute probabilities from logits with temperature, top-k, and top-p applied. + """ + logits = apply_temperature(logits, temperatures) + logits = apply_top_k_top_p(logits, top_k, top_p) + return logits.softmax(dim=-1, dtype=torch.float32) + + +def rejection_sampling_one_model( + draft_probs: torch.Tensor, + draft_token_ids: torch.Tensor, + target_probs: torch.Tensor, + deterministic: bool = True, + seed: Optional[int] = None, + offset: Optional[int] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + CUDA-graph compatible rejection sampling for one-model speculative decoding. + """ + batch_size = draft_token_ids.shape[0] + device = draft_token_ids.device + + output_accepted_token_num = torch.zeros(batch_size, dtype=torch.int32, device=device) + output_emitted_draft_token_num = torch.zeros(batch_size, dtype=torch.int32, device=device) + + if chain_speculative_sampling is None: + raise RuntimeError( + "Rejection sampling for one-model speculative decoding requires flashinfer" + ) + + accepted_tokens, _, output_emitted_draft_token_num = chain_speculative_sampling( + draft_probs, + draft_token_ids, + target_probs, + maybe_output_accepted_token_num=output_accepted_token_num, + maybe_output_emitted_draft_token_num=output_emitted_draft_token_num, + deterministic=deterministic, + generator=None, + seed=seed, + offset=offset, + ) + + return accepted_tokens, output_emitted_draft_token_num + 1 diff --git a/tensorrt_llm/_torch/speculative/pard.py b/tensorrt_llm/_torch/speculative/pard.py index 628de48022d7..53c6a0e7cb09 100644 --- a/tensorrt_llm/_torch/speculative/pard.py +++ b/tensorrt_llm/_torch/speculative/pard.py @@ -36,6 +36,7 @@ def __post_init__(self): self.is_spec_dec_dynamic_tree = False def prepare(self): + super().prepare() assert self.request_ids is not None num_seqs = len(self.request_ids) @@ -210,7 +211,7 @@ def forward( else: logits_for_accept = logits - accepted_tokens, num_accepted_tokens = self._sample_and_accept_draft_tokens_base( + accepted_tokens, num_accepted_tokens = self._accept_draft_tokens( logits_for_accept, draft_tokens, num_contexts, batch_size, spec_metadata ) @@ -297,6 +298,12 @@ def forward( gen_draft_tokens ) + if spec_metadata.use_rejection_sampling: + draft_logits_list = [gen_logits[:, i, :] for i in range(self.max_draft_len)] + d2t_param = getattr(draft_model.model, "d2t", None) + spec_metadata.d2t = d2t_param.data if d2t_param is not None else None + self._compute_and_store_draft_probs(draft_logits_list, spec_metadata, num_gens) + # Pad from (num_gens, K) to (num_gens, 2K-1). if K > 1: pad = torch.zeros((num_gens, K - 1), dtype=torch.int32, device="cuda") diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 4243f27bfe7f..0be688b2030b 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -40,6 +40,9 @@ def get_spec_metadata(spec_config, spec_resource_manager=None, is_draft_model=False, max_seq_len=262144): + use_rejection_sampling = getattr(spec_config, "use_rejection_sampling", + False) + vocab_size = getattr(model_config, "vocab_size", 0) if spec_config.spec_dec_mode.is_mtp_one_model(): return MTPSpecMetadata( max_draft_len=spec_config.max_draft_len, @@ -49,6 +52,8 @@ def get_spec_metadata(spec_config, max_num_requests=max_num_requests, mtp_hidden_states_manager=spec_resource_manager, allow_advanced_sampling=spec_config.allow_advanced_sampling, + use_rejection_sampling=use_rejection_sampling, + vocab_size=vocab_size, ) if spec_config.spec_dec_mode.is_mtp_eagle(): return Eagle3SpecMetadata( @@ -95,6 +100,8 @@ def get_spec_metadata(spec_config, max_num_tokens=max_num_tokens, layers_to_capture=spec_config.eagle3_layers_to_capture, allow_advanced_sampling=spec_config.allow_advanced_sampling, + use_rejection_sampling=use_rejection_sampling, + vocab_size=vocab_size, spec_resource_manager=spec_resource_manager, use_dynamic_tree=spec_config.use_dynamic_tree, eagle_choices=spec_config.eagle_choices, @@ -106,6 +113,8 @@ def get_spec_metadata(spec_config, spec_dec_mode=spec_config.spec_dec_mode, max_num_requests=max_num_requests, allow_advanced_sampling=spec_config.allow_advanced_sampling, + use_rejection_sampling=use_rejection_sampling, + vocab_size=vocab_size, spec_resource_manager=spec_resource_manager, ) if spec_config.spec_dec_mode.is_dflash(): @@ -129,6 +138,8 @@ def get_spec_metadata(spec_config, max_num_requests=max_num_requests, max_num_tokens=max_num_tokens, allow_advanced_sampling=spec_config.allow_advanced_sampling, + use_rejection_sampling=use_rejection_sampling, + vocab_size=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 3db16e5cabfb..be9aa97c08ea 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -914,6 +914,14 @@ class DecodingBaseConfig(StrictBaseModel): "to 1-model code paths; non-greedy sampling is always enabled on 2-model paths." ) + use_rejection_sampling: bool = Field( + default=False, + status="prototype", + description= + "If true, enables rejection sampling for one-model speculative decoding paths. " + "This is intended for non-greedy sampling configurations on the PyTorch backend." + ) + # If set, drafting is allowed to use chain drafter. _allow_chain_drafter: bool = PrivateAttr(True) # If set, drafting uses greedy sampling, irrespective of sampling parameters. From 674059046c6222fb0763fd39d68d5747cadc9e9b Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Mon, 20 Apr 2026 23:43:20 -0700 Subject: [PATCH 02/17] [None][spec] add dynamic-tree rejection sampling support for Eagle3 one-model decoding Signed-off-by: ZhaoyangWang --- .../speculativeDecoding/dynamicTreeKernels.cu | 1116 ++++++++++++++++- .../speculativeDecoding/dynamicTreeKernels.h | 77 ++ cpp/tensorrt_llm/thop/dynamicTreeOp.cpp | 227 ++++ .../_torch/pyexecutor/model_engine.py | 19 + .../_torch/speculative/dynamic_tree_ops.py | 173 ++- .../_torch/speculative/eagle3_dynamic_tree.py | 225 +++- tensorrt_llm/_torch/speculative/interface.py | 102 +- .../_torch/speculative/one_model_sampler.py | 19 +- 8 files changed, 1902 insertions(+), 56 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu index 1dcfdde71671..400daee9f21d 100644 --- a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu +++ b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu @@ -15,17 +15,685 @@ * limitations under the License. */ +#ifndef CUDART_VERSION +#error CUDART_VERSION Undefined! +#elif (CUDART_VERSION >= 11050) +#include +#else +#include "3rdparty/cub/cub.cuh" +#endif + #include "dynamicTreeKernels.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" - +#include "tensorrt_llm/common/reduceKernelUtils.cuh" +#include "tensorrt_llm/kernels/decodingCommon.h" +#include +#include +#include +#include TRTLLM_NAMESPACE_BEGIN +using namespace tensorrt_llm::common; using namespace tensorrt_llm::runtime; namespace kernels::speculative_decoding { +// --------------------------------------------------------------------------- +// Two-stage top-k / top-p masking kernels +// Mirrors the approach in invokeBatchTopKSampling (samplingTopKKernels.cu), +// but outputs a masked logits tensor instead of sampling a token. +// --------------------------------------------------------------------------- + +// Stage 1: Parallel top-k reduction across BLOCKS_PER_BEAM_ blocks per row. +// Each block handles (vocabSize / BLOCKS_PER_BEAM_) elements and finds its +// local top-k, writing (global_index, logit_value) pairs into the tmp buffers. +template +__global__ void topKProbStage1(T const* __restrict__ logits, T* tmpLogProbs, int32_t* topKTmpIdBuf, T* topKTmpValBuf, + int32_t maxTopK, int32_t const* topKs, int32_t vocabSize) +{ + typedef cub::BlockReduce, BLOCK_SIZE_> BlockReduce; + __shared__ typename BlockReduce::TempStorage tempStorage; + + auto const tid = static_cast(threadIdx.x); + auto const bid = static_cast(blockIdx.x); + auto const rowId = bid / BLOCKS_PER_BEAM_; + auto const blockLane = bid % BLOCKS_PER_BEAM_; // chunk index within the row + + auto const k = (topKs != nullptr) ? topKs[rowId] : maxTopK; + + bool const IS_FP16 = std::is_same::value; + T const MAX_T_VAL = IS_FP16 ? HALF_FLT_MAX : FLT_MAX; + + // Base offset into the flat (nRows * vocabSize) logits array for this row. + auto const rowOffset = rowId * vocabSize; + // Base offset into the tmp buffers for this (row, blockLane). + auto const tmpIdxBase = rowId * BLOCKS_PER_BEAM_ * maxTopK + blockLane * k; + + // Copy this block's chunk of logits into tmpLogProbs scratch space. + for (auto elemId = tid + blockLane * BLOCK_SIZE_; elemId < vocabSize; elemId += BLOCK_SIZE_ * BLOCKS_PER_BEAM_) + { + tmpLogProbs[rowOffset + elemId] = logits[rowOffset + elemId]; + } + __syncthreads(); + + // Iteratively find the top-k values via max-reduction, zeroing each found max. + TopK_2 partial; + for (int32_t ite = 0; ite < k; ite++) + { + partial.init(); + for (auto elemId = tid + blockLane * BLOCK_SIZE_; elemId < vocabSize; elemId += BLOCK_SIZE_ * BLOCKS_PER_BEAM_) + { + partial.insert(tmpLogProbs[rowOffset + elemId], rowOffset + elemId); + } + + TopK_2 total = BlockReduce(tempStorage).Reduce(partial, reduce_topk_op_2); + + if (tid == 0) + { + topKTmpIdBuf[tmpIdxBase + ite] = total.p; // global index (rowOffset + vocabIdx) + topKTmpValBuf[tmpIdxBase + ite] = total.u; // logit value + if (total.p >= 0) + { + tmpLogProbs[total.p] = -MAX_T_VAL; // zero out so next iteration finds next-best + } + } + __syncthreads(); + } +} + +// Stage 2: Merge BLOCKS_PER_BEAM_ * k candidates per row, apply optional top-p, +// then scatter selected logit values back to an output logits tensor (all other +// positions are set to -inf so that a subsequent softmax produces 0 probability). +template +__global__ void topKProbStage2ForLogits(int32_t const* __restrict__ topKTmpIdBuf, T* topKTmpValBuf, float* outputLogits, + int32_t maxTopK, int32_t const* topKs, float const* topPs, int32_t vocabSize) +{ + bool const IS_FP16 = std::is_same::value; + T const MAX_T_VAL = IS_FP16 ? HALF_FLT_MAX : FLT_MAX; + + auto const tid = static_cast(threadIdx.x); + auto const rowId = static_cast(blockIdx.x); + + auto const k = (topKs != nullptr) ? topKs[rowId] : maxTopK; + // size: number of valid candidates written by Stage 1 for this row. + // stride: row pitch in the tmp buffers (same as in invokeBatchTopKSampling). + auto const size = k * BLOCKS_PER_BEAM_; + auto const stride = maxTopK * BLOCKS_PER_BEAM_; + + typedef cub::BlockReduce, BLOCK_SIZE_> BlockReduce; + __shared__ typename BlockReduce::TempStorage tempStorage; + extern __shared__ char sharedArray[]; + // Shared layout: sId[maxTopK] | sVal2[maxTopK] + auto* sId = reinterpret_cast(sharedArray); + auto* sVal2 = reinterpret_cast(sId + maxTopK); + + // Pointer to this row's candidates in the tmp value buffer (modified in-place during reduction). + T* sVal = topKTmpValBuf + rowId * stride; + + // Step 1: Initialize output row to -inf (all threads cooperate for bandwidth). + float* outRow = outputLogits + rowId * vocabSize; + float const negInf = -std::numeric_limits::infinity(); + for (int32_t i = tid; i < vocabSize; i += BLOCK_SIZE_) + { + outRow[i] = negInf; + } + __syncthreads(); + + // Step 2: k-round block-reduction over the k * BLOCKS_PER_BEAM_ valid candidates. + // (Only the first 'size' entries of the row's tmp buffer were written by Stage 1.) + TopK_2 partial; + __shared__ float sMaxLogit; + for (int32_t ite = 0; ite < k; ite++) + { + partial.init(); + for (int32_t i = tid; i < size; i += BLOCK_SIZE_) + { + partial.insert(static_cast(sVal[i]), i); + } + + TopK_2 total = BlockReduce(tempStorage).Reduce(partial, reduce_topk_op_2); + + if (tid == 0) + { + if (ite == 0) + { + sMaxLogit = total.u; + } + sId[ite] = total.p; + sVal[total.p] = -MAX_T_VAL; // zero out so next iteration finds next-best + sVal2[ite] = total.u; // store raw logit value (not exponentiated) + } + __syncthreads(); + } + + // Step 3: Determine top-p cutoff (tid=0 only). + // sVal2 contains logit values in descending order; we exponentiate to get unnormalized probs. + if (tid == 0) + { + int32_t cutoff = k; + if (topPs != nullptr) + { + float const topP = topPs[rowId]; + if (topP < 1.0f) + { + // Compute unnormalized probabilities and their sum. + float sSum = 0.0f; + for (int32_t ki = 0; ki < k; ki++) + { + sVal2[ki] = __expf(sVal2[ki] - sMaxLogit); // reuse sVal2 to hold exp probs + sSum += sVal2[ki]; + } + // Walk in descending-probability order; stop as soon as cumulative prob >= topP. + float cumProb = 0.0f; + for (int32_t ki = 0; ki < k; ki++) + { + cumProb += sVal2[ki] / sSum; + if (cumProb >= topP) + { + cutoff = ki + 1; // always keep at least this token + break; + } + } + } + } + + // Step 4: Scatter selected logit values back to output. + // topKTmpIdBuf stores (rowOffset + vocabIdx); recover vocabIdx with % vocabSize. + auto const rowStride = rowId * stride; + for (int32_t ki = 0; ki < cutoff; ki++) + { + auto const candidateIdx = sId[ki]; + auto const globalIdx = topKTmpIdBuf[rowStride + candidateIdx]; + if (globalIdx >= 0) + { + auto const vocabIdx = globalIdx % vocabSize; + // sVal2 was overwritten with exp probs when topP < 1; we need the original logit. + // Re-read from the original tmp buffer — the stored value IS the logit (set in Stage 1). + // However sVal[candidateIdx] was zeroed during Stage 2 reduction; but + // topKTmpValBuf still holds the original value at that index (sVal points there). + // We stored the logit as sVal2[ite] = total.u BEFORE any exp, so if topP was not + // applied we can use sVal2[ki] directly. If topP was applied, sVal2[ki] now holds + // the exp prob — we cannot recover the logit. To handle both cases cleanly we use + // log(sVal2[ki]) + sMaxLogit when topPs was applied, otherwise sVal2[ki] directly. + float logitVal; + if (topPs != nullptr && topPs[rowId] < 1.0f) + { + // sVal2[ki] = exp(logit - sMaxLogit), so logit = log(sVal2[ki]) + sMaxLogit + logitVal = __logf(sVal2[ki]) + sMaxLogit; + } + else + { + logitVal = sVal2[ki]; // still the raw logit + } + outRow[vocabIdx] = logitVal; + } + } + } +} + +#define CASE_K_PROB(K_MAX, BLOCK_SIZE_1_, BLOCK_SIZE_2_, BLOCKS_PER_BEAM_) \ + do \ + { \ + topKProbStage1 \ + <<>>( \ + logits, tmpLogProbs, topKTmpIdBuf, topKTmpValBuf, maxTopK, topKs, vocabSize); \ + topKProbStage2ForLogits \ + <<>>( \ + topKTmpIdBuf, topKTmpValBuf, outputLogits, maxTopK, topKs, topPs, vocabSize); \ + } while (0) + +// Host launcher: allocates workspace tensors internally and dispatches the two-stage kernels. +// logits [nRows, vocabSize] – temperature-scaled input (float or half) +// outputLogits [nRows, vocabSize] – output: -inf everywhere except selected top-k-p positions +// topKs [nRows] – per-row k values (int32, on device) +// topPs [nRows] or nullptr – per-row p values (float, on device) +// maxTopK – maximum k across all rows (CPU scalar, 1–1024) +template +void invokeTopKTopPMaskingForProbs(T const* logits, float* outputLogits, int32_t const* topKs, float const* topPs, + int32_t maxTopK, int32_t nRows, int32_t vocabSize, cudaStream_t stream) +{ + constexpr int32_t BLOCKS_PER_BEAM = 8; + + // Workspace buffers (allocated as CUDA device tensors via ATen). + auto opts = at::TensorOptions().dtype(torch::kFloat32).device(at::kCUDA); + auto tmpLogProbsTensor = torch::empty({nRows * vocabSize}, opts); + auto topKTmpIdBufTensor + = torch::empty({nRows * BLOCKS_PER_BEAM * maxTopK}, at::TensorOptions().dtype(torch::kInt32).device(at::kCUDA)); + // topKTmpValBuf uses the same dtype as T; we allocate as float and reinterpret for half if needed. + auto topKTmpValBufTensor = torch::empty({nRows * BLOCKS_PER_BEAM * maxTopK}, opts); + + T* tmpLogProbs = reinterpret_cast(tmpLogProbsTensor.data_ptr()); + int32_t* topKTmpIdBuf = topKTmpIdBufTensor.data_ptr(); + T* topKTmpValBuf = reinterpret_cast(topKTmpValBufTensor.data_ptr()); + + int32_t logMaxTopK = 0; + int32_t recursor = maxTopK - 1; + while (recursor >>= 1) + { + ++logMaxTopK; + } + + switch (logMaxTopK) + { + case 0: + case 1: + case 2: + case 3: // 0 < maxTopK <= 16 + CASE_K_PROB(16, 128, 128, 8); + break; + case 4: // 16 < maxTopK <= 32 + CASE_K_PROB(32, 256, 128, 8); + break; + case 5: // 32 < maxTopK <= 64 + CASE_K_PROB(64, 256, 256, 8); + break; + case 6: + case 7: + case 8: + case 9: // 64 < maxTopK <= 1024 + CASE_K_PROB(1024, 256, 256, 8); + break; + default: TLLM_CHECK_WITH_INFO(false, "topKProbMasking supports 1 <= k <= 1024 but got k=%d", maxTopK); + } +} + +#undef CASE_K_PROB + +namespace +{ +constexpr double kGreedyTempThreshold = 1e-4; + +torch::Tensor computeSoftmaxForProbOp(torch::Tensor logits) +{ + TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA tensor"); + TORCH_CHECK(logits.dim() == 2, "logits must be a 2D tensor"); + + auto probs = logits.contiguous().to(torch::kFloat32); + auto stream = at::cuda::getCurrentCUDAStream(probs.device().index()); + + BiasSoftmaxParams biasSoftmaxParams; + biasSoftmaxParams.logits = probs.data_ptr(); + biasSoftmaxParams.probs = probs.data_ptr(); + biasSoftmaxParams.batchSize = static_cast(probs.size(0)); + biasSoftmaxParams.maxBatchSize = static_cast(probs.size(0)); + biasSoftmaxParams.maxBeamWidth = 1; + biasSoftmaxParams.vocabSize = static_cast(probs.size(1)); + biasSoftmaxParams.vocabSizePadded = static_cast(probs.size(1)); + biasSoftmaxParams.skipSoftMax = false; + biasSoftmaxParams.batchSlotsLogits = false; + biasSoftmaxParams.checkParams(); + + invokeAddBiasSoftMax(biasSoftmaxParams, stream); + return probs; +} + +// Fast path for top-K (and optional top-P) filtering using torch::topk instead of a +// full vocab-size sort. kMax must be provided as a CPU integer (the caller computes it +// via topK.max().item() on the Python side). When kMax == 0 or kMax >= vocabSize the +// function falls back to the original sort-based path. +// +// Key advantages over the full-sort path: +// 1. torch::topk with small kMax is O(V * log kMax) vs O(V * log V) for full sort. +// 2. The topk index tensor is [nRows, kMax] instead of [nRows, V] — much smaller. +// 3. No scatter-back of sorted indices needed; masking is done directly on logits. +// 4. For combined top-K + top-P, softmax/cumsum are computed on kMax values (not V). +torch::Tensor applyTopKTopPForProbOp(torch::Tensor logits, torch::optional const& topK, + torch::optional const& topP, int32_t kMax) +{ + bool const hasTopK = topK.has_value() && topK->defined(); + bool const hasTopP = topP.has_value() && topP->defined(); + + if (!hasTopK && !hasTopP) + { + return logits; + } + + int64_t const vocabSize = logits.size(1); + + if (hasTopK && kMax > 0 && kMax < vocabSize) + { + // Fast topk path ───────────────────────────────────────────────────────────── + // topKValues/topKIdx: [nRows, kMax], values in descending order + auto [topKValues, topKIdx] = logits.topk(kMax, /*dim=*/-1, /*largest=*/true, /*sorted=*/true); + + // validTopK[i, j]: True when position j falls within top-K[i] for row i + auto kArange = torch::arange(kMax, torch::TensorOptions().dtype(torch::kInt64).device(logits.device())) + .unsqueeze(0); // [1, kMax] + auto kVals = topK->to(torch::kInt64).unsqueeze(1); // [nRows, 1] + auto validTopK = kArange < kVals; // [nRows, kMax] + + // Start with everything masked; scatter will unmark the kept positions. + auto mask = torch::ones( + {logits.size(0), vocabSize}, torch::TensorOptions().dtype(torch::kBool).device(logits.device())); + + if (hasTopP) + { + // Compute top-P on the kMax descending-sorted values only (much cheaper). + // Positions beyond K[i] are treated as -inf so their probability ≈ 0. + auto validTopKValues = topKValues.masked_fill(~validTopK, -std::numeric_limits::infinity()); + auto sortedProbs = validTopKValues.softmax(/*dim=*/-1); // [nRows, kMax] + auto cumsum = sortedProbs.cumsum(/*dim=*/-1); // [nRows, kMax] + // Mask positions where the cumulative probability *before* this token + // already reaches topP — i.e. we have enough probability mass already. + auto topPMask = (cumsum - sortedProbs) >= topP->unsqueeze(1); // [nRows, kMax] + topPMask.select(/*dim=*/1, /*index=*/0).fill_(false); // always keep the top-1 token + // combinedMask: True → mask this vocab position + // False → keep this vocab position + auto combinedMask = topPMask | (~validTopK); // [nRows, kMax] + mask.scatter_(/*dim=*/1, /*index=*/topKIdx, /*src=*/combinedMask); + } + else + { + // Top-K only: unmark the first K[i] positions (those within validTopK). + // ~validTopK is True for positions j >= K[i] → they should stay masked. + mask.scatter_(/*dim=*/1, /*index=*/topKIdx, /*src=*/(~validTopK)); + } + + return logits.masked_fill(mask, -std::numeric_limits::infinity()); + } + + // Fallback: full-sort path (used for top-P only, or when kMax == 0) ──────────── + auto sortResult = logits.sort(/*dim=*/-1, /*descending=*/false); + auto logitsSort = std::get<0>(sortResult); + auto logitsIdx = std::get<1>(sortResult); + + if (hasTopK) + { + auto topKMask = logitsSort.size(1) - topK->to(torch::kLong); + topKMask = topKMask.clamp_min(0); + auto topKThreshold = logitsSort.gather(1, topKMask.unsqueeze(1)); + auto mask = logitsSort < topKThreshold; + logitsSort.masked_fill_(mask, -std::numeric_limits::infinity()); + } + + if (hasTopP) + { + auto probsSort = logitsSort.softmax(/*dim=*/-1); + auto probsSum = probsSort.cumsum(/*dim=*/-1, /*dtype=*/probsSort.scalar_type()); + auto topPMask = probsSum <= (1.0 - topP->unsqueeze(1)); + topPMask.select(/*dim=*/1, /*index=*/logitsSort.size(1) - 1).fill_(false); + logitsSort.masked_fill_(topPMask, -std::numeric_limits::infinity()); + } + + return logitsSort.scatter(/*dim=*/-1, /*index=*/logitsIdx, /*src=*/logitsSort); +} + +} // namespace + +torch::Tensor computeProbsFromLogits(torch::Tensor const& logits, torch::Tensor const& temperatures, + torch::optional const& topK, torch::optional const& topP, bool skipTemperature, + int32_t kMax) +{ + TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA tensor"); + TORCH_CHECK(temperatures.is_cuda(), "temperatures must be a CUDA tensor"); + TORCH_CHECK(logits.dim() == 2, "logits must be a 2D tensor"); + TORCH_CHECK(temperatures.dim() == 1, "temperatures must be a 1D tensor"); + TORCH_CHECK(logits.size(0) == temperatures.size(0), "logits and temperatures size mismatch"); + if (topK.has_value() && topK->defined()) + { + TORCH_CHECK(topK->is_cuda(), "top_k must be a CUDA tensor"); + TORCH_CHECK(topK->dim() == 1, "top_k must be a 1D tensor"); + TORCH_CHECK(topK->size(0) == logits.size(0), "top_k and logits size mismatch"); + } + if (topP.has_value() && topP->defined()) + { + TORCH_CHECK(topP->is_cuda(), "top_p must be a CUDA tensor"); + TORCH_CHECK(topP->dim() == 1, "top_p must be a 1D tensor"); + TORCH_CHECK(topP->size(0) == logits.size(0), "top_p and logits size mismatch"); + } + + auto scaledLogits + = (skipTemperature ? logits : logits.div(temperatures.unsqueeze(1))).contiguous().to(torch::kFloat32); + + bool const hasTopK = topK.has_value() && topK->defined(); + int64_t const vocabSize = scaledLogits.size(1); + int64_t const nRows = scaledLogits.size(0); + + torch::Tensor maskedLogits; + if (hasTopK && kMax > 0 && kMax < vocabSize) + { + // Two-stage CUDA top-k/top-p masking (mirrors invokeBatchTopKSampling). + maskedLogits = torch::empty_like(scaledLogits); + auto stream = at::cuda::getCurrentCUDAStream(scaledLogits.device().index()); + invokeTopKTopPMaskingForProbs(scaledLogits.data_ptr(), maskedLogits.data_ptr(), + topK->to(torch::kInt32).contiguous().data_ptr(), + (topP.has_value() && topP->defined()) ? topP->to(torch::kFloat32).contiguous().data_ptr() : nullptr, + kMax, static_cast(nRows), static_cast(vocabSize), stream); + } + else + { + // Fallback: PyTorch-based sort path (top-P only or kMax == 0). + maskedLogits = applyTopKTopPForProbOp(scaledLogits, topK, topP, kMax); + } + + auto probs = computeSoftmaxForProbOp(maskedLogits); + + auto isGreedy = temperatures <= kGreedyTempThreshold; + auto argmaxIds = maskedLogits.argmax(/*dim=*/-1, /*keepdim=*/true); + auto oneHot = torch::zeros_like(probs).scatter_(1, argmaxIds, 1.0); + return torch::where(isGreedy.unsqueeze(1), oneHot, probs); +} + +torch::Tensor computeDraftProbsForDynamicTreeRejection(torch::Tensor const& draftLogits, + torch::Tensor const& temperatures, SizeType32 const numDraftProbRows, torch::optional const& topK, + torch::optional const& topP, SizeType32 const targetVocabSize, bool skipTemperature, + torch::optional const& d2t, SizeType32 const kMax) +{ + TORCH_CHECK(draftLogits.is_cuda(), "draftLogits must be a CUDA tensor"); + TORCH_CHECK(temperatures.is_cuda(), "temperatures must be a CUDA tensor"); + TORCH_CHECK(draftLogits.dim() == 2, "draftLogits must be a 2D tensor"); + TORCH_CHECK(temperatures.dim() == 1, "temperatures must be a 1D tensor"); + TORCH_CHECK(numDraftProbRows > 0, "numDraftProbRows must be positive"); + + auto const batchSize = temperatures.size(0); + auto const draftVocabSize = draftLogits.size(1); + + TORCH_CHECK(batchSize > 0, "batchSize must be positive"); + TORCH_CHECK( + draftLogits.size(0) == batchSize * numDraftProbRows, "draftLogits row count does not match numDraftProbRows"); + TORCH_CHECK(targetVocabSize >= draftVocabSize, "targetVocabSize must be >= draft vocab size"); + + if (topK.has_value() && topK->defined()) + { + TORCH_CHECK(topK->is_cuda(), "top_k must be a CUDA tensor"); + TORCH_CHECK(topK->dim() == 1, "top_k must be a 1D tensor"); + TORCH_CHECK(topK->size(0) == batchSize, "top_k size mismatch"); + } + if (topP.has_value() && topP->defined()) + { + TORCH_CHECK(topP->is_cuda(), "top_p must be a CUDA tensor"); + TORCH_CHECK(topP->dim() == 1, "top_p must be a 1D tensor"); + TORCH_CHECK(topP->size(0) == batchSize, "top_p size mismatch"); + } + if (d2t.has_value() && d2t->defined()) + { + TORCH_CHECK(d2t->is_cuda(), "d2t must be a CUDA tensor"); + TORCH_CHECK(d2t->dim() == 1, "d2t must be a 1D tensor"); + TORCH_CHECK(d2t->size(0) >= draftVocabSize, "d2t size mismatch"); + } + + auto draftTemps = temperatures.repeat_interleave(numDraftProbRows); + auto draftTopK = topK.has_value() && topK->defined() + ? torch::optional(topK->repeat_interleave(numDraftProbRows)) + : torch::optional(); + auto draftTopP = topP.has_value() && topP->defined() + ? torch::optional(topP->repeat_interleave(numDraftProbRows)) + : torch::optional(); + + auto draftProbs = computeProbsFromLogits(draftLogits, draftTemps, draftTopK, draftTopP, skipTemperature, kMax) + .reshape({batchSize, numDraftProbRows, draftVocabSize}); + + if (draftVocabSize == targetVocabSize) + { + return draftProbs; + } + + auto fullDraftProbs = torch::zeros({batchSize, numDraftProbRows, targetVocabSize}, + torch::TensorOptions().dtype(torch::kFloat32).device(draftProbs.device())); + if (d2t.has_value() && d2t->defined()) + { + auto srcIdx + = torch::arange(draftVocabSize, torch::TensorOptions().dtype(torch::kInt64).device(draftProbs.device())); + auto targetIdx = srcIdx + d2t->slice(0, 0, draftVocabSize).to(torch::kInt64); + auto expandedTargetIdx + = targetIdx.view({1, 1, draftVocabSize}).expand({batchSize, numDraftProbRows, draftVocabSize}); + fullDraftProbs.scatter_(2, expandedTargetIdx, draftProbs); + } + else + { + fullDraftProbs.slice(/*dim=*/2, /*start=*/0, /*end=*/draftVocabSize).copy_(draftProbs); + } + + return fullDraftProbs; +} + +std::tuple computeTargetProbsForDynamicTreeRejection( + torch::Tensor const& targetLogits, torch::Tensor const& temperatures, SizeType32 const numDraftTokens, + torch::optional const& topK, torch::optional const& topP, bool skipTemperature, + SizeType32 const kMax) +{ + TORCH_CHECK(targetLogits.is_cuda(), "targetLogits must be a CUDA tensor"); + TORCH_CHECK(temperatures.is_cuda(), "temperatures must be a CUDA tensor"); + TORCH_CHECK(targetLogits.dim() == 2, "targetLogits must be a 2D tensor"); + TORCH_CHECK(temperatures.dim() == 1, "temperatures must be a 1D tensor"); + TORCH_CHECK(numDraftTokens > 1, "numDraftTokens must be greater than 1"); + + auto const batchSize = temperatures.size(0); + auto const targetVocabSize = targetLogits.size(1); + auto const nRows = batchSize * numDraftTokens; + + TORCH_CHECK(batchSize > 0, "batchSize must be positive"); + TORCH_CHECK( + targetLogits.size(0) == batchSize * numDraftTokens, "targetLogits row count does not match numDraftTokens"); + + if (topK.has_value() && topK->defined()) + { + TORCH_CHECK(topK->is_cuda(), "top_k must be a CUDA tensor"); + TORCH_CHECK(topK->dim() == 1, "top_k must be a 1D tensor"); + TORCH_CHECK(topK->size(0) == batchSize, "top_k size mismatch"); + } + if (topP.has_value() && topP->defined()) + { + TORCH_CHECK(topP->is_cuda(), "top_p must be a CUDA tensor"); + TORCH_CHECK(topP->dim() == 1, "top_p must be a 1D tensor"); + TORCH_CHECK(topP->size(0) == batchSize, "top_p size mismatch"); + } + + auto targetTemps = temperatures.repeat_interleave(numDraftTokens); + auto targetTopK = topK.has_value() && topK->defined() + ? torch::optional(topK->repeat_interleave(numDraftTokens)) + : torch::optional(); + auto targetTopP = topP.has_value() && topP->defined() + ? torch::optional(topP->repeat_interleave(numDraftTokens)) + : torch::optional(); + + bool const hasTopK = targetTopK.has_value() && targetTopK->defined(); + bool const hasTopP = targetTopP.has_value() && targetTopP->defined(); + bool const hasFiltering = hasTopK || hasTopP; + + auto scaledTargetLogits = (skipTemperature ? targetLogits : targetLogits.div(targetTemps.unsqueeze(1))) + .contiguous() + .to(torch::kFloat32); + + torch::Tensor maskedTargetLogits; + torch::Tensor targetSupportIndices; + torch::Tensor targetSupportLengths; + + if (!hasFiltering) + { + // No filtering: use full-vocab probs; sparse support is not applicable. + maskedTargetLogits = scaledTargetLogits; + targetSupportIndices + = torch::empty({0}, torch::TensorOptions().dtype(torch::kInt32).device(targetLogits.device())); + targetSupportLengths + = torch::empty({0}, torch::TensorOptions().dtype(torch::kInt32).device(targetLogits.device())); + } + else if (hasTopK && kMax > 0 && static_cast(kMax) < targetVocabSize) + { + // Fast two-stage CUDA path for masked logits. + maskedTargetLogits = torch::empty_like(scaledTargetLogits); + auto stream = at::cuda::getCurrentCUDAStream(scaledTargetLogits.device().index()); + invokeTopKTopPMaskingForProbs(scaledTargetLogits.data_ptr(), maskedTargetLogits.data_ptr(), + targetTopK->to(torch::kInt32).contiguous().data_ptr(), + hasTopP ? targetTopP->to(torch::kFloat32).contiguous().data_ptr() : nullptr, kMax, + static_cast(nRows), static_cast(targetVocabSize), stream); + + // Extract support indices: the finite positions after masking (at most kMax per row). + auto [topKVals, topKIdx] = maskedTargetLogits.topk(kMax, /*dim=*/-1, /*largest=*/true, /*sorted=*/true); + auto validMask = topKVals.isfinite(); // [nRows, kMax] + auto supportLengthsLong = validMask.sum(/*dim=*/-1, /*keepdim=*/false, torch::kLong); // [nRows] + auto supportIndicesRaw + = torch::where(validMask, topKIdx.to(torch::kInt32), torch::full_like(topKIdx.to(torch::kInt32), -1)); + + targetSupportIndices = supportIndicesRaw.reshape({batchSize, numDraftTokens, kMax}); + targetSupportLengths = supportLengthsLong.to(torch::kInt32).reshape({batchSize, numDraftTokens}); + } + else + { + // Sort-based fallback: top-P only, or kMax == 0 / kMax >= vocabSize. + auto sortResult = scaledTargetLogits.sort(/*dim=*/-1, /*descending=*/false); + auto logitsSort = std::get<0>(sortResult); + auto logitsIdx = std::get<1>(sortResult); + + if (hasTopK) + { + auto topKMask = logitsSort.size(1) - targetTopK->to(torch::kLong); + topKMask = topKMask.clamp_min(0); + auto topKThreshold = logitsSort.gather(1, topKMask.unsqueeze(1)); + auto mask = logitsSort < topKThreshold; + logitsSort.masked_fill_(mask, -std::numeric_limits::infinity()); + } + + if (hasTopP) + { + auto probsSort = logitsSort.softmax(/*dim=*/-1); + auto probsSum = probsSort.cumsum(/*dim=*/-1, /*dtype=*/probsSort.scalar_type()); + auto topPMask = probsSum <= (1.0 - targetTopP->unsqueeze(1)); + topPMask.select(/*dim=*/1, /*index=*/logitsSort.size(1) - 1).fill_(false); + logitsSort.masked_fill_(topPMask, -std::numeric_limits::infinity()); + } + + maskedTargetLogits = logitsSort.scatter(/*dim=*/-1, /*index=*/logitsIdx, /*src=*/logitsSort); + + // Compact support indices: finite values are at the END of the ascending-sorted logitsSort. + auto supportLengthsLong + = logitsSort.isfinite().sum(/*dim=*/-1, /*keepdim=*/false, /*dtype=*/torch::kLong); // [nRows] + auto supportLengths1D = supportLengthsLong.to(torch::kInt32); + + int64_t maxSupportSize = targetVocabSize; + if (hasTopK && targetTopK->defined()) + { + maxSupportSize = std::min(targetVocabSize, targetTopK->max().item()); + } + + auto compactPositions + = torch::arange(maxSupportSize, torch::TensorOptions().dtype(torch::kLong).device(targetLogits.device())) + .unsqueeze(0) + .expand({nRows, maxSupportSize}); + auto supportStart = targetVocabSize - supportLengthsLong.unsqueeze(1); + auto gatherPositions = (supportStart + compactPositions).clamp_max(targetVocabSize - 1); + auto gatheredSupportIndices = logitsIdx.gather(1, gatherPositions); + auto validMask = compactPositions < supportLengthsLong.unsqueeze(1); + auto invalidFill = torch::full_like(gatheredSupportIndices, -1L); + targetSupportIndices = torch::where(validMask, gatheredSupportIndices, invalidFill) + .to(torch::kInt32) + .reshape({batchSize, numDraftTokens, maxSupportSize}); + targetSupportLengths = supportLengths1D.reshape({batchSize, numDraftTokens}); + } + + auto targetProbs = computeSoftmaxForProbOp(maskedTargetLogits); + + auto isGreedy = targetTemps <= kGreedyTempThreshold; + auto argmaxIds = maskedTargetLogits.argmax(/*dim=*/-1, /*keepdim=*/true); + auto oneHot = torch::zeros_like(targetProbs).scatter_(1, argmaxIds, 1.0); + targetProbs = torch::where(isGreedy.unsqueeze(1), oneHot, targetProbs); + + return std::make_tuple( + targetProbs.reshape({batchSize, numDraftTokens, targetVocabSize}), targetSupportIndices, targetSupportLengths); +} + //! \param parentList [in] layer-wise parent indices [bs, topK*(depth-1)+1] //! \param selectedIndex [in] resampled history buffer indices [bs, draftTokenNum-1] //! \param treeMask [out] attention mask (which nodes each node can see) @@ -264,6 +932,49 @@ void invokeBuildDynamicTree(int64_t const* parentList, int64_t const* selectedIn sync_check_cuda_error(stream); } +__global__ void buildDraftProbIndicesKernel( + int64_t const* topkScoreIndices, int32_t* draftProbIndices, SizeType32 topK, SizeType32 numDraftTokens) +{ + int32_t const batchIdx = blockIdx.x; + int32_t const tokenIdx = threadIdx.x; + + if (tokenIdx > numDraftTokens) + { + return; + } + + int32_t* draftProbIndicesRow = draftProbIndices + batchIdx * (numDraftTokens + 1); + + if (tokenIdx == 0) + { + draftProbIndicesRow[0] = 0; + return; + } + + int64_t const histIdx = topkScoreIndices[batchIdx * numDraftTokens + (tokenIdx - 1)]; + int32_t draftProbRow = 0; + + if (histIdx >= topK) + { + int64_t const relative = histIdx - topK; + int64_t const depthBucket = relative / (topK * topK); + int64_t const parentK = (relative % (topK * topK)) / topK; + draftProbRow = static_cast(1 + depthBucket * topK + parentK); + } + + draftProbIndicesRow[tokenIdx] = draftProbRow; +} + +void invokeBuildDraftProbIndices(int64_t const* topkScoreIndices, int32_t* draftProbIndices, SizeType32 batchSize, + SizeType32 topK, SizeType32 numDraftTokens, cudaStream_t stream) +{ + dim3 const grid(batchSize); + dim3 const block(numDraftTokens + 1); + + buildDraftProbIndicesKernel<<>>(topkScoreIndices, draftProbIndices, topK, numDraftTokens); + sync_check_cuda_error(stream); +} + //! \param predicts [out] accepted token ids + bonus token [bs * numDraftTokens] //! \param acceptIndex [out] accepted path as local tree positions [bs, numSpeculativeTokens] //! \param acceptTokenNum [out] number of accepted draft tokens per batch [bs] @@ -273,7 +984,8 @@ void invokeBuildDynamicTree(int64_t const* parentList, int64_t const* selectedIn //! \param retrieveNextSibling [in] next-sibling pointer [bs, numDraftTokens], -1=none //! \param targetPredict [in] target model prediction per position [bs * numDraftTokens] //! \param batchSize batch size -//! \param numSpeculativeTokens second dim of acceptIndex (>= max possible accepts + 1) +//! \param numSpeculativeTokens second dim of acceptIndex/acceptToken +//! (= numSpecStep = max_path_len, >= max possible accepts + 1) //! \param numDraftTokens total tree nodes per batch (including root) __global__ void verifyDynamicTreeGreedyKernel(int64_t* predicts, int64_t* acceptIndex, int64_t* acceptTokenNum, int64_t* acceptToken, int64_t const* candidates, int32_t const* retrieveIndex, int32_t const* retrieveNextToken, @@ -351,6 +1063,406 @@ void invokeVerifyDynamicTreeGreedy(int64_t* predicts, int64_t* acceptIndex, int6 sync_check_cuda_error(stream); } +// ------------------------------------------------------------ +// Background: Speculative Sampling Theory +// ------------------------------------------------------------ +// +// Goal: reuse draft model samples to speed up generation while keeping the +// final output distribution strictly equal to the target distribution q. +// +// For a given token x: +// p(x) = draft_probs[x] (draft model probability) +// q(x) = target_probs[x] (target model probability) +// +// Step 1 - The draft model proposes token x sampled from p. +// Step 2 - Accept x with probability min(1, q(x)/p(x)). +// Equivalently: accept when u * p(x) < q(x), where u ~ Uniform(0,1). +// +// Why does this work? +// x is proposed with probability p(x) and then accepted with probability +// min(1, q(x)/p(x)), so its total probability mass reaching the output is: +// p(x) * min(1, q(x)/p(x)) = min(p(x), q(x)) +// +// This covers only the min(p, q) portion of the target mass. +// The remaining portion q - min(p, q) = relu(q - p) is not yet covered. +// +// Therefore, if the draft token is rejected, we must resample from the +// residual distribution relu(q - p) (normalised) to fill the gap and +// restore the full target distribution. +// +// Example: +// p = [0.6, 0.3, 0.1] tokens [A, B, C] +// q = [0.2, 0.5, 0.3] +// +// Accept probabilities: +// A: min(1, 0.2/0.6) = 1/3 B: min(1, 0.5/0.3) = 1 C: min(1, 0.3/0.1) = 1 +// +// Case 1 - draft proposes A (prob 0.6): +// Accept (1/3): contributes 0.6 * 1/3 = 0.2 to output A. +// Reject (2/3): total rejected mass = 0.6 * 2/3 = 0.4. +// relu(q-p) = [0, 0.2, 0.2] -> normalised [0, 0.5, 0.5] +// contributes 0.4*0.5 = 0.2 to B and 0.4*0.5 = 0.2 to C. +// Case 2 - draft proposes B (prob 0.3): always accepted -> 0.3 to B. +// Case 3 - draft proposes C (prob 0.1): always accepted -> 0.1 to C. +// +// Final output distribution: +// A = 0.2, B = 0.3 + 0.2 = 0.5, C = 0.1 + 0.2 = 0.3 -> exactly q. +// +// Tree extension: +// The same logic applies depth-by-depth along the draft tree. At each +// depth the kernel tries siblings in score order; the first accepted +// sibling extends the current path. If every sibling at a depth is +// rejected the kernel samples a correction token from relu(q-p) and +// terminates traversal for that request. +// ------------------------------------------------------------ + +#include + +__device__ int64_t sampleFromDistribution(curandStatePhilox4_32_10_t& state, float const* probs, uint32_t vocabSize) +{ + float r = curand_uniform(&state); + float cumsum = 0.0f; + int64_t sampledTok = static_cast(vocabSize) - 1; // fallback: last vocab token + + for (uint32_t v = 0; v < vocabSize; ++v) + { + cumsum += probs[v]; + if (r <= cumsum) + { + sampledTok = static_cast(v); + break; + } + } + + return sampledTok; +} + +__device__ int64_t sampleFromIndexedDistribution(curandStatePhilox4_32_10_t& state, float const* probs, + int32_t const* supportIndices, uint32_t supportSize, uint32_t vocabSize) +{ + float r = curand_uniform(&state); + float cumsum = 0.0f; + int64_t sampledTok = static_cast(vocabSize) - 1; // fallback: last vocab token + + for (uint32_t i = 0; i < supportSize; ++i) + { + int32_t const tok = supportIndices[i]; + cumsum += probs[tok]; + if (r <= cumsum) + { + sampledTok = static_cast(tok); + break; + } + } + + return sampledTok; +} + +//! \param acceptIndex [out] accepted path as tree positions [bs, numSpecStep]. int64. +//! \param acceptTokenNum [out] number of accepted draft tokens (excl. root) [bs]. int64. +//! \param acceptToken [out] emitted token ids [bs, numSpecStep]. int64. +//! \param candidates [in] candidate token ids [bs, numDraftTokens]; col 0 = root. int64. +//! \param draftProbs [in] unique draft probs [bs, numDraftProbRows, vocabSize]. float32. +//! \param targetProbs [in] target probs [bs, numDraftTokens, vocabSize]; index 0 = root. float32. +//! \param targetSupportIndices [in] compact target support per tree position +//! [bs, numDraftTokens, maxTargetSupportSize]. int32, or nullptr. +//! \param targetSupportLengths [in] support length per tree position [bs, numDraftTokens]. int32, or nullptr. +//! \param draftProbIndices [in] tree position -> draftProbs row [bs, numDraftTokens], root unused. int32. +//! \param retrieveNextToken [in] first-child pointer [bs, numDraftTokens], -1=none. int32. +//! \param retrieveNextSibling [in] next-sibling pointer [bs, numDraftTokens], -1=none. int32. +//! \param batchSize batch size. +//! \param numDraftProbRows unique draft-prob rows per request. +//! \param maxTargetSupportSize support-array width. Zero when targetSupportIndices is null. +//! \param numSpecStep second dim of acceptIndex/acceptToken +//! (= max_path_len = max_draft_len + 1). +//! \param numDraftTokens total tree nodes per batch (including root). +//! \param vocabSize vocabulary size. +//! \param seed [1] int64 on GPU. Philox RNG seed. +//! \param offset [1] int64 on GPU. Philox RNG offset. +__global__ void verifyDynamicTreeRejectionKernel(int64_t* acceptIndex, int64_t* acceptTokenNum, int64_t* acceptToken, + int64_t const* candidates, float const* draftProbs, float const* targetProbs, int32_t const* targetSupportIndices, + int32_t const* targetSupportLengths, int32_t const* draftProbIndices, int32_t const* retrieveNextToken, + int32_t const* retrieveNextSibling, uint32_t batchSize, uint32_t numDraftProbRows, uint32_t maxTargetSupportSize, + uint32_t numSpeculativeTokens, uint32_t numDraftTokens, uint32_t vocabSize, int64_t const* seed, + int64_t const* offset) +{ + uint32_t bx = blockIdx.x; + if (bx >= batchSize) + { + return; + } + + uint32_t batchOffset = bx * numDraftTokens; + + curandStatePhilox4_32_10_t state; + curand_init(static_cast(seed[0]), static_cast(bx), static_cast(offset[0]), &state); + + // Root (depth 0): initialize path state at tree position 0. + // + // Example tree used in code review discussions: + // root: E + // children of E: F1, F2, F3 + // children of F1: G1, G2 + // children of F2: G3 + // + // In that example the per-request inputs are conceptually: + // candidates = [E, F1, F2, F3, G1, G2, G3] + // draftProbs = [p(.|E), p(.|F1), p(.|F2)] + // draftProbIndices = [0, 0, 0, 0, 1, 1, 2] + // targetProbs = [q(.|E), q(.|F1), q(.|F2), q(.|F3), q(.|G1), q(.|G2), q(.|G3)] + // + // draftProbs stores one row per unique parent context, so siblings that share + // the same parent also share the same draftProbs row via draftProbIndices. + // targetProbs remains aligned to all tree positions, including the root at slot 0. + // + // Output convention: + // - acceptIndex stores the accepted draft path as tree positions, with slot 0 + // reserved for the root position. + // - acceptToken stores the emitted token sequence, matching the greedy kernel: + // slot 0 = first emitted token + // slot numAcceptedTokens = final bonus/correction token + // - acceptTokenNum stores the number of accepted draft tokens only. The caller + // adds 1 to obtain the total number of emitted tokens. + int32_t lastAcceptedLocalIdx = 0; + acceptIndex[bx * numSpeculativeTokens] = lastAcceptedLocalIdx; + uint32_t numAcceptedTokens = 0; + int32_t curIndex = 0; + bool hasTerminalToken = false; + bool const hasCompactTargetSupport = targetSupportIndices != nullptr && targetSupportLengths != nullptr; + + for (uint32_t j = 1; j < numSpeculativeTokens; ++j) + { + // Advance to the first child of the last accepted node. + // Continuing the example above: + // j = 1, curIndex = 0 (E) -> firstChild = F1 + // j = 2, curIndex = 1 (F1) -> firstChild = G1 + int32_t firstChild = retrieveNextToken[batchOffset + curIndex]; + curIndex = firstChild; + + while (curIndex != -1) + { + int32_t draftLocalIdx = curIndex; // retrieveIndex is identity: draftLocalIdx == curIndex + int64_t draftTokenId = candidates[batchOffset + curIndex]; + // draftProbs: tree position curIndex maps to a unique parent-context row. + // For the example: + // curIndex = 1 (F1) -> draftProbIndices[1] = 0 -> p(F1|E) + // curIndex = 2 (F2) -> draftProbIndices[2] = 0 -> p(F2|E) + // curIndex = 4 (G1) -> draftProbIndices[4] = 1 -> p(G1|F1) + int32_t const draftProbRow = draftProbIndices[batchOffset + curIndex]; + float pDraft + = draftProbs[(static_cast(bx) * numDraftProbRows + draftProbRow) * vocabSize + draftTokenId]; + // Rejection sampling compares draft siblings under the target + // distribution of the currently accepted parent node. + float pTarget = targetProbs[(static_cast(bx) * numDraftTokens + lastAcceptedLocalIdx) * vocabSize + + draftTokenId]; + + // Rejection test for the current sibling: + // accept with probability min(1, q(x) / p(x)) + // where: + // p(x) = pDraft from the draft model proposal + // q(x) = pTarget from the target model verification distribution + float acceptProb = fminf(1.0f, pTarget / (pDraft + 1e-10f)); + + float u = curand_uniform(&state); + + if (u < acceptProb) + { + // Accepted sibling: extend the path and continue with this node's children. + // Example: + // if F1 is accepted at j = 1, then on the next iteration we descend to + // F1's first child (G1) and compare G1/G2 in sibling order. + // Emit the accepted draft token immediately. If we later stop at a leaf + // (or hit max depth), the final bonus token from the target distribution + // will be written at slot numAcceptedTokens. + acceptToken[bx * numSpeculativeTokens + numAcceptedTokens] = draftTokenId; + ++numAcceptedTokens; + acceptIndex[bx * numSpeculativeTokens + numAcceptedTokens] = draftLocalIdx; + lastAcceptedLocalIdx = draftLocalIdx; + break; + } + else + { + // Reject this sibling and try the next sibling at the same depth. + // Example: + // reject F1 -> move to F2 + // reject F2 -> move to F3 + curIndex = retrieveNextSibling[batchOffset + curIndex]; + } + } + + if (curIndex == -1) + { + // All siblings exhausted. Two sub-cases: + // (a) firstChild == -1: leaf node, no draft tokens at this depth. + // Emit the final bonus token sampled from q(.|lastAcceptedLocalIdx). + // (b) firstChild != -1: every sibling was rejected -> sample correction token + // from relu(q - p) at firstChild's position to restore the target distribution. + if (firstChild == -1) + { + float const* tProbs + = targetProbs + (static_cast(bx) * numDraftTokens + lastAcceptedLocalIdx) * vocabSize; + if (hasCompactTargetSupport) + { + uint32_t const supportOffset + = (static_cast(bx) * numDraftTokens + lastAcceptedLocalIdx) * maxTargetSupportSize; + uint32_t const supportSize + = static_cast(targetSupportLengths[batchOffset + lastAcceptedLocalIdx]); + acceptToken[bx * numSpeculativeTokens + numAcceptedTokens] = sampleFromIndexedDistribution( + state, tProbs, targetSupportIndices + supportOffset, supportSize, vocabSize); + } + else + { + acceptToken[bx * numSpeculativeTokens + numAcceptedTokens] + = sampleFromDistribution(state, tProbs, vocabSize); + } + hasTerminalToken = true; + } + else + { + int32_t const draftProbRow = draftProbIndices[batchOffset + firstChild]; + float const* tProbs + = targetProbs + (static_cast(bx) * numDraftTokens + lastAcceptedLocalIdx) * vocabSize; + float const* dProbs + = draftProbs + (static_cast(bx) * numDraftProbRows + draftProbRow) * vocabSize; + int32_t const* tProbIndices = nullptr; + uint32_t targetSupportSize = vocabSize; + if (hasCompactTargetSupport) + { + uint32_t const supportOffset + = (static_cast(bx) * numDraftTokens + lastAcceptedLocalIdx) * maxTargetSupportSize; + tProbIndices = targetSupportIndices + supportOffset; + targetSupportSize = static_cast(targetSupportLengths[batchOffset + lastAcceptedLocalIdx]); + } + + // Pass 1: compute diffSum = sum of relu(q - p) over the target support. + float diffSum = 0.0f; + if (hasCompactTargetSupport) + { + for (uint32_t i = 0; i < targetSupportSize; ++i) + { + uint32_t const v = static_cast(tProbIndices[i]); + float diff = tProbs[v] - dProbs[v]; + if (diff > 0.0f) + { + diffSum += diff; + } + } + } + else + { + for (uint32_t v = 0; v < targetSupportSize; ++v) + { + float diff = tProbs[v] - dProbs[v]; + if (diff > 0.0f) + { + diffSum += diff; + } + } + } + + // Pass 2: CDF inversion over the normalised residual distribution, + // traversing only the target support. + // Falls back to sampling directly from target when diffSum ~= 0. + // + // Example: + // if F1/F2/F3 are all rejected under parent E, then we sample a correction + // token from relu(q(.|E) - p(.|E)). The sampled token terminates traversal: + // it is emitted as the next output token, and there is no further descent. + int64_t corrTok = static_cast(vocabSize) - 1; // fallback: last vocab token + bool useDiff = (diffSum > 1e-10f); + + if (useDiff) + { + float r = curand_uniform(&state); + float cumsum = 0.0f; + if (hasCompactTargetSupport) + { + for (uint32_t i = 0; i < targetSupportSize; ++i) + { + uint32_t const v = static_cast(tProbIndices[i]); + float diff = tProbs[v] - dProbs[v]; + float prob = (diff > 0.0f) ? diff / diffSum : 0.0f; + cumsum += prob; + if (r <= cumsum) + { + corrTok = static_cast(v); + break; + } + } + } + else + { + for (uint32_t v = 0; v < targetSupportSize; ++v) + { + float diff = tProbs[v] - dProbs[v]; + float prob = (diff > 0.0f) ? diff / diffSum : 0.0f; + cumsum += prob; + if (r <= cumsum) + { + corrTok = static_cast(v); + break; + } + } + } + } + else + { + corrTok = hasCompactTargetSupport + ? sampleFromIndexedDistribution(state, tProbs, tProbIndices, targetSupportSize, vocabSize) + : sampleFromDistribution(state, tProbs, vocabSize); + } + + acceptToken[bx * numSpeculativeTokens + numAcceptedTokens] = corrTok; + // acceptIndex at correction depth left as 0 (zero-initialized by caller) + hasTerminalToken = true; + } + break; + } + } + + if (!hasTerminalToken) + { + // Reached max speculative depth while continuing to accept the draft path. + // Emit the final bonus token from the last accepted position. + float const* tProbs + = targetProbs + (static_cast(bx) * numDraftTokens + lastAcceptedLocalIdx) * vocabSize; + if (hasCompactTargetSupport) + { + uint32_t const supportOffset + = (static_cast(bx) * numDraftTokens + lastAcceptedLocalIdx) * maxTargetSupportSize; + uint32_t const supportSize + = static_cast(targetSupportLengths[batchOffset + lastAcceptedLocalIdx]); + acceptToken[bx * numSpeculativeTokens + numAcceptedTokens] = sampleFromIndexedDistribution( + state, tProbs, targetSupportIndices + supportOffset, supportSize, vocabSize); + } + else + { + acceptToken[bx * numSpeculativeTokens + numAcceptedTokens] + = sampleFromDistribution(state, tProbs, vocabSize); + } + } + + acceptTokenNum[bx] = numAcceptedTokens; +} + +void invokeVerifyDynamicTreeRejection(int64_t* acceptIndex, int64_t* acceptTokenNum, int64_t* acceptToken, + int64_t const* candidates, float const* draftProbs, float const* targetProbs, int32_t const* targetSupportIndices, + int32_t const* targetSupportLengths, int32_t const* draftProbIndices, int32_t const* retrieveNextToken, + int32_t const* retrieveNextSibling, SizeType32 batchSize, SizeType32 numDraftProbRows, + SizeType32 maxTargetSupportSize, SizeType32 numDraftTokens, SizeType32 numSpecStep, SizeType32 vocabSize, + int64_t const* seed, int64_t const* offset, cudaStream_t stream) +{ + dim3 grid(batchSize); + dim3 block(1); + + verifyDynamicTreeRejectionKernel<<>>(acceptIndex, acceptTokenNum, acceptToken, candidates, + draftProbs, targetProbs, targetSupportIndices, targetSupportLengths, draftProbIndices, retrieveNextToken, + retrieveNextSibling, batchSize, numDraftProbRows, maxTargetSupportSize, numSpecStep, numDraftTokens, vocabSize, + seed, offset); + + sync_check_cuda_error(stream); +} + } // namespace kernels::speculative_decoding TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h index c48b0f25005a..d2fe371d3823 100644 --- a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h +++ b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h @@ -19,6 +19,7 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/runtime/common.h" #include +#include TRTLLM_NAMESPACE_BEGIN @@ -60,6 +61,18 @@ void invokeBuildDynamicTree(int64_t const* parentList, int64_t const* selectedIn runtime::SizeType32 topK, runtime::SizeType32 depth, runtime::SizeType32 numDraftTokens, TreeMaskMode treeMaskMode, cudaStream_t stream, runtime::SizeType32 numInt32PerRow); +//! \brief Build tree-position -> unique draft-prob row mapping for rejection sampling. +//! \param topkScoreIndices [batchSize, numDraftTokens], on GPU. int64. +//! History-buffer indices selected for each final tree position. +//! \param draftProbIndices output [batchSize, numDraftTokens + 1], on GPU. int32. +//! Column 0 is reserved for the root and set to 0. +//! \param batchSize runtime::SizeType32. Batch size. +//! \param topK runtime::SizeType32. Tree top-K branching factor. +//! \param numDraftTokens runtime::SizeType32. Total number of non-root draft positions. +//! \param stream cuda stream. +void invokeBuildDraftProbIndices(int64_t const* topkScoreIndices, int32_t* draftProbIndices, + runtime::SizeType32 batchSize, runtime::SizeType32 topK, runtime::SizeType32 numDraftTokens, cudaStream_t stream); + //! \brief Verify dynamic tree using greedy strategy //! Verifies draft tokens against target model predictions using tree traversal. //! All index/token pointer parameters use int64. @@ -93,6 +106,70 @@ void invokeVerifyDynamicTreeGreedy(int64_t* predicts, int64_t* acceptIndex, int6 runtime::SizeType32 batchSize, runtime::SizeType32 numDraftTokens, runtime::SizeType32 numSpecStep, cudaStream_t stream); +//! \brief Verify dynamic tree using rejection sampling. +//! For each request, traverses the tree depth-by-depth. At each depth, siblings are tried +//! in order; the first sibling accepted by rejection sampling (p_target/p_draft) continues +//! the path. If all siblings are rejected, a correction token is sampled from (target-draft)_+. +//! \param acceptIndex output [batchSize, numSpecStep] int64 — tree positions of accepted tokens. +//! \param acceptTokenNum output [batchSize] int64 — # accepted draft tokens (excl. root). +//! \param acceptToken output [batchSize, numSpecStep] int64 — accepted/correction token ids. +//! \param candidates [batchSize, numDraftTokens] int64; col 0 = root (target sample). +//! \param draftProbs [batchSize, numDraftTokens-1, vocabSize] float32; index 0 = tree pos 1. +//! \param targetProbs [batchSize, numDraftTokens, vocabSize] float32; index 0 = root. +//! \param targetSupportIndices [batchSize, numDraftTokens, maxTargetSupportSize] int32; compact token ids that +//! survive top-k/top-p filtering for each row, padded with -1. May be empty when no filtering is active. +//! \param targetSupportLengths [batchSize, numDraftTokens] int32; valid support length per row. May be empty when +//! no filtering is active. +//! \param retrieveNextToken [batchSize, numDraftTokens] int32 first-child pointer, -1=none. +//! \param retrieveNextSibling [batchSize, numDraftTokens] int32 next-sibling pointer, -1=none. +//! \param batchSize runtime::SizeType32. +//! \param maxTargetSupportSize runtime::SizeType32. Third dim of targetSupportIndices. Can be zero. +//! \param numDraftTokens runtime::SizeType32. Total tree nodes per request (including root). +//! \param numSpecStep runtime::SizeType32. Second dim of acceptIndex/acceptToken. +//! \param vocabSize runtime::SizeType32. Vocabulary size. +//! \param seed [1] int64 on GPU. Philox RNG seed. +//! \param offset [1] int64 on GPU. Philox RNG offset. +//! \param stream cudaStream_t. +void invokeVerifyDynamicTreeRejection(int64_t* acceptIndex, int64_t* acceptTokenNum, int64_t* acceptToken, + int64_t const* candidates, float const* draftProbs, float const* targetProbs, int32_t const* targetSupportIndices, + int32_t const* targetSupportLengths, int32_t const* draftProbIndices, int32_t const* retrieveNextToken, + int32_t const* retrieveNextSibling, runtime::SizeType32 batchSize, runtime::SizeType32 numDraftProbRows, + runtime::SizeType32 maxTargetSupportSize, runtime::SizeType32 numDraftTokens, runtime::SizeType32 numSpecStep, + runtime::SizeType32 vocabSize, int64_t const* seed, int64_t const* offset, cudaStream_t stream); + +//! \brief Compute draft probabilities for dynamic-tree rejection sampling from logits. +//! \param draftLogits [batchSize * numDraftProbRows, draftVocabSize], on GPU. +//! \param temperatures [batchSize], on GPU. +//! \param numDraftProbRows runtime::SizeType32. Unique draft-prob rows per request. +//! \param topK Optional [batchSize], on GPU. +//! \param topP Optional [batchSize], on GPU. +//! \param targetVocabSize runtime::SizeType32. Output vocabulary size after optional d2t expansion. +//! \param d2t Optional [draftVocabSize], on GPU. +//! \return [batchSize, numDraftProbRows, targetVocabSize] float32 probabilities. +torch::Tensor computeDraftProbsForDynamicTreeRejection(torch::Tensor const& draftLogits, + torch::Tensor const& temperatures, runtime::SizeType32 numDraftProbRows, torch::optional const& topK, + torch::optional const& topP, runtime::SizeType32 targetVocabSize, bool skipTemperature, + torch::optional const& d2t, runtime::SizeType32 kMax = 0); + +//! \brief Compute target probabilities for dynamic-tree rejection sampling from logits. +//! \param targetLogits [batchSize * numDraftTokens, targetVocabSize], on GPU. +//! \param temperatures [batchSize], on GPU. +//! \param numDraftTokens runtime::SizeType32. Total tree nodes per request (including root). +//! \param topK Optional [batchSize], on GPU. +//! \param topP Optional [batchSize], on GPU. +//! \param kMax runtime::SizeType32. Max top-K value across the batch; enables the fast topk +//! path when > 0. Must be computed on CPU (e.g. topK.max().item()). Default 0 = fallback +//! to full sort. +//! \return Tuple of: +//! 1. [batchSize, numDraftTokens, targetVocabSize] float32 probabilities. +//! 2. [batchSize, numDraftTokens, maxTargetSupportSize] int32 compact token ids that survive +//! top-k/top-p filtering for each row, padded with -1. Empty when no filtering is active. +//! 3. [batchSize, numDraftTokens] int32 support lengths. Empty when no filtering is active. +std::tuple computeTargetProbsForDynamicTreeRejection( + torch::Tensor const& targetLogits, torch::Tensor const& temperatures, runtime::SizeType32 numDraftTokens, + torch::optional const& topK, torch::optional const& topP, bool skipTemperature, + runtime::SizeType32 kMax = 0); + } // namespace kernels::speculative_decoding TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp b/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp index 2c70fcfc5cb1..550bb8cde230 100644 --- a/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp +++ b/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp @@ -25,9 +25,44 @@ namespace tk = tensorrt_llm::kernels::speculative_decoding; TRTLLM_NAMESPACE_BEGIN +namespace kernels::speculative_decoding +{ +th::Tensor computeProbsFromLogits(th::Tensor const& logits, th::Tensor const& temperatures, + th::optional const& topK, th::optional const& topP, bool skipTemperature, + runtime::SizeType32 kMax); +void invokeBuildDraftProbIndices(int64_t const* topkScoreIndices, int32_t* draftProbIndices, + runtime::SizeType32 batchSize, runtime::SizeType32 topK, runtime::SizeType32 numDraftTokens, cudaStream_t stream); +th::Tensor computeDraftProbsForDynamicTreeRejection(th::Tensor const& draftLogits, th::Tensor const& temperatures, + runtime::SizeType32 numDraftProbRows, th::optional const& topK, th::optional const& topP, + runtime::SizeType32 targetVocabSize, bool skipTemperature, th::optional const& d2t, + runtime::SizeType32 kMax); +std::tuple computeTargetProbsForDynamicTreeRejection(th::Tensor const& targetLogits, + th::Tensor const& temperatures, runtime::SizeType32 numDraftTokens, th::optional const& topK, + th::optional const& topP, bool skipTemperature, runtime::SizeType32 kMax); +} // namespace kernels::speculative_decoding + namespace torch_ext { +void build_draft_prob_indices_out_op( + th::Tensor& topkScoreIndices, th::Tensor& draftProbIndices, int64_t topK, int64_t numDraftTokens) +{ + TORCH_CHECK(topkScoreIndices.is_cuda(), "topkScoreIndices must be a CUDA tensor"); + TORCH_CHECK(draftProbIndices.is_cuda(), "draftProbIndices must be a CUDA tensor"); + TORCH_CHECK(topkScoreIndices.dim() == 2, "topkScoreIndices must be a 2D tensor"); + TORCH_CHECK(draftProbIndices.dim() == 2, "draftProbIndices must be a 2D tensor"); + TORCH_CHECK(topkScoreIndices.scalar_type() == torch::kInt64, "topkScoreIndices must be int64 tensor"); + TORCH_CHECK(draftProbIndices.scalar_type() == torch::kInt32, "draftProbIndices must be int32 tensor"); + TORCH_CHECK(topkScoreIndices.size(1) == numDraftTokens, "topkScoreIndices size mismatch"); + TORCH_CHECK(draftProbIndices.size(0) == topkScoreIndices.size(0), "Batch size mismatch"); + TORCH_CHECK(draftProbIndices.size(1) == numDraftTokens + 1, "draftProbIndices size mismatch"); + TORCH_CHECK(numDraftTokens + 1 <= 1024, "numDraftTokens + 1 exceeds CUDA block size limit of 1024"); + + auto stream = at::cuda::getCurrentCUDAStream(topkScoreIndices.device().index()); + tk::invokeBuildDraftProbIndices(topkScoreIndices.data_ptr(), draftProbIndices.data_ptr(), + topkScoreIndices.size(0), topK, numDraftTokens, stream); +} + //! \brief Build dynamic tree structure (in-place, writes to pre-allocated output buffers) //! All index tensors use int64 to match PyTorch's default integer dtype. void build_dynamic_tree_op(th::Tensor& parentList, th::Tensor& selectedIndex, th::Tensor& treeMask, @@ -175,12 +210,145 @@ void verify_dynamic_tree_greedy_out_op(th::Tensor& candidates, th::Tensor& retri batchSize, numDraftTokens, numSpecStep, stream); } +//! \brief In-place tree rejection sampling verify op. +//! Accepts draft tokens by rejection sampling at each depth using pre-computed probabilities. +void verify_dynamic_tree_rejection_out_op(th::Tensor& candidates, th::Tensor& draftProbs, th::Tensor& targetProbs, + th::Tensor& targetSupportIndices, th::Tensor& targetSupportLengths, th::Tensor& draftProbIndices, + th::Tensor& retrieveNextToken, th::Tensor& retrieveNextSibling, th::Tensor& acceptIndex, th::Tensor& acceptTokenNum, + th::Tensor& acceptToken, int64_t numSpecStep, th::Tensor& seed, th::Tensor& offset) +{ + TORCH_CHECK(candidates.dim() == 2, "candidates must be 2D tensor"); + TORCH_CHECK(draftProbs.dim() == 3, "draftProbs must be 3D tensor"); + TORCH_CHECK(targetProbs.dim() == 3, "targetProbs must be 3D tensor"); + TORCH_CHECK(targetSupportIndices.dim() == 1 || targetSupportIndices.dim() == 3, + "targetSupportIndices must be 1D or 3D tensor"); + TORCH_CHECK(targetSupportLengths.dim() == 1 || targetSupportLengths.dim() == 2, + "targetSupportLengths must be 1D or 2D tensor"); + TORCH_CHECK(draftProbIndices.dim() == 2, "draftProbIndices must be 2D tensor"); + TORCH_CHECK(retrieveNextToken.dim() == 2, "retrieveNextToken must be 2D tensor"); + TORCH_CHECK(retrieveNextSibling.dim() == 2, "retrieveNextSibling must be 2D tensor"); + TORCH_CHECK(candidates.scalar_type() == torch::kInt64, "candidates must be int64 tensor"); + TORCH_CHECK(draftProbs.scalar_type() == torch::kFloat32, "draftProbs must be float32 tensor"); + TORCH_CHECK(targetProbs.scalar_type() == torch::kFloat32, "targetProbs must be float32 tensor"); + TORCH_CHECK(targetSupportIndices.scalar_type() == torch::kInt32, "targetSupportIndices must be int32 tensor"); + TORCH_CHECK(targetSupportLengths.scalar_type() == torch::kInt32, "targetSupportLengths must be int32 tensor"); + TORCH_CHECK(draftProbIndices.scalar_type() == torch::kInt32, "draftProbIndices must be int32 tensor"); + + int64_t batchSize = candidates.size(0); + int64_t numDraftProbRows = draftProbs.size(1); + int64_t numDraftTokens = candidates.size(1); + int64_t vocabSize = targetProbs.size(2); + int64_t maxTargetSupportSize = targetSupportIndices.dim() == 3 ? targetSupportIndices.size(2) : 0; + + TORCH_CHECK(draftProbs.size(0) == batchSize, "draftProbs batch size mismatch"); + TORCH_CHECK(draftProbs.size(2) == vocabSize, "draftProbs vocabSize mismatch"); + TORCH_CHECK(targetProbs.size(0) == batchSize, "targetProbs batch size mismatch"); + TORCH_CHECK(targetProbs.size(1) == numDraftTokens, "targetProbs numDraftTokens mismatch"); + if (targetSupportIndices.numel() > 0) + { + TORCH_CHECK(targetSupportIndices.dim() == 3, "targetSupportIndices must be 3D when non-empty"); + TORCH_CHECK(targetSupportIndices.size(0) == batchSize, "targetSupportIndices batch size mismatch"); + TORCH_CHECK(targetSupportIndices.size(1) == numDraftTokens, "targetSupportIndices numDraftTokens mismatch"); + } + if (targetSupportLengths.numel() > 0) + { + TORCH_CHECK(targetSupportLengths.dim() == 2, "targetSupportLengths must be 2D when non-empty"); + TORCH_CHECK(targetSupportLengths.size(0) == batchSize, "targetSupportLengths batch size mismatch"); + TORCH_CHECK(targetSupportLengths.size(1) == numDraftTokens, "targetSupportLengths numDraftTokens mismatch"); + } + TORCH_CHECK((targetSupportIndices.numel() == 0) == (targetSupportLengths.numel() == 0), + "targetSupportIndices and targetSupportLengths must both be empty or both be non-empty"); + TORCH_CHECK(draftProbIndices.size(0) == batchSize, "draftProbIndices batch size mismatch"); + TORCH_CHECK(draftProbIndices.size(1) == numDraftTokens, "draftProbIndices size mismatch"); + TORCH_CHECK(retrieveNextToken.size(0) == batchSize, "retrieveNextToken batch size mismatch"); + TORCH_CHECK(retrieveNextToken.size(1) == numDraftTokens, "retrieveNextToken size mismatch"); + TORCH_CHECK(retrieveNextSibling.size(0) == batchSize, "retrieveNextSibling batch size mismatch"); + TORCH_CHECK(retrieveNextSibling.size(1) == numDraftTokens, "retrieveNextSibling size mismatch"); + TORCH_CHECK(acceptIndex.scalar_type() == torch::kInt64, "acceptIndex must be int64 tensor"); + TORCH_CHECK(acceptTokenNum.scalar_type() == torch::kInt64, "acceptTokenNum must be int64 tensor"); + TORCH_CHECK(acceptToken.scalar_type() == torch::kInt64, "acceptToken must be int64 tensor"); + TORCH_CHECK(acceptIndex.size(0) >= batchSize && acceptIndex.size(1) >= numSpecStep, "acceptIndex buffer too small"); + TORCH_CHECK(acceptTokenNum.size(0) >= batchSize, "acceptTokenNum buffer too small"); + TORCH_CHECK(acceptToken.size(0) >= batchSize && acceptToken.size(1) >= numSpecStep, "acceptToken buffer too small"); + TORCH_CHECK(seed.scalar_type() == torch::kInt64, "seed must be int64 tensor"); + TORCH_CHECK(offset.scalar_type() == torch::kInt64, "offset must be int64 tensor"); + TORCH_CHECK(seed.numel() >= 1, "seed tensor must have at least one element"); + TORCH_CHECK(offset.numel() >= 1, "offset tensor must have at least one element"); + TORCH_CHECK(seed.is_cuda(), "seed must be CUDA tensor"); + TORCH_CHECK(offset.is_cuda(), "offset must be CUDA tensor"); + TORCH_CHECK(seed.device() == candidates.device(), "seed must be on the same device as candidates"); + TORCH_CHECK(offset.device() == candidates.device(), "offset must be on the same device as candidates"); + + auto stream = at::cuda::getCurrentCUDAStream(candidates.device().index()); + + acceptIndex.zero_(); + acceptTokenNum.zero_(); + acceptToken.zero_(); + + tk::invokeVerifyDynamicTreeRejection(acceptIndex.data_ptr(), acceptTokenNum.data_ptr(), + acceptToken.data_ptr(), candidates.data_ptr(), draftProbs.data_ptr(), + targetProbs.data_ptr(), + targetSupportIndices.numel() > 0 ? targetSupportIndices.data_ptr() : nullptr, + targetSupportLengths.numel() > 0 ? targetSupportLengths.data_ptr() : nullptr, + draftProbIndices.data_ptr(), retrieveNextToken.data_ptr(), + retrieveNextSibling.data_ptr(), batchSize, numDraftProbRows, maxTargetSupportSize, numDraftTokens, + numSpecStep, vocabSize, seed.data_ptr(), offset.data_ptr(), stream); +} + +th::Tensor compute_draft_probs_for_dynamic_tree_rejection_op(th::Tensor draftLogits, th::Tensor temperatures, + int64_t numDraftProbRows, int64_t targetVocabSize, th::optional topK, th::optional topP, + bool skipTemperature, th::optional d2t, int64_t topKMax) +{ + return tk::computeDraftProbsForDynamicTreeRejection( + draftLogits, temperatures, numDraftProbRows, topK, topP, targetVocabSize, skipTemperature, d2t, topKMax); +} + +std::tuple compute_target_probs_for_dynamic_tree_rejection_op( + th::Tensor targetLogits, th::Tensor temperatures, int64_t numDraftTokens, th::optional topK, + th::optional topP, bool skipTemperature, int64_t topKMax) +{ + return tk::computeTargetProbsForDynamicTreeRejection( + targetLogits, temperatures, numDraftTokens, topK, topP, skipTemperature, topKMax); +} + +th::Tensor compute_probs_from_logits_op(th::Tensor logits, th::Tensor temperatures, th::optional topK, + th::optional topP, bool skipTemperature) +{ + TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA tensor"); + TORCH_CHECK(temperatures.is_cuda(), "temperatures must be a CUDA tensor"); + TORCH_CHECK(logits.dim() == 2, "logits must be a 2D tensor"); + TORCH_CHECK(temperatures.dim() == 1, "temperatures must be a 1D tensor"); + TORCH_CHECK(logits.size(0) == temperatures.size(0), "logits and temperatures size mismatch"); + if (topK.has_value() && topK->defined()) + { + TORCH_CHECK(topK->is_cuda(), "top_k must be a CUDA tensor"); + } + if (topP.has_value() && topP->defined()) + { + TORCH_CHECK(topP->is_cuda(), "top_p must be a CUDA tensor"); + } + + return tk::computeProbsFromLogits(logits, temperatures, topK, topP, skipTemperature, /*kMax=*/0); +} + } // namespace torch_ext TRTLLM_NAMESPACE_END //////////////////////////////////////////////////////////////////////////////////////////////////////////// +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "build_draft_prob_indices_out_op(Tensor topkScoreIndices, Tensor(a!) draftProbIndices, " + "int topK, int numDraftTokens) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("build_draft_prob_indices_out_op", &tensorrt_llm::torch_ext::build_draft_prob_indices_out_op); +} + TORCH_LIBRARY_FRAGMENT(trtllm, m) { m.def( @@ -227,3 +395,62 @@ TORCH_LIBRARY_IMPL(trtllm, CUDA, m) { m.impl("verify_dynamic_tree_greedy_out_op", &tensorrt_llm::torch_ext::verify_dynamic_tree_greedy_out_op); } + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "verify_dynamic_tree_rejection_out_op(" + "Tensor candidates, Tensor draftProbs, Tensor targetProbs, Tensor targetSupportIndices, " + "Tensor targetSupportLengths, Tensor draftProbIndices, " + "Tensor retrieveNextToken, Tensor retrieveNextSibling, " + "Tensor(a!) acceptIndex, Tensor(b!) acceptTokenNum, Tensor(c!) acceptToken, " + "int numSpecStep, Tensor seed, Tensor offset) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("verify_dynamic_tree_rejection_out_op", &tensorrt_llm::torch_ext::verify_dynamic_tree_rejection_out_op); +} + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "compute_draft_probs_for_dynamic_tree_rejection_op(" + "Tensor draftLogits, Tensor temperatures, int numDraftProbRows, int targetVocabSize, " + "Tensor? top_k=None, Tensor? top_p=None, bool skip_temperature=False, Tensor? d2t=None, " + "int top_k_max=0) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("compute_draft_probs_for_dynamic_tree_rejection_op", + &tensorrt_llm::torch_ext::compute_draft_probs_for_dynamic_tree_rejection_op); +} + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "compute_target_probs_for_dynamic_tree_rejection_op(" + "Tensor targetLogits, Tensor temperatures, int numDraftTokens, " + "Tensor? top_k=None, Tensor? top_p=None, bool skip_temperature=False, int top_k_max=0) -> (Tensor, Tensor, " + "Tensor)"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("compute_target_probs_for_dynamic_tree_rejection_op", + &tensorrt_llm::torch_ext::compute_target_probs_for_dynamic_tree_rejection_op); +} + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "compute_probs_from_logits_op(" + "Tensor logits, Tensor temperatures, Tensor? top_k=None, Tensor? top_p=None, " + "bool skip_temperature=False) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("compute_probs_from_logits_op", &tensorrt_llm::torch_ext::compute_probs_from_logits_op); +} diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index b75c60424537..f7b84d31ccca 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -748,6 +748,11 @@ def _get_max_shape_warmup_requests( max_batch_size = min( self.batch_size, curr_max_num_tokens // (1 + self.max_total_draft_tokens) // self.max_beam_width) + dynamic_tree_warmup_max_batch_size = self._get_dynamic_tree_warmup_max_batch_size( + ) + if dynamic_tree_warmup_max_batch_size is not None: + max_batch_size = min(max_batch_size, + dynamic_tree_warmup_max_batch_size) warmup_requests_configs = [ (curr_max_num_tokens, 0), # max_num_tokens, pure context @@ -777,6 +782,14 @@ def _get_full_general_warmup_requests( # Deduplicate the warmup_configs while keeping the order. return list(dict.fromkeys(warmup_configs)) + def _get_dynamic_tree_warmup_max_batch_size(self) -> Optional[int]: + """Return the dynamic-tree worker batch capacity for warmup batch clamping.""" + if (self.spec_config is None or self.is_draft_model + or not self.spec_config.spec_dec_mode.use_one_engine() + or not getattr(self.spec_config, "use_dynamic_tree", False)): + return None + return getattr(self.spec_config, "max_batch_size", None) + @with_warmup_flag @warmup_with_kv_cache_cleanup def warmup(self, resource_manager: ResourceManager) -> None: @@ -1030,6 +1043,12 @@ def _capture_generation_cuda_graphs(self, # Determine which graphs to capture graphs_to_capture = self._get_graphs_to_capture(cuda_graph_batch_sizes, spec_resource_manager) + dynamic_tree_warmup_max_batch_size = self._get_dynamic_tree_warmup_max_batch_size( + ) + if dynamic_tree_warmup_max_batch_size is not None: + graphs_to_capture = [(bs, draft_len) + for bs, draft_len in graphs_to_capture + if bs <= dynamic_tree_warmup_max_batch_size] graphs_to_capture = sorted(graphs_to_capture, reverse=True) # Create CUDA graphs for short and long sequences separately for sparse attention. # self.max_seq_len is the global max sequence length. For Helix CP each diff --git a/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py b/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py index 2565e034aaf1..849cd87b3e3a 100644 --- a/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py +++ b/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py @@ -70,6 +70,38 @@ def __init__( max_batch_size, max_path_len, dtype=torch.int64, device=device ) + # Pre-allocated output buffers for verify_dynamic_tree_rejection_out + self._rej_accept_index_buf = torch.zeros( + max_batch_size, max_path_len, dtype=torch.int64, device=device + ) + self._rej_accept_token_num_buf = torch.zeros( + max_batch_size, dtype=torch.int64, device=device + ) + self._rej_accept_token_buf = torch.zeros( + max_batch_size, max_path_len, dtype=torch.int64, device=device + ) + self._rej_seed_buf = torch.zeros(1, dtype=torch.int64, device=device) + self._rej_offset_buf = torch.zeros(1, dtype=torch.int64, device=device) + + def _get_rejection_rng_tensor( + self, + value: int | torch.Tensor, + buffer: torch.Tensor, + name: str, + ) -> torch.Tensor: + if isinstance(value, int): + buffer.fill_(value) + return buffer + if not isinstance(value, torch.Tensor): + raise TypeError(f"{name} must be an int or torch.Tensor, got {type(value)!r}") + if value.dtype != torch.int64: + raise TypeError(f"{name} must be int64 tensor, got {value.dtype}") + if value.numel() < 1: + raise ValueError(f"{name} tensor must have at least one element") + if value.device != buffer.device: + raise ValueError(f"{name} tensor must be on device {buffer.device}, got {value.device}") + return value.reshape(-1)[:1] + def build_dynamic_tree( self, history_draft_tokens_parent_buffer: torch.Tensor, @@ -175,19 +207,23 @@ def verify_dynamic_tree_greedy_out( tree_valid = torch.ones(num_gens, dtype=torch.bool, device=candidates.device) try: - torch.ops.trtllm.verify_dynamic_tree_greedy_out_op( - candidates, - retrieve_index, - retrieve_next_token, - retrieve_next_sibling, - target_predict, - predicts, - accept_index, - accept_token_num, - accept_token, - tree_valid, - num_spec_step, - ) + torch.cuda.nvtx.range_push("verify_dynamic_tree_greedy_out") + try: + torch.ops.trtllm.verify_dynamic_tree_greedy_out_op( + candidates, + retrieve_index, + retrieve_next_token, + retrieve_next_sibling, + target_predict, + predicts, + accept_index, + accept_token_num, + accept_token, + tree_valid, + num_spec_step, + ) + finally: + torch.cuda.nvtx.range_pop() except Exception as e: raise RuntimeError( f"verify_dynamic_tree_greedy_out_op failed: {e}\n" @@ -196,3 +232,114 @@ def verify_dynamic_tree_greedy_out( ) from e return predicts, accept_index, accept_token_num, accept_token + + def verify_dynamic_tree_rejection_from_logits_out( + self, + candidates: torch.Tensor, + draft_logits_tree: torch.Tensor, + target_logits_tree: torch.Tensor, + draft_prob_indices: torch.Tensor, + retrieve_next_token: torch.Tensor, + retrieve_next_sibling: torch.Tensor, + temperatures: torch.Tensor, + top_k: torch.Tensor | None, + top_p: torch.Tensor | None, + skip_temperature: bool, + num_gens: int, + num_spec_step: int, + seed: int | torch.Tensor = 0, + offset: int | torch.Tensor = 0, + d2t: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Tree-aware rejection sampling from logits (three CUDA ops). + + This path keeps draft/target logits as inputs, computes unique draft + and target probabilities with separate CUDA ops, then runs the tree + rejection kernel as a third CUDA op. `draft_prob_indices` maps each + tree position to its shared draft-prob row. + """ + accept_index = self._rej_accept_index_buf[:num_gens] + accept_token = self._rej_accept_token_buf[:num_gens] + accept_tok_num = self._rej_accept_token_num_buf[:num_gens] + seed_tensor = self._get_rejection_rng_tensor(seed, self._rej_seed_buf, "seed") + offset_tensor = self._get_rejection_rng_tensor(offset, self._rej_offset_buf, "offset") + num_draft_tokens = candidates.shape[1] + num_draft_prob_rows = draft_logits_tree.shape[0] // num_gens + target_vocab_size = target_logits_tree.shape[-1] + + top_k_max = int(top_k.max().item()) if top_k is not None else 0 + + try: + torch.cuda.nvtx.range_push("verify_dynamic_tree_rejection_from_logits_out") + try: + torch.cuda.nvtx.range_push( + "trtllm::compute_draft_probs_for_dynamic_tree_rejection_op" + ) + try: + draft_probs_tree = ( + torch.ops.trtllm.compute_draft_probs_for_dynamic_tree_rejection_op( + draft_logits_tree, + temperatures, + num_draft_prob_rows, + target_vocab_size, + top_k, + top_p, + skip_temperature, + d2t=d2t, + top_k_max=top_k_max, + ) + ) + finally: + torch.cuda.nvtx.range_pop() + + torch.cuda.nvtx.range_push( + "trtllm::compute_target_probs_for_dynamic_tree_rejection_op" + ) + try: + ( + target_probs_tree, + target_support_indices, + target_support_lengths, + ) = torch.ops.trtllm.compute_target_probs_for_dynamic_tree_rejection_op( + target_logits_tree, + temperatures, + num_draft_tokens, + top_k, + top_p, + skip_temperature, + top_k_max=top_k_max, + ) + finally: + torch.cuda.nvtx.range_pop() + + torch.cuda.nvtx.range_push("trtllm::verify_dynamic_tree_rejection_out_op") + try: + torch.ops.trtllm.verify_dynamic_tree_rejection_out_op( + candidates, + draft_probs_tree, + target_probs_tree, + target_support_indices, + target_support_lengths, + draft_prob_indices, + retrieve_next_token, + retrieve_next_sibling, + accept_index, + accept_tok_num, + accept_token, + num_spec_step, + seed_tensor, + offset_tensor, + ) + finally: + torch.cuda.nvtx.range_pop() + finally: + torch.cuda.nvtx.range_pop() + except Exception as e: + raise RuntimeError( + f"dynamic tree rejection op chain failed: {e}\n" + f"Inputs: num_gens={num_gens}, N={candidates.shape[1]}, " + f"draft_vocab={draft_logits_tree.shape[-1]}, " + f"target_vocab={target_logits_tree.shape[-1]}, num_spec_step={num_spec_step}" + ) from e + + return target_support_indices, accept_index, accept_tok_num, accept_token diff --git a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py index f33a5824a455..09995b04fd99 100644 --- a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py @@ -14,7 +14,7 @@ # limitations under the License. """Eagle3 one-model dynamic tree speculative decoding.""" -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional import torch import triton @@ -279,6 +279,26 @@ def __init__( sm = get_sm_version() self._needs_mask_repack = sm < 100 or sm in (120, 121) + # Rejection sampling buffers + # Unique draft logits per request, one row per distinct parent context. + # Shape: [max_batch_size, 1 + (max_draft_len - 1) * K, vocab_size] + # row 0 -> p(.|root) + # rows [1 : 1 + K] -> p(.|depth-0 parent_k) + # rows [1 + K : 1 + 2K] -> p(.|depth-1 parent_k) + # This avoids materializing repeated per-tree-position logits such as + # [p(.|E), p(.|E), p(.|E), p(.|F1), p(.|F1), p(.|F2)]. + self._draft_depth_logits_cat: Optional[torch.Tensor] = None + # topk_score_indices from resampling_final_draft_tokens for path tracing. + # Shape: [max_batch_size, max_total_draft_tokens] + self._topk_score_indices_buf = torch.zeros( + max_batch_size, self.max_total_draft_tokens, dtype=torch.int64, device="cuda" + ) + # Tree position -> unique draft-prob row mapping for rejection sampling. + # Shape: [max_batch_size, max_total_draft_tokens + 1], root row is unused and kept at 0. + self._draft_prob_indices_buf = torch.zeros( + max_batch_size, self.max_total_draft_tokens + 1, dtype=torch.int32, device="cuda" + ) + def _repack_mask_padded_to_packed(self, mask_buf, n_req, n_tok): """XQA indexes mask flat via cuQSeqLens; padded [n_req, buf_dim, maskW] has batch stride buf_dim*maskW, so when n_tok < buf_dim and n_req > 1 packed @@ -371,10 +391,6 @@ def _relocate_kv_eagerly(self, attn_metadata, batch_size): if cache_mgr is None: return - assert len(set(cache_mgr.num_kv_heads_per_layer)) == 1, ( - "update_kv_cache_draft_token_location_2d requires uniform num_kv_heads across all layers, " - f"but got {cache_mgr.num_kv_heads_per_layer}" - ) torch.ops.tensorrt_llm.update_kv_cache_draft_token_location_2d( self._accepted_draft_indices_tensor[:batch_size], self._num_accepted_tokens_buf[:batch_size], @@ -524,6 +540,7 @@ 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}" @@ -577,6 +594,11 @@ def _forward_draft_loop( hidden_states[gather_ids], draft_model.lm_head, attn_metadata, True ) + # Capture unique draft logits for the root parent context. + if spec_metadata.use_rejection_sampling and num_gens > 0: + self._lazy_alloc_draft_logits_buf(logits.shape[-1], logits.dtype, logits.device) + self._draft_depth_logits_cat[:num_gens, 0].copy_(logits[num_contexts:]) + new_draft_tokens, new_draft_scores = self.sample( logits, self.K, draft_model=draft_model ) @@ -643,6 +665,13 @@ def _forward_draft_loop( selected_hs, draft_model.lm_head, attn_metadata, True ) + # Capture unique draft logits for the K parent contexts at this depth. + if spec_metadata.use_rejection_sampling and num_gens > 0: + row_start = 1 + (layer_idx - 1) * self.K + self._draft_depth_logits_cat[:num_gens, row_start : row_start + self.K].copy_( + logits[num_contexts * self.K :].reshape(num_gens, self.K, -1) + ) + new_draft_tokens, new_draft_scores = self.sample( logits, self.K, draft_model=draft_model ) @@ -672,6 +701,13 @@ def _forward_draft_loop( # Resample final tokens and build tree real_draft_tokens, topk_score_indices = self.resampling_final_draft_tokens(batch_size) + # Save topk_score_indices for rejection sampling path tracing. + # Rejection sampling needs to map each final draft token back to its history + # buffer index to retrieve the corresponding draft logits. Greedy verification + # doesn't need this mapping since it only compares token IDs. + if spec_metadata.use_rejection_sampling: + self._topk_score_indices_buf[:batch_size].copy_(topk_score_indices) + if spec_tree_manager is not None: # Gen-only; spec-dec batch is [:num_gens] (not num_contexts-offset). self.tree_ops_converter.build_dynamic_tree( @@ -698,6 +734,11 @@ def _sample_and_accept_dynamic_tree( self, logits, attn_metadata, spec_metadata, batch_size, num_contexts, num_gens ): """Dynamic tree verification using CUDA kernel.""" + if num_gens > self._max_batch_size: + raise RuntimeError( + f"Dynamic tree batch size {num_gens} exceeds configured " + f"max_batch_size {self._max_batch_size}" + ) N = self.tokens_per_gen_step max_path_len = self._max_path_len @@ -739,6 +780,76 @@ def _sample_and_accept_dynamic_tree( ] tree_valid = spec_tree_manager.slot_has_tree[gen_slot_ids] + if self._can_use_rejection_sampling(spec_metadata): + vocab_size = logits.shape[-1] + num_ctx_tokens = logits.shape[0] - num_gens * N + device = logits.device + + draft_logits_tree = self._get_unique_draft_logits(num_gens) + draft_prob_indices = self._build_draft_prob_indices(num_gens) + target_logits_tree = logits[num_ctx_tokens:].reshape(-1, vocab_size) + gen_slice = slice(num_contexts, num_contexts + num_gens) + skip_top_k = getattr(spec_metadata, "skip_top_k", False) + skip_top_p = getattr(spec_metadata, "skip_top_p", False) + skip_temperature = getattr(spec_metadata, "skip_temperature", False) + + if spec_metadata.temperatures is None: + temps = torch.ones(num_gens, dtype=torch.float32, device=device) + skip_temperature = True + else: + temps = spec_metadata.temperatures[gen_slice] + + top_ks = None + if not skip_top_k and spec_metadata.top_ks is not None: + top_ks = spec_metadata.top_ks[gen_slice] + + top_ps = None + if not skip_top_p and spec_metadata.top_ps is not None: + top_ps = spec_metadata.top_ps[gen_slice] + + # Lazily initialize seed/offset tensors on correct device + 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) + # Use in-place operations for CUDA graph compatibility + self.seed.add_(1).remainder_(2**31) + + _, accept_index, accept_token_num, accept_token = ( + self.tree_ops_converter.verify_dynamic_tree_rejection_from_logits_out( + candidates, + draft_logits_tree, + target_logits_tree, + draft_prob_indices, + spec_tree_manager.retrieve_next_token[:num_gens], + spec_tree_manager.retrieve_next_sibling[:num_gens], + temps, + top_ks, + top_ps, + skip_temperature, + num_gens, + self._max_path_len, + seed=self.seed, + offset=self.offset, + d2t=self._d2t, + ) + ) + + self._finalize_dynamic_tree_verify_outputs( + accept_index=accept_index, + accept_token_num=accept_token_num, + accept_token=accept_token, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_contexts=num_contexts, + batch_size=batch_size, + num_gens=num_gens, + max_path_len=max_path_len, + ) + num_accepted_tokens = self._apply_force_accepted_tokens( + num_accepted_tokens, num_contexts, self.max_draft_len + ) + return accepted_tokens, num_accepted_tokens + _, accept_index, accept_token_num, accept_token = ( self.tree_ops_converter.verify_dynamic_tree_greedy_out( candidates, @@ -752,14 +863,17 @@ def _sample_and_accept_dynamic_tree( ) ) - self._accept_token = accept_token - n_acc_draft = accept_token_num[:num_gens] - num_accepted_tokens[num_contexts:batch_size] = (n_acc_draft + 1).to(torch.int32) - accepted_tokens[num_contexts:batch_size] = accept_token[:num_gens].to(torch.int32) - # accept_index is 0-based from kernel; -1 converts padding (0) to sentinel (-1) - self._accepted_draft_indices_tensor[num_contexts:batch_size] = ( - accept_index[:num_gens, 1:max_path_len] - 1 - ).to(torch.int32) + self._finalize_dynamic_tree_verify_outputs( + accept_index=accept_index, + accept_token_num=accept_token_num, + accept_token=accept_token, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_contexts=num_contexts, + batch_size=batch_size, + num_gens=num_gens, + max_path_len=max_path_len, + ) num_accepted_tokens = self._apply_force_accepted_tokens( num_accepted_tokens, num_contexts, self.max_draft_len @@ -767,6 +881,91 @@ def _sample_and_accept_dynamic_tree( return accepted_tokens, num_accepted_tokens + def _can_use_rejection_sampling(self, spec_metadata) -> bool: + """Check if rejection sampling can be used for dynamic tree verification. + + Dynamic tree uses its own compact unique-logit buffer + (_draft_depth_logits_cat) instead of spec_metadata.draft_logits, so we + check that buffer's allocation status rather than draft_logits_valid. + The buffer is lazily allocated and populated during the first forward + pass with generation requests. + + Args: + spec_metadata: Speculative decoding metadata + + Returns: + True if rejection sampling is enabled and the draft logit buffer is allocated + """ + return spec_metadata.use_rejection_sampling and self._draft_depth_logits_cat is not None + + def _finalize_dynamic_tree_verify_outputs( + self, + accept_index: torch.Tensor, + accept_token_num: torch.Tensor, + accept_token: torch.Tensor, + accepted_tokens: torch.Tensor, + num_accepted_tokens: torch.Tensor, + num_contexts: int, + batch_size: int, + num_gens: int, + max_path_len: int, + ) -> None: + """Write dynamic-tree verify outputs back to the shared step buffers.""" + self._accept_token = accept_token + accepted_draft_count = accept_token_num[:num_gens] + num_accepted_tokens[num_contexts:batch_size] = (accepted_draft_count + 1).to(torch.int32) + accepted_tokens[num_contexts:batch_size] = accept_token[:num_gens].to(torch.int32) + # accept_index stores root at slot 0. Subtract 1 so padding/root value 0 becomes sentinel -1. + self._accepted_draft_indices_tensor[num_contexts:batch_size] = ( + accept_index[:num_gens, 1:max_path_len] - 1 + ).to(torch.int32) + + def _lazy_alloc_draft_logits_buf(self, vocab_size: int, dtype, device): + """Lazily allocate unique draft-logit capture buffer.""" + if self._draft_depth_logits_cat is None: + rows = 1 + (self.max_draft_len - 1) * self.K + self._draft_depth_logits_cat = torch.empty( + self._max_batch_size, rows, vocab_size, dtype=dtype, device=device + ) + + def _get_unique_draft_logits( + self, + num_gens: int, + ) -> torch.Tensor: + """Return compact unique draft logits for rejection sampling. + + The returned rows are unique per parent context, not duplicated per + final tree position. + + Args: + num_gens: Number of generation requests. + + Returns: + draft_logits: [num_gens * (1 + (max_draft_len - 1) * K), draft_vocab_size] + """ + return self._draft_depth_logits_cat[:num_gens].reshape( + -1, self._draft_depth_logits_cat.shape[-1] + ) + + def _build_draft_prob_indices( + self, + num_gens: int, + ) -> torch.Tensor: + """Build tree-position -> unique draft-prob row mapping. + + For a final tree position: + - depth 0 children map to row 0, which stores p(.|root) + - deeper nodes map to row 1 + depth_bucket * K + parent_k + """ + draft_prob_indices = self._draft_prob_indices_buf[:num_gens] + torch.ops.trtllm.build_draft_prob_indices_out_op( + self._topk_score_indices_buf[:num_gens], + draft_prob_indices, + self.K, + self.max_total_draft_tokens, + ) + return draft_prob_indices + @nvtx_range("eagle3_dyn.sample") def sample( self, logits: torch.Tensor, max_top_k: int, draft_model=None diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index e4586b2142eb..509bd01c72cc 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -444,6 +444,10 @@ class SpecMetadata: temperatures: Optional[torch.Tensor] = None top_ks: Optional[torch.Tensor] = None top_ps: Optional[torch.Tensor] = None + # Whether top-k/top-p/temperature are globally disabled for the current batch. + skip_temperature: bool = False + skip_top_k: bool = False + skip_top_p: bool = False # Sampling parameters indexed per request. request_temperatures: Optional[torch.Tensor] = None request_top_ks: Optional[torch.Tensor] = None @@ -528,6 +532,9 @@ def populate_sampling_params_for_one_model( request_temperatures = [] request_top_ks = [] request_top_ps = [] + top_k_enabled = False + top_p_enabled = False + temperature_enabled = False # Need to use a very small value for temperature when disabled to avoid division by 0 DISABLE_TEMP_VAL = 1e-5 @@ -535,29 +542,67 @@ def populate_sampling_params_for_one_model( DISABLE_TOPK_VAL = torch.iinfo(torch.int32).max DISABLE_TOPP_VAL = 1.0 - for request in requests: - sampling_config = request.sampling_config - temp = sampling_config.temperature - temp_val = temp[0] if temp is not None and len(temp) > 0 else None + def _first_or_none(values): + return values[0] if values is not None and len(values) > 0 else None + + def _normalize_request_sampling_params( + *, + temperature: Optional[float], + top_k: Optional[int], + top_p: Optional[float], + ) -> tuple[float, int, float, bool, bool, bool]: + is_greedy = SamplingParams.params_imply_greedy_decoding( + temperature=temperature, + top_k=top_k, + top_p=top_p, + use_beam_search=False) - tk = sampling_config.top_k - tk_val = tk[0] if tk is not None and len(tk) > 0 else None + use_temperature = (not is_greedy + and temperature not in (None, 0, 1)) + use_top_k = not is_greedy and top_k is not None and top_k > 0 + use_top_p = not is_greedy and top_p is not None and top_p < 1.0 + + normalized_temperature = (DISABLE_TEMP_VAL + if is_greedy or temperature is None + or temperature == 0 else temperature) + normalized_top_k = DISABLE_TOPK_VAL if not use_top_k else top_k + normalized_top_p = (DISABLE_TOPP_VAL + if is_greedy or top_p is None else top_p) + + return ( + normalized_temperature, + normalized_top_k, + normalized_top_p, + use_temperature, + use_top_k, + use_top_p, + ) - tp = sampling_config.top_p - tp_val = tp[0] if tp is not None and len(tp) > 0 else None + for request in requests: + sampling_config = request.sampling_config + temp_val = _first_or_none(sampling_config.temperature) + tk_val = _first_or_none(sampling_config.top_k) + tp_val = _first_or_none(sampling_config.top_p) # Context requests have no draft tokens yet. num_tokens = 1 + self.runtime_draft_len if request.state == LlmRequestState.GENERATION_IN_PROGRESS else 1 - is_greedy = SamplingParams.params_imply_greedy_decoding( + ( + temp_val, + tk_val, + tp_val, + use_temperature, + use_top_k, + use_top_p, + ) = _normalize_request_sampling_params( temperature=temp_val, top_k=tk_val, top_p=tp_val, - use_beam_search=False) + ) - temp_val = DISABLE_TEMP_VAL if is_greedy or temp_val is None or temp_val == 0 else temp_val - tk_val = DISABLE_TOPK_VAL if is_greedy or tk_val is None or tk_val <= 0 else tk_val - tp_val = DISABLE_TOPP_VAL if is_greedy or tp_val is None else tp_val + temperature_enabled |= use_temperature + top_k_enabled |= use_top_k + top_p_enabled |= use_top_p request_temperatures.append(temp_val) request_top_ks.append(tk_val) @@ -566,19 +611,21 @@ def populate_sampling_params_for_one_model( top_ks.extend(tk_val for _ in range(num_tokens)) top_ps.extend(tp_val for _ in range(num_tokens)) - if self.temperatures is None: - self.temperatures = torch.ones( - (self.max_draft_len + 1) * self.max_num_requests, - dtype=torch.float32, - device='cuda') - self.top_ks = torch.zeros( - (self.max_draft_len + 1) * self.max_num_requests, - dtype=torch.int32, - device='cuda') - self.top_ps = torch.ones( - (self.max_draft_len + 1) * self.max_num_requests, - dtype=torch.float32, - device='cuda') + tokens_per_request = (self.max_total_draft_tokens + 1 if + self.is_spec_dec_tree else self.max_draft_len + 1) + required_flat_size = tokens_per_request * self.max_num_requests + + if self.temperatures is None or self.temperatures.numel( + ) < required_flat_size: + self.temperatures = torch.ones(required_flat_size, + dtype=torch.float32, + device='cuda') + self.top_ks = torch.zeros(required_flat_size, + dtype=torch.int32, + device='cuda') + self.top_ps = torch.ones(required_flat_size, + dtype=torch.float32, + device='cuda') self.request_temperatures = torch.ones(self.max_num_requests, dtype=torch.float32, device='cuda') @@ -609,6 +656,9 @@ def populate_sampling_params_for_one_model( self.request_top_ps[:len(request_top_ps)].copy_(torch.tensor( request_top_ps, dtype=torch.float32, pin_memory=prefer_pinned()), non_blocking=True) + self.skip_temperature = not temperature_enabled + self.skip_top_k = not top_k_enabled + self.skip_top_p = not top_p_enabled class SpecWorkerBase(nn.Module, ABC): diff --git a/tensorrt_llm/_torch/speculative/one_model_sampler.py b/tensorrt_llm/_torch/speculative/one_model_sampler.py index 9949a6f588ee..ae1dda63e658 100644 --- a/tensorrt_llm/_torch/speculative/one_model_sampler.py +++ b/tensorrt_llm/_torch/speculative/one_model_sampler.py @@ -111,13 +111,28 @@ def compute_probs_from_logits( temperatures: torch.Tensor, top_k: Optional[torch.Tensor], top_p: Optional[torch.Tensor], + skip_temperature: bool = False, ) -> torch.Tensor: """ Compute probabilities from logits with temperature, top-k, and top-p applied. """ - logits = apply_temperature(logits, temperatures) + if logits.is_cuda and hasattr(torch.ops.trtllm, "compute_probs_from_logits_op"): + return torch.ops.trtllm.compute_probs_from_logits_op( + logits, temperatures, top_k, top_p, skip_temperature + ) + + if not skip_temperature: + logits = apply_temperature(logits, temperatures) logits = apply_top_k_top_p(logits, top_k, top_p) - return logits.softmax(dim=-1, dtype=torch.float32) + probs = logits.softmax(dim=-1, dtype=torch.float32) + + # Greedy rows should remain exactly one-hot so rejection sampling does not + # spuriously reject numerically-near argmax tokens. + greedy_temp_threshold = 1e-4 + is_greedy = temperatures <= greedy_temp_threshold + argmax_ids = logits.argmax(dim=-1, keepdim=True) + one_hot = torch.zeros_like(probs).scatter_(1, argmax_ids, 1.0) + return torch.where(is_greedy.unsqueeze(1), one_hot, probs) def rejection_sampling_one_model( From c60e39c504df80a48f86425fb2a61d35b356696c Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 22 Apr 2026 00:26:21 -0700 Subject: [PATCH 03/17] [None][perf] optimize and instrument Eagle3 dynamic-tree rejection verification Signed-off-by: ZhaoyangWang --- .../speculativeDecoding/dynamicTreeKernels.cu | 372 +++++++++++------- .../_torch/speculative/eagle3_dynamic_tree.py | 334 +++++++++------- 2 files changed, 417 insertions(+), 289 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu index 400daee9f21d..e2a4dd27a225 100644 --- a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu +++ b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu @@ -1179,6 +1179,7 @@ __device__ int64_t sampleFromIndexedDistribution(curandStatePhilox4_32_10_t& sta //! \param vocabSize vocabulary size. //! \param seed [1] int64 on GPU. Philox RNG seed. //! \param offset [1] int64 on GPU. Philox RNG offset. +template __global__ void verifyDynamicTreeRejectionKernel(int64_t* acceptIndex, int64_t* acceptTokenNum, int64_t* acceptToken, int64_t const* candidates, float const* draftProbs, float const* targetProbs, int32_t const* targetSupportIndices, int32_t const* targetSupportLengths, int32_t const* draftProbIndices, int32_t const* retrieveNextToken, @@ -1187,15 +1188,42 @@ __global__ void verifyDynamicTreeRejectionKernel(int64_t* acceptIndex, int64_t* int64_t const* offset) { uint32_t bx = blockIdx.x; + int32_t const tid = static_cast(threadIdx.x); if (bx >= batchSize) { return; } + using BlockReduce = cub::BlockReduce; + using BlockScan = cub::BlockScan; + + __shared__ union + { + typename BlockReduce::TempStorage reduce; + typename BlockScan::TempStorage scan; + } tempStorage; + + __shared__ int32_t sLastAcceptedLocalIdx; + __shared__ uint32_t sNumAcceptedTokens; + __shared__ int32_t sCurIndex; + __shared__ int32_t sFirstChild; + __shared__ bool sHasTerminalToken; + __shared__ bool sAcceptedSibling; + __shared__ float sDiffSum; + __shared__ float sTargetMass; + __shared__ float sPrefixBase; + __shared__ int32_t sWinnerIndex; + __shared__ int64_t sSampledToken; + uint32_t batchOffset = bx * numDraftTokens; curandStatePhilox4_32_10_t state; - curand_init(static_cast(seed[0]), static_cast(bx), static_cast(offset[0]), &state); + if (tid == 0) + { + curand_init( + static_cast(seed[0]), static_cast(bx), static_cast(offset[0]), &state); + } + __syncthreads(); // Root (depth 0): initialize path state at tree position 0. // @@ -1223,11 +1251,15 @@ __global__ void verifyDynamicTreeRejectionKernel(int64_t* acceptIndex, int64_t* // slot numAcceptedTokens = final bonus/correction token // - acceptTokenNum stores the number of accepted draft tokens only. The caller // adds 1 to obtain the total number of emitted tokens. - int32_t lastAcceptedLocalIdx = 0; - acceptIndex[bx * numSpeculativeTokens] = lastAcceptedLocalIdx; - uint32_t numAcceptedTokens = 0; - int32_t curIndex = 0; - bool hasTerminalToken = false; + if (tid == 0) + { + sLastAcceptedLocalIdx = 0; + acceptIndex[bx * numSpeculativeTokens] = sLastAcceptedLocalIdx; + sNumAcceptedTokens = 0; + sCurIndex = 0; + sHasTerminalToken = false; + } + __syncthreads(); bool const hasCompactTargetSupport = targetSupportIndices != nullptr && targetSupportLengths != nullptr; for (uint32_t j = 1; j < numSpeculativeTokens; ++j) @@ -1236,92 +1268,83 @@ __global__ void verifyDynamicTreeRejectionKernel(int64_t* acceptIndex, int64_t* // Continuing the example above: // j = 1, curIndex = 0 (E) -> firstChild = F1 // j = 2, curIndex = 1 (F1) -> firstChild = G1 - int32_t firstChild = retrieveNextToken[batchOffset + curIndex]; - curIndex = firstChild; + if (tid == 0) + { + sFirstChild = retrieveNextToken[batchOffset + sLastAcceptedLocalIdx]; + sCurIndex = sFirstChild; + sAcceptedSibling = false; + } + __syncthreads(); - while (curIndex != -1) + while (sCurIndex != -1 && !sAcceptedSibling) { - int32_t draftLocalIdx = curIndex; // retrieveIndex is identity: draftLocalIdx == curIndex - int64_t draftTokenId = candidates[batchOffset + curIndex]; - // draftProbs: tree position curIndex maps to a unique parent-context row. - // For the example: - // curIndex = 1 (F1) -> draftProbIndices[1] = 0 -> p(F1|E) - // curIndex = 2 (F2) -> draftProbIndices[2] = 0 -> p(F2|E) - // curIndex = 4 (G1) -> draftProbIndices[4] = 1 -> p(G1|F1) - int32_t const draftProbRow = draftProbIndices[batchOffset + curIndex]; - float pDraft - = draftProbs[(static_cast(bx) * numDraftProbRows + draftProbRow) * vocabSize + draftTokenId]; - // Rejection sampling compares draft siblings under the target - // distribution of the currently accepted parent node. - float pTarget = targetProbs[(static_cast(bx) * numDraftTokens + lastAcceptedLocalIdx) * vocabSize - + draftTokenId]; - - // Rejection test for the current sibling: - // accept with probability min(1, q(x) / p(x)) - // where: - // p(x) = pDraft from the draft model proposal - // q(x) = pTarget from the target model verification distribution - float acceptProb = fminf(1.0f, pTarget / (pDraft + 1e-10f)); - - float u = curand_uniform(&state); - - if (u < acceptProb) + if (tid == 0) { - // Accepted sibling: extend the path and continue with this node's children. - // Example: - // if F1 is accepted at j = 1, then on the next iteration we descend to - // F1's first child (G1) and compare G1/G2 in sibling order. - // Emit the accepted draft token immediately. If we later stop at a leaf - // (or hit max depth), the final bonus token from the target distribution - // will be written at slot numAcceptedTokens. - acceptToken[bx * numSpeculativeTokens + numAcceptedTokens] = draftTokenId; - ++numAcceptedTokens; - acceptIndex[bx * numSpeculativeTokens + numAcceptedTokens] = draftLocalIdx; - lastAcceptedLocalIdx = draftLocalIdx; - break; - } - else - { - // Reject this sibling and try the next sibling at the same depth. - // Example: - // reject F1 -> move to F2 - // reject F2 -> move to F3 - curIndex = retrieveNextSibling[batchOffset + curIndex]; + int32_t const draftLocalIdx = sCurIndex; // retrieveIndex is identity: draftLocalIdx == curIndex + int64_t const draftTokenId = candidates[batchOffset + sCurIndex]; + int32_t const draftProbRow = draftProbIndices[batchOffset + sCurIndex]; + float const pDraft + = draftProbs[(static_cast(bx) * numDraftProbRows + draftProbRow) * vocabSize + + draftTokenId]; + float const pTarget + = targetProbs[(static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * vocabSize + + draftTokenId]; + + float const acceptProb = fminf(1.0f, pTarget / (pDraft + 1e-10f)); + float const u = curand_uniform(&state); + + if (u < acceptProb) + { + acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] = draftTokenId; + ++sNumAcceptedTokens; + acceptIndex[bx * numSpeculativeTokens + sNumAcceptedTokens] = draftLocalIdx; + sLastAcceptedLocalIdx = draftLocalIdx; + sAcceptedSibling = true; + } + else + { + sCurIndex = retrieveNextSibling[batchOffset + sCurIndex]; + } } + __syncthreads(); } - if (curIndex == -1) + if (sCurIndex == -1) { // All siblings exhausted. Two sub-cases: // (a) firstChild == -1: leaf node, no draft tokens at this depth. // Emit the final bonus token sampled from q(.|lastAcceptedLocalIdx). // (b) firstChild != -1: every sibling was rejected -> sample correction token // from relu(q - p) at firstChild's position to restore the target distribution. - if (firstChild == -1) + if (sFirstChild == -1) { - float const* tProbs - = targetProbs + (static_cast(bx) * numDraftTokens + lastAcceptedLocalIdx) * vocabSize; - if (hasCompactTargetSupport) + if (tid == 0) { - uint32_t const supportOffset - = (static_cast(bx) * numDraftTokens + lastAcceptedLocalIdx) * maxTargetSupportSize; - uint32_t const supportSize - = static_cast(targetSupportLengths[batchOffset + lastAcceptedLocalIdx]); - acceptToken[bx * numSpeculativeTokens + numAcceptedTokens] = sampleFromIndexedDistribution( - state, tProbs, targetSupportIndices + supportOffset, supportSize, vocabSize); - } - else - { - acceptToken[bx * numSpeculativeTokens + numAcceptedTokens] - = sampleFromDistribution(state, tProbs, vocabSize); + float const* tProbs = targetProbs + + (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * vocabSize; + if (hasCompactTargetSupport) + { + uint32_t const supportOffset + = (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) + * maxTargetSupportSize; + uint32_t const supportSize + = static_cast(targetSupportLengths[batchOffset + sLastAcceptedLocalIdx]); + acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] = sampleFromIndexedDistribution( + state, tProbs, targetSupportIndices + supportOffset, supportSize, vocabSize); + } + else + { + acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] + = sampleFromDistribution(state, tProbs, vocabSize); + } + sHasTerminalToken = true; } - hasTerminalToken = true; } else { - int32_t const draftProbRow = draftProbIndices[batchOffset + firstChild]; + int32_t const draftProbRow = draftProbIndices[batchOffset + sFirstChild]; float const* tProbs - = targetProbs + (static_cast(bx) * numDraftTokens + lastAcceptedLocalIdx) * vocabSize; + = targetProbs + (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * vocabSize; float const* dProbs = draftProbs + (static_cast(bx) * numDraftProbRows + draftProbRow) * vocabSize; int32_t const* tProbIndices = nullptr; @@ -1329,120 +1352,168 @@ __global__ void verifyDynamicTreeRejectionKernel(int64_t* acceptIndex, int64_t* if (hasCompactTargetSupport) { uint32_t const supportOffset - = (static_cast(bx) * numDraftTokens + lastAcceptedLocalIdx) * maxTargetSupportSize; + = (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * maxTargetSupportSize; tProbIndices = targetSupportIndices + supportOffset; - targetSupportSize = static_cast(targetSupportLengths[batchOffset + lastAcceptedLocalIdx]); + targetSupportSize + = static_cast(targetSupportLengths[batchOffset + sLastAcceptedLocalIdx]); } - // Pass 1: compute diffSum = sum of relu(q - p) over the target support. - float diffSum = 0.0f; if (hasCompactTargetSupport) { - for (uint32_t i = 0; i < targetSupportSize; ++i) + if (tid == 0) { - uint32_t const v = static_cast(tProbIndices[i]); - float diff = tProbs[v] - dProbs[v]; - if (diff > 0.0f) + float diffSum = 0.0f; + for (uint32_t i = 0; i < targetSupportSize; ++i) + { + uint32_t const v = static_cast(tProbIndices[i]); + float const diff = tProbs[v] - dProbs[v]; + if (diff > 0.0f) + { + diffSum += diff; + } + } + + int64_t corrTok = static_cast(vocabSize) - 1; + bool const useDiff = (diffSum > 1e-10f); + + if (useDiff) { - diffSum += diff; + float const r = curand_uniform(&state); + float cumsum = 0.0f; + for (uint32_t i = 0; i < targetSupportSize; ++i) + { + uint32_t const v = static_cast(tProbIndices[i]); + float const diff = tProbs[v] - dProbs[v]; + float const prob = (diff > 0.0f) ? diff / diffSum : 0.0f; + cumsum += prob; + if (r <= cumsum) + { + corrTok = static_cast(v); + break; + } + } } + else + { + corrTok = sampleFromIndexedDistribution( + state, tProbs, tProbIndices, targetSupportSize, vocabSize); + } + acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] = corrTok; + sHasTerminalToken = true; } } else { - for (uint32_t v = 0; v < targetSupportSize; ++v) + float threadDiffSum = 0.0f; + for (uint32_t v = static_cast(tid); v < targetSupportSize; v += BLOCK_SIZE) { - float diff = tProbs[v] - dProbs[v]; + float const diff = tProbs[v] - dProbs[v]; if (diff > 0.0f) { - diffSum += diff; + threadDiffSum += diff; } } - } - // Pass 2: CDF inversion over the normalised residual distribution, - // traversing only the target support. - // Falls back to sampling directly from target when diffSum ~= 0. - // - // Example: - // if F1/F2/F3 are all rejected under parent E, then we sample a correction - // token from relu(q(.|E) - p(.|E)). The sampled token terminates traversal: - // it is emitted as the next output token, and there is no further descent. - int64_t corrTok = static_cast(vocabSize) - 1; // fallback: last vocab token - bool useDiff = (diffSum > 1e-10f); - - if (useDiff) - { - float r = curand_uniform(&state); - float cumsum = 0.0f; - if (hasCompactTargetSupport) + float const diffSum = BlockReduce(tempStorage.reduce).Sum(threadDiffSum); + if (tid == 0) { - for (uint32_t i = 0; i < targetSupportSize; ++i) + sDiffSum = diffSum; + sPrefixBase = 0.0f; + sWinnerIndex = static_cast(targetSupportSize); + sSampledToken = static_cast(vocabSize) - 1; + if (diffSum > 1e-10f) { - uint32_t const v = static_cast(tProbIndices[i]); - float diff = tProbs[v] - dProbs[v]; - float prob = (diff > 0.0f) ? diff / diffSum : 0.0f; - cumsum += prob; - if (r <= cumsum) - { - corrTok = static_cast(v); - break; - } + sTargetMass = curand_uniform(&state) * diffSum; + } + else + { + sSampledToken = sampleFromDistribution(state, tProbs, vocabSize); } } - else + __syncthreads(); + + if (sDiffSum > 1e-10f) { - for (uint32_t v = 0; v < targetSupportSize; ++v) + for (uint32_t tileStart = 0; tileStart < targetSupportSize; tileStart += BLOCK_SIZE) { - float diff = tProbs[v] - dProbs[v]; - float prob = (diff > 0.0f) ? diff / diffSum : 0.0f; - cumsum += prob; - if (r <= cumsum) + float value = 0.0f; + uint32_t const v = tileStart + static_cast(tid); + if (v < targetSupportSize) + { + float const diff = tProbs[v] - dProbs[v]; + value = diff > 0.0f ? diff : 0.0f; + } + + float inclusive = 0.0f; + float tileSum = 0.0f; + BlockScan(tempStorage.scan).InclusiveSum(value, inclusive, tileSum); + float const threshold = sTargetMass; + float const prefixBase = sPrefixBase; + + if (value > 0.0f && prefixBase + inclusive >= threshold) + { + atomicMin(&sWinnerIndex, static_cast(v)); + } + __syncthreads(); + + if (tid == 0) + { + if (sWinnerIndex < static_cast(targetSupportSize)) + { + sSampledToken = static_cast(sWinnerIndex); + } + sPrefixBase += tileSum; + } + __syncthreads(); + + if (sWinnerIndex < static_cast(targetSupportSize)) { - corrTok = static_cast(v); break; } } } - } - else - { - corrTok = hasCompactTargetSupport - ? sampleFromIndexedDistribution(state, tProbs, tProbIndices, targetSupportSize, vocabSize) - : sampleFromDistribution(state, tProbs, vocabSize); - } - acceptToken[bx * numSpeculativeTokens + numAcceptedTokens] = corrTok; - // acceptIndex at correction depth left as 0 (zero-initialized by caller) - hasTerminalToken = true; + if (tid == 0) + { + acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] = sSampledToken; + sHasTerminalToken = true; + } + } } + __syncthreads(); break; } } - if (!hasTerminalToken) + if (!sHasTerminalToken) { // Reached max speculative depth while continuing to accept the draft path. // Emit the final bonus token from the last accepted position. - float const* tProbs - = targetProbs + (static_cast(bx) * numDraftTokens + lastAcceptedLocalIdx) * vocabSize; - if (hasCompactTargetSupport) - { - uint32_t const supportOffset - = (static_cast(bx) * numDraftTokens + lastAcceptedLocalIdx) * maxTargetSupportSize; - uint32_t const supportSize - = static_cast(targetSupportLengths[batchOffset + lastAcceptedLocalIdx]); - acceptToken[bx * numSpeculativeTokens + numAcceptedTokens] = sampleFromIndexedDistribution( - state, tProbs, targetSupportIndices + supportOffset, supportSize, vocabSize); - } - else + if (tid == 0) { - acceptToken[bx * numSpeculativeTokens + numAcceptedTokens] - = sampleFromDistribution(state, tProbs, vocabSize); + float const* tProbs + = targetProbs + (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * vocabSize; + if (hasCompactTargetSupport) + { + uint32_t const supportOffset + = (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * maxTargetSupportSize; + uint32_t const supportSize + = static_cast(targetSupportLengths[batchOffset + sLastAcceptedLocalIdx]); + acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] = sampleFromIndexedDistribution( + state, tProbs, targetSupportIndices + supportOffset, supportSize, vocabSize); + } + else + { + acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] + = sampleFromDistribution(state, tProbs, vocabSize); + } } } - acceptTokenNum[bx] = numAcceptedTokens; + if (tid == 0) + { + acceptTokenNum[bx] = sNumAcceptedTokens; + } } void invokeVerifyDynamicTreeRejection(int64_t* acceptIndex, int64_t* acceptTokenNum, int64_t* acceptToken, @@ -1452,13 +1523,14 @@ void invokeVerifyDynamicTreeRejection(int64_t* acceptIndex, int64_t* acceptToken SizeType32 maxTargetSupportSize, SizeType32 numDraftTokens, SizeType32 numSpecStep, SizeType32 vocabSize, int64_t const* seed, int64_t const* offset, cudaStream_t stream) { + constexpr int32_t kVerifyDynamicTreeRejectionBlockSize = 128; dim3 grid(batchSize); - dim3 block(1); + dim3 block(kVerifyDynamicTreeRejectionBlockSize); - verifyDynamicTreeRejectionKernel<<>>(acceptIndex, acceptTokenNum, acceptToken, candidates, - draftProbs, targetProbs, targetSupportIndices, targetSupportLengths, draftProbIndices, retrieveNextToken, - retrieveNextSibling, batchSize, numDraftProbRows, maxTargetSupportSize, numSpecStep, numDraftTokens, vocabSize, - seed, offset); + verifyDynamicTreeRejectionKernel<<>>(acceptIndex, + acceptTokenNum, acceptToken, candidates, draftProbs, targetProbs, targetSupportIndices, targetSupportLengths, + draftProbIndices, retrieveNextToken, retrieveNextSibling, batchSize, numDraftProbRows, maxTargetSupportSize, + numSpecStep, numDraftTokens, vocabSize, seed, offset); sync_check_cuda_error(stream); } diff --git a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py index 09995b04fd99..32fb46d822b4 100644 --- a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py @@ -734,152 +734,208 @@ def _sample_and_accept_dynamic_tree( self, logits, attn_metadata, spec_metadata, batch_size, num_contexts, num_gens ): """Dynamic tree verification using CUDA kernel.""" - if num_gens > self._max_batch_size: - raise RuntimeError( - f"Dynamic tree batch size {num_gens} exceeds configured " - f"max_batch_size {self._max_batch_size}" - ) - N = self.tokens_per_gen_step - max_path_len = self._max_path_len - - # Reset output buffers - self._accepted_tokens_buf[:batch_size].zero_() - accepted_tokens = self._accepted_tokens_buf[:batch_size, :max_path_len] - self._num_accepted_tokens_buf[:batch_size].fill_(1) - num_accepted_tokens = self._num_accepted_tokens_buf[:batch_size] - self._accepted_draft_indices_tensor[:batch_size].fill_(-1) - - num_flat_tokens = logits.shape[0] - torch.argmax(logits, dim=-1, out=self._target_tokens_buf[:num_flat_tokens]) - target_tokens = self._target_tokens_buf[:num_flat_tokens] - - # Context requests: accept sampled token - accepted_tokens[:num_contexts, 0] = target_tokens[:num_contexts].to(torch.int32) - - # Generation requests: tree verification - if num_gens > 0: - spec_tree_manager = self.spec_tree_manager - - target_predict = self._target_predict_buf[:num_gens] - target_predict[:] = target_tokens[num_contexts:].reshape(num_gens, N) - - if spec_tree_manager is None: - # CUDA graph warmup: accept only the first token per request - num_accepted_tokens[num_contexts:batch_size] = 1 - accepted_tokens[num_contexts:batch_size, 0] = target_predict[:, 0].to(torch.int32) - self._accepted_draft_indices_tensor[num_contexts:batch_size] = -1 - return accepted_tokens, num_accepted_tokens - - candidates = self._candidates_buf[:num_gens] - candidates[:, 1:] = spec_metadata.draft_tokens.reshape(num_gens, N - 1).to(torch.int64) - candidates[:, 0] = target_predict[:, 0] - - # Slots for gen rows: real py_seq_slot vs dummy; dummy -> slot_has_tree False. - gen_slot_ids = spec_tree_manager._all_slot_ids_buf[ - num_contexts : num_contexts + num_gens - ] - tree_valid = spec_tree_manager.slot_has_tree[gen_slot_ids] - - if self._can_use_rejection_sampling(spec_metadata): - vocab_size = logits.shape[-1] - num_ctx_tokens = logits.shape[0] - num_gens * N - device = logits.device - - draft_logits_tree = self._get_unique_draft_logits(num_gens) - draft_prob_indices = self._build_draft_prob_indices(num_gens) - target_logits_tree = logits[num_ctx_tokens:].reshape(-1, vocab_size) - gen_slice = slice(num_contexts, num_contexts + num_gens) - skip_top_k = getattr(spec_metadata, "skip_top_k", False) - skip_top_p = getattr(spec_metadata, "skip_top_p", False) - skip_temperature = getattr(spec_metadata, "skip_temperature", False) - - if spec_metadata.temperatures is None: - temps = torch.ones(num_gens, dtype=torch.float32, device=device) - skip_temperature = True - else: - temps = spec_metadata.temperatures[gen_slice] - - top_ks = None - if not skip_top_k and spec_metadata.top_ks is not None: - top_ks = spec_metadata.top_ks[gen_slice] - - top_ps = None - if not skip_top_p and spec_metadata.top_ps is not None: - top_ps = spec_metadata.top_ps[gen_slice] - - # Lazily initialize seed/offset tensors on correct device - 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) - # Use in-place operations for CUDA graph compatibility - self.seed.add_(1).remainder_(2**31) - - _, accept_index, accept_token_num, accept_token = ( - self.tree_ops_converter.verify_dynamic_tree_rejection_from_logits_out( - candidates, - draft_logits_tree, - target_logits_tree, - draft_prob_indices, - spec_tree_manager.retrieve_next_token[:num_gens], - spec_tree_manager.retrieve_next_sibling[:num_gens], - temps, - top_ks, - top_ps, - skip_temperature, - num_gens, - self._max_path_len, - seed=self.seed, - offset=self.offset, - d2t=self._d2t, - ) + torch.cuda.nvtx.range_push("eagle3_dynamic_tree::_sample_and_accept_dynamic_tree") + try: + if num_gens > self._max_batch_size: + raise RuntimeError( + f"Dynamic tree batch size {num_gens} exceeds configured " + f"max_batch_size {self._max_batch_size}" ) + N = self.tokens_per_gen_step + max_path_len = self._max_path_len - self._finalize_dynamic_tree_verify_outputs( - accept_index=accept_index, - accept_token_num=accept_token_num, - accept_token=accept_token, - accepted_tokens=accepted_tokens, - num_accepted_tokens=num_accepted_tokens, - num_contexts=num_contexts, - batch_size=batch_size, - num_gens=num_gens, - max_path_len=max_path_len, + torch.cuda.nvtx.range_push( + "eagle3_dynamic_tree::_sample_and_accept_dynamic_tree::prepare_outputs" + ) + try: + # Reset output buffers + self._accepted_tokens_buf[:batch_size].zero_() + accepted_tokens = self._accepted_tokens_buf[:batch_size, :max_path_len] + self._num_accepted_tokens_buf[:batch_size].fill_(1) + num_accepted_tokens = self._num_accepted_tokens_buf[:batch_size] + self._accepted_draft_indices_tensor[:batch_size].fill_(-1) + finally: + torch.cuda.nvtx.range_pop() + + torch.cuda.nvtx.range_push( + "eagle3_dynamic_tree::_sample_and_accept_dynamic_tree::compute_target_tokens" + ) + try: + num_flat_tokens = logits.shape[0] + torch.argmax(logits, dim=-1, out=self._target_tokens_buf[:num_flat_tokens]) + target_tokens = self._target_tokens_buf[:num_flat_tokens] + + # Context requests: accept sampled token + accepted_tokens[:num_contexts, 0] = target_tokens[:num_contexts].to(torch.int32) + finally: + torch.cuda.nvtx.range_pop() + + # Generation requests: tree verification + if num_gens > 0: + torch.cuda.nvtx.range_push( + "eagle3_dynamic_tree::_sample_and_accept_dynamic_tree::generation_verify" ) + try: + spec_tree_manager = self.spec_tree_manager + + target_predict = self._target_predict_buf[:num_gens] + target_predict[:] = target_tokens[num_contexts:].reshape(num_gens, N) + + if spec_tree_manager is None: + # CUDA graph warmup: accept only the first token per request + num_accepted_tokens[num_contexts:batch_size] = 1 + accepted_tokens[num_contexts:batch_size, 0] = target_predict[:, 0].to( + torch.int32 + ) + self._accepted_draft_indices_tensor[num_contexts:batch_size] = -1 + return accepted_tokens, num_accepted_tokens + + candidates = self._candidates_buf[:num_gens] + candidates[:, 1:] = spec_metadata.draft_tokens.reshape(num_gens, N - 1).to( + torch.int64 + ) + candidates[:, 0] = target_predict[:, 0] + + # Slots for gen rows: real py_seq_slot vs dummy; dummy -> slot_has_tree False. + gen_slot_ids = spec_tree_manager._all_slot_ids_buf[ + num_contexts : num_contexts + num_gens + ] + tree_valid = spec_tree_manager.slot_has_tree[gen_slot_ids] + + if self._can_use_rejection_sampling(spec_metadata): + torch.cuda.nvtx.range_push( + "eagle3_dynamic_tree::_sample_and_accept_dynamic_tree::rejection_verify" + ) + try: + vocab_size = logits.shape[-1] + num_ctx_tokens = logits.shape[0] - num_gens * N + device = logits.device + + draft_logits_tree = self._get_unique_draft_logits(num_gens) + draft_prob_indices = self._build_draft_prob_indices(num_gens) + target_logits_tree = logits[num_ctx_tokens:].reshape(-1, vocab_size) + gen_slice = slice(num_contexts, num_contexts + num_gens) + skip_top_k = getattr(spec_metadata, "skip_top_k", False) + skip_top_p = getattr(spec_metadata, "skip_top_p", False) + skip_temperature = getattr(spec_metadata, "skip_temperature", False) + + if spec_metadata.temperatures is None: + temps = torch.ones(num_gens, dtype=torch.float32, device=device) + skip_temperature = True + else: + temps = spec_metadata.temperatures[gen_slice] + + top_ks = None + if not skip_top_k and spec_metadata.top_ks is not None: + top_ks = spec_metadata.top_ks[gen_slice] + + top_ps = None + if not skip_top_p and spec_metadata.top_ps is not None: + top_ps = spec_metadata.top_ps[gen_slice] + + # Lazily initialize seed/offset tensors on correct device + 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) + # Use in-place operations for CUDA graph compatibility + self.seed.add_(1).remainder_(2**31) + + _, accept_index, accept_token_num, accept_token = ( + self.tree_ops_converter.verify_dynamic_tree_rejection_from_logits_out( + candidates, + draft_logits_tree, + target_logits_tree, + draft_prob_indices, + spec_tree_manager.retrieve_next_token[:num_gens], + spec_tree_manager.retrieve_next_sibling[:num_gens], + temps, + top_ks, + top_ps, + skip_temperature, + num_gens, + self._max_path_len, + seed=self.seed, + offset=self.offset, + d2t=self._d2t, + ) + ) + finally: + torch.cuda.nvtx.range_pop() + + torch.cuda.nvtx.range_push( + "eagle3_dynamic_tree::_sample_and_accept_dynamic_tree::finalize_rejection" + ) + try: + self._finalize_dynamic_tree_verify_outputs( + accept_index=accept_index, + accept_token_num=accept_token_num, + accept_token=accept_token, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_contexts=num_contexts, + batch_size=batch_size, + num_gens=num_gens, + max_path_len=max_path_len, + ) + num_accepted_tokens = self._apply_force_accepted_tokens( + num_accepted_tokens, num_contexts, self.max_draft_len + ) + return accepted_tokens, num_accepted_tokens + finally: + torch.cuda.nvtx.range_pop() + + torch.cuda.nvtx.range_push( + "eagle3_dynamic_tree::_sample_and_accept_dynamic_tree::greedy_verify" + ) + try: + _, accept_index, accept_token_num, accept_token = ( + self.tree_ops_converter.verify_dynamic_tree_greedy_out( + candidates, + spec_tree_manager.retrieve_index[:num_gens], + spec_tree_manager.retrieve_next_token[:num_gens], + spec_tree_manager.retrieve_next_sibling[:num_gens], + target_predict, + num_gens, + self._max_path_len, + tree_valid=tree_valid, + ) + ) + finally: + torch.cuda.nvtx.range_pop() + + torch.cuda.nvtx.range_push( + "eagle3_dynamic_tree::_sample_and_accept_dynamic_tree::finalize_greedy" + ) + try: + self._finalize_dynamic_tree_verify_outputs( + accept_index=accept_index, + accept_token_num=accept_token_num, + accept_token=accept_token, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_contexts=num_contexts, + batch_size=batch_size, + num_gens=num_gens, + max_path_len=max_path_len, + ) + finally: + torch.cuda.nvtx.range_pop() + finally: + torch.cuda.nvtx.range_pop() + + torch.cuda.nvtx.range_push( + "eagle3_dynamic_tree::_sample_and_accept_dynamic_tree::apply_force_accepted_tokens" + ) + try: num_accepted_tokens = self._apply_force_accepted_tokens( num_accepted_tokens, num_contexts, self.max_draft_len ) - return accepted_tokens, num_accepted_tokens - - _, accept_index, accept_token_num, accept_token = ( - self.tree_ops_converter.verify_dynamic_tree_greedy_out( - candidates, - spec_tree_manager.retrieve_index[:num_gens], - spec_tree_manager.retrieve_next_token[:num_gens], - spec_tree_manager.retrieve_next_sibling[:num_gens], - target_predict, - num_gens, - self._max_path_len, - tree_valid=tree_valid, - ) - ) + finally: + torch.cuda.nvtx.range_pop() - self._finalize_dynamic_tree_verify_outputs( - accept_index=accept_index, - accept_token_num=accept_token_num, - accept_token=accept_token, - accepted_tokens=accepted_tokens, - num_accepted_tokens=num_accepted_tokens, - num_contexts=num_contexts, - batch_size=batch_size, - num_gens=num_gens, - max_path_len=max_path_len, - ) - - num_accepted_tokens = self._apply_force_accepted_tokens( - num_accepted_tokens, num_contexts, self.max_draft_len - ) - - return accepted_tokens, num_accepted_tokens + return accepted_tokens, num_accepted_tokens + finally: + torch.cuda.nvtx.range_pop() def _can_use_rejection_sampling(self, spec_metadata) -> bool: """Check if rejection sampling can be used for dynamic tree verification. From 66ef38560a229c3930526aa5264501621496d9d3 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 22 Apr 2026 01:55:07 -0700 Subject: [PATCH 04/17] [None][test] add EAGLE3 rejection sampling coverage for dynamic tree Signed-off-by: ZhaoyangWang --- .../defs/accuracy/test_llm_api_pytorch.py | 48 ++++++++++++++ .../test_lists/test-db/l0_dgx_b200.yml | 5 ++ .../_torch/speculative/test_eagle3.py | 64 +++++++++++++++++++ 3 files changed, 117 insertions(+) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 51504b530938..4a2bb5a8d37b 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -187,6 +187,54 @@ def test_dummy_load_format(self): task = MMLU(self.MODEL_NAME) task.evaluate(llm, is_integration_test=True) + @pytest.mark.skip_less_device(4) + @pytest.mark.parametrize("use_dynamic_tree", [False, True], + ids=["no_dynamic_tree", "dynamic_tree"]) + def test_eagle3_rejection_dynamic_tree_smoke(self, use_dynamic_tree, + mocker): + mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", 128) + + eagle_model_dir = f"{llm_models_root()}/EAGLE3-LLaMA3.1-Instruct-8B" + spec_config_kwargs = dict( + max_draft_len=4, + speculative_model=eagle_model_dir, + eagle3_one_model=True, + allow_advanced_sampling=True, + use_rejection_sampling=True, + ) + if use_dynamic_tree: + spec_config_kwargs.update( + use_dynamic_tree=True, + dynamic_tree_max_topK=4, + max_total_draft_tokens=16, + max_batch_size=4, + ) + + llm = LLM( + self.MODEL_PATH, + tensor_parallel_size=4, + pipeline_parallel_size=1, + attn_backend="TRTLLM", + disable_overlap_scheduler=True, + cuda_graph_config=None, + kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.4, + dtype="auto"), + max_seq_len=4096, + max_batch_size=4, + speculative_config=Eagle3DecodingConfig(**spec_config_kwargs), + ) + + with llm: + task = GSM8K(self.MODEL_NAME) + sampling_params = SamplingParams(temperature=1.0, + top_p=1.0, + max_tokens=128, + truncate_prompt_tokens=2048) + task.evaluate(llm, + sampling_params=sampling_params, + extra_evaluator_kwargs=dict(apply_chat_template=True), + is_integration_test=True) + @pytest.mark.skip_less_device_memory(32000) @parametrize_with_ids("torch_compile", [False, True]) @parametrize_with_ids("attn_backend", ["TRTLLM", "FLASHINFER"]) diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index ab3f8c18d5d8..dea6041d46e0 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -291,6 +291,11 @@ l0_dgx_b200: - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-two_model-overlap_scheduler] - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-two_model-no_overlap_scheduler] - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-two_model-no_overlap_scheduler] + - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3_rejection_dynamic_tree_smoke[no_dynamic_tree] + - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3_rejection_dynamic_tree_smoke[dynamic_tree] + - disaggregated/test_disaggregated.py::test_disaggregated_benchmark_on_diff_backends[llama-v3-8b-hf] + - disaggregated/test_disaggregated.py::test_disaggregated_benchmark_on_diff_backends[DeepSeek-V3-Lite-bf16] + - disaggregated/test_disaggregated.py::test_disaggregated_benchmark_on_diff_backends[llama-3.1-8b-instruct-hf-fp8] - unittest/_torch/multi_gpu_modeling -k "deepseek" - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[pp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTEDSL-mtp_nextn=2-ep4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index d7d8f0a6e813..2ff2f6f5e868 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -825,6 +825,70 @@ def test_eagle3_cdl_sampling(disable_overlap_scheduler: bool): llm_spec.shutdown() +@pytest.mark.parametrize("use_dynamic_tree", [False, True], + ids=["no_dynamic_tree", "dynamic_tree"]) +@pytest.mark.parametrize("use_cuda_graph", [False, True]) +@pytest.mark.high_cuda_memory +@skip_blackwell +@with_mocked_hf_download_for_single_gpu +def test_llama_eagle3_rejection_sampling_modes(use_dynamic_tree: bool, + use_cuda_graph: bool): + """Test one-model rejection sampling with and without dynamic tree.""" + 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() + target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" + eagle_model = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" + + max_batch_size = 1 + max_draft_len = 6 + dynamic_tree_max_top_k = 10 + max_total_draft_tokens = 60 + kv_cache_config = KvCacheConfig(enable_block_reuse=False, max_tokens=8192) + cuda_graph_config = CudaGraphConfig( + batch_sizes=[1]) if use_cuda_graph else None + + llm_common_config = dict( + model=target_model_dir, + attn_backend="TRTLLM", + disable_overlap_scheduler=True, + cuda_graph_config=cuda_graph_config, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + max_seq_len=8192, + ) + + spec_config_kwargs = dict( + max_draft_len=max_draft_len, + speculative_model=eagle_model, + eagle3_one_model=True, + allow_advanced_sampling=True, + use_rejection_sampling=True, + ) + if use_dynamic_tree: + spec_config_kwargs.update( + use_dynamic_tree=True, + dynamic_tree_max_topK=dynamic_tree_max_top_k, + max_total_draft_tokens=max_total_draft_tokens, + max_batch_size=max_batch_size, + ) + + llm_spec = LLM(**llm_common_config, + speculative_config=Eagle3DecodingConfig( + **spec_config_kwargs)) + + prompts = ["The president of the United States is"] + sampling_params = SamplingParams(max_tokens=20, temperature=1.0, top_p=1.0) + + results = llm_spec.generate(prompts, sampling_params) + llm_spec.shutdown() + + assert len(results) == len(prompts) + assert len(results[0].outputs[0].token_ids) > 0 + + @pytest.mark.parametrize("use_cuda_graph", [True, False]) def test_eagle3_lora(use_cuda_graph: bool): """Test LoRA with 3 requests and max_batch_size=4. From 16f026718eaba46002916d68bc00d12f80acf6e9 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 22 Apr 2026 05:50:23 -0700 Subject: [PATCH 05/17] [None][fix] address rejection sampling review fixes Signed-off-by: ZhaoyangWang --- .../speculativeDecoding/dynamicTreeKernels.cu | 50 +++++++++++++------ .../speculativeDecoding/dynamicTreeKernels.h | 7 ++- cpp/tensorrt_llm/thop/dynamicTreeOp.cpp | 28 +++++++++++ .../_torch/pyexecutor/model_engine.py | 7 +++ .../_torch/speculative/dynamic_tree_ops.py | 6 ++- .../_torch/speculative/eagle3_dynamic_tree.py | 12 ++--- tensorrt_llm/_torch/speculative/interface.py | 41 +++++++++------ tensorrt_llm/_torch/speculative/mtp.py | 3 +- tensorrt_llm/llmapi/llm_args.py | 24 +++++++++ 9 files changed, 139 insertions(+), 39 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu index e2a4dd27a225..54fb366f54b7 100644 --- a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu +++ b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu @@ -351,8 +351,17 @@ torch::Tensor applyTopKTopPForProbOp(torch::Tensor logits, torch::optionalto(torch::kLong); + effectiveTopK + = torch::where(topKLong > 0, topKLong, torch::full_like(topKLong, vocabSize)).clamp_max(vocabSize); + hasDisabledTopKRows = topKLong.le(0).any().item(); + } - if (hasTopK && kMax > 0 && kMax < vocabSize) + if (hasTopK && !hasDisabledTopKRows && kMax > 0 && kMax < vocabSize) { // Fast topk path ───────────────────────────────────────────────────────────── // topKValues/topKIdx: [nRows, kMax], values in descending order @@ -360,9 +369,9 @@ torch::Tensor applyTopKTopPForProbOp(torch::Tensor logits, torch::optionalto(torch::kInt64).unsqueeze(1); // [nRows, 1] - auto validTopK = kArange < kVals; // [nRows, kMax] + .unsqueeze(0); // [1, kMax] + auto kVals = effectiveTopK.to(torch::kInt64).unsqueeze(1); // [nRows, 1] + auto validTopK = kArange < kVals; // [nRows, kMax] // Start with everything masked; scatter will unmark the kept positions. auto mask = torch::ones( @@ -401,7 +410,7 @@ torch::Tensor applyTopKTopPForProbOp(torch::Tensor logits, torch::optionalto(torch::kLong); + auto topKMask = logitsSort.size(1) - effectiveTopK; topKMask = topKMask.clamp_min(0); auto topKThreshold = logitsSort.gather(1, topKMask.unsqueeze(1)); auto mask = logitsSort < topKThreshold; @@ -444,8 +453,10 @@ torch::Tensor computeProbsFromLogits(torch::Tensor const& logits, torch::Tensor TORCH_CHECK(topP->size(0) == logits.size(0), "top_p and logits size mismatch"); } + auto const isGreedy = temperatures <= kGreedyTempThreshold; + auto const safeTemperatures = torch::where(isGreedy, torch::ones_like(temperatures), temperatures); auto scaledLogits - = (skipTemperature ? logits : logits.div(temperatures.unsqueeze(1))).contiguous().to(torch::kFloat32); + = (skipTemperature ? logits : logits.div(safeTemperatures.unsqueeze(1))).contiguous().to(torch::kFloat32); bool const hasTopK = topK.has_value() && topK->defined(); int64_t const vocabSize = scaledLogits.size(1); @@ -470,7 +481,6 @@ torch::Tensor computeProbsFromLogits(torch::Tensor const& logits, torch::Tensor auto probs = computeSoftmaxForProbOp(maskedLogits); - auto isGreedy = temperatures <= kGreedyTempThreshold; auto argmaxIds = maskedLogits.argmax(/*dim=*/-1, /*keepdim=*/true); auto oneHot = torch::zeros_like(probs).scatter_(1, argmaxIds, 1.0); return torch::where(isGreedy.unsqueeze(1), oneHot, probs); @@ -592,11 +602,24 @@ std::tuple computeTargetProbsForDyn bool const hasTopK = targetTopK.has_value() && targetTopK->defined(); bool const hasTopP = targetTopP.has_value() && targetTopP->defined(); bool const hasFiltering = hasTopK || hasTopP; + torch::Tensor effectiveTargetTopK; + bool hasDisabledTopKRows = false; - auto scaledTargetLogits = (skipTemperature ? targetLogits : targetLogits.div(targetTemps.unsqueeze(1))) + auto const isGreedy = targetTemps <= kGreedyTempThreshold; + auto const safeTargetTemps = torch::where(isGreedy, torch::ones_like(targetTemps), targetTemps); + auto scaledTargetLogits = (skipTemperature ? targetLogits : targetLogits.div(safeTargetTemps.unsqueeze(1))) .contiguous() .to(torch::kFloat32); + if (hasTopK) + { + auto targetTopKLong = targetTopK->to(torch::kLong); + effectiveTargetTopK + = torch::where(targetTopKLong > 0, targetTopKLong, torch::full_like(targetTopKLong, targetVocabSize)) + .clamp_max(targetVocabSize); + hasDisabledTopKRows = targetTopKLong.le(0).any().item(); + } + torch::Tensor maskedTargetLogits; torch::Tensor targetSupportIndices; torch::Tensor targetSupportLengths; @@ -610,13 +633,13 @@ std::tuple computeTargetProbsForDyn targetSupportLengths = torch::empty({0}, torch::TensorOptions().dtype(torch::kInt32).device(targetLogits.device())); } - else if (hasTopK && kMax > 0 && static_cast(kMax) < targetVocabSize) + else if (hasTopK && !hasDisabledTopKRows && kMax > 0 && static_cast(kMax) < targetVocabSize) { // Fast two-stage CUDA path for masked logits. maskedTargetLogits = torch::empty_like(scaledTargetLogits); auto stream = at::cuda::getCurrentCUDAStream(scaledTargetLogits.device().index()); invokeTopKTopPMaskingForProbs(scaledTargetLogits.data_ptr(), maskedTargetLogits.data_ptr(), - targetTopK->to(torch::kInt32).contiguous().data_ptr(), + effectiveTargetTopK.to(torch::kInt32).contiguous().data_ptr(), hasTopP ? targetTopP->to(torch::kFloat32).contiguous().data_ptr() : nullptr, kMax, static_cast(nRows), static_cast(targetVocabSize), stream); @@ -639,7 +662,7 @@ std::tuple computeTargetProbsForDyn if (hasTopK) { - auto topKMask = logitsSort.size(1) - targetTopK->to(torch::kLong); + auto topKMask = logitsSort.size(1) - effectiveTargetTopK; topKMask = topKMask.clamp_min(0); auto topKThreshold = logitsSort.gather(1, topKMask.unsqueeze(1)); auto mask = logitsSort < topKThreshold; @@ -663,9 +686,9 @@ std::tuple computeTargetProbsForDyn auto supportLengths1D = supportLengthsLong.to(torch::kInt32); int64_t maxSupportSize = targetVocabSize; - if (hasTopK && targetTopK->defined()) + if (hasTopK && effectiveTargetTopK.defined()) { - maxSupportSize = std::min(targetVocabSize, targetTopK->max().item()); + maxSupportSize = std::min(targetVocabSize, effectiveTargetTopK.max().item()); } auto compactPositions @@ -685,7 +708,6 @@ std::tuple computeTargetProbsForDyn auto targetProbs = computeSoftmaxForProbOp(maskedTargetLogits); - auto isGreedy = targetTemps <= kGreedyTempThreshold; auto argmaxIds = maskedTargetLogits.argmax(/*dim=*/-1, /*keepdim=*/true); auto oneHot = torch::zeros_like(targetProbs).scatter_(1, argmaxIds, 1.0); targetProbs = torch::where(isGreedy.unsqueeze(1), oneHot, targetProbs); diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h index d2fe371d3823..78eecfdcbe96 100644 --- a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h +++ b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h @@ -114,15 +114,20 @@ void invokeVerifyDynamicTreeGreedy(int64_t* predicts, int64_t* acceptIndex, int6 //! \param acceptTokenNum output [batchSize] int64 — # accepted draft tokens (excl. root). //! \param acceptToken output [batchSize, numSpecStep] int64 — accepted/correction token ids. //! \param candidates [batchSize, numDraftTokens] int64; col 0 = root (target sample). -//! \param draftProbs [batchSize, numDraftTokens-1, vocabSize] float32; index 0 = tree pos 1. +//! \param draftProbs [batchSize, numDraftProbRows, vocabSize] float32. +//! Unique draft probability rows per request; tree positions map into this +//! tensor via draftProbIndices. //! \param targetProbs [batchSize, numDraftTokens, vocabSize] float32; index 0 = root. //! \param targetSupportIndices [batchSize, numDraftTokens, maxTargetSupportSize] int32; compact token ids that //! survive top-k/top-p filtering for each row, padded with -1. May be empty when no filtering is active. //! \param targetSupportLengths [batchSize, numDraftTokens] int32; valid support length per row. May be empty when //! no filtering is active. +//! \param draftProbIndices [batchSize, numDraftTokens] int32; maps each tree position to the +//! corresponding row in draftProbs. Root is unused. //! \param retrieveNextToken [batchSize, numDraftTokens] int32 first-child pointer, -1=none. //! \param retrieveNextSibling [batchSize, numDraftTokens] int32 next-sibling pointer, -1=none. //! \param batchSize runtime::SizeType32. +//! \param numDraftProbRows runtime::SizeType32. Number of unique draft-prob rows per request. //! \param maxTargetSupportSize runtime::SizeType32. Third dim of targetSupportIndices. Can be zero. //! \param numDraftTokens runtime::SizeType32. Total tree nodes per request (including root). //! \param numSpecStep runtime::SizeType32. Second dim of acceptIndex/acceptToken. diff --git a/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp b/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp index 550bb8cde230..e795a04c96a1 100644 --- a/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp +++ b/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp @@ -56,6 +56,7 @@ void build_draft_prob_indices_out_op( TORCH_CHECK(topkScoreIndices.size(1) == numDraftTokens, "topkScoreIndices size mismatch"); TORCH_CHECK(draftProbIndices.size(0) == topkScoreIndices.size(0), "Batch size mismatch"); TORCH_CHECK(draftProbIndices.size(1) == numDraftTokens + 1, "draftProbIndices size mismatch"); + TORCH_CHECK(topK > 0, "topK must be positive"); TORCH_CHECK(numDraftTokens + 1 <= 1024, "numDraftTokens + 1 exceeds CUDA block size limit of 1024"); auto stream = at::cuda::getCurrentCUDAStream(topkScoreIndices.device().index()); @@ -233,6 +234,15 @@ void verify_dynamic_tree_rejection_out_op(th::Tensor& candidates, th::Tensor& dr TORCH_CHECK(targetSupportIndices.scalar_type() == torch::kInt32, "targetSupportIndices must be int32 tensor"); TORCH_CHECK(targetSupportLengths.scalar_type() == torch::kInt32, "targetSupportLengths must be int32 tensor"); TORCH_CHECK(draftProbIndices.scalar_type() == torch::kInt32, "draftProbIndices must be int32 tensor"); + TORCH_CHECK(candidates.is_cuda(), "candidates must be a CUDA tensor"); + TORCH_CHECK(draftProbs.is_cuda(), "draftProbs must be a CUDA tensor"); + TORCH_CHECK(targetProbs.is_cuda(), "targetProbs must be a CUDA tensor"); + TORCH_CHECK(draftProbIndices.is_cuda(), "draftProbIndices must be a CUDA tensor"); + TORCH_CHECK(retrieveNextToken.is_cuda(), "retrieveNextToken must be a CUDA tensor"); + TORCH_CHECK(retrieveNextSibling.is_cuda(), "retrieveNextSibling must be a CUDA tensor"); + TORCH_CHECK(acceptIndex.is_cuda(), "acceptIndex must be a CUDA tensor"); + TORCH_CHECK(acceptTokenNum.is_cuda(), "acceptTokenNum must be a CUDA tensor"); + TORCH_CHECK(acceptToken.is_cuda(), "acceptToken must be a CUDA tensor"); int64_t batchSize = candidates.size(0); int64_t numDraftProbRows = draftProbs.size(1); @@ -249,12 +259,18 @@ void verify_dynamic_tree_rejection_out_op(th::Tensor& candidates, th::Tensor& dr TORCH_CHECK(targetSupportIndices.dim() == 3, "targetSupportIndices must be 3D when non-empty"); TORCH_CHECK(targetSupportIndices.size(0) == batchSize, "targetSupportIndices batch size mismatch"); TORCH_CHECK(targetSupportIndices.size(1) == numDraftTokens, "targetSupportIndices numDraftTokens mismatch"); + TORCH_CHECK(targetSupportIndices.is_cuda(), "targetSupportIndices must be a CUDA tensor when non-empty"); + TORCH_CHECK(targetSupportIndices.device() == candidates.device(), + "targetSupportIndices must be on the same device as candidates"); } if (targetSupportLengths.numel() > 0) { TORCH_CHECK(targetSupportLengths.dim() == 2, "targetSupportLengths must be 2D when non-empty"); TORCH_CHECK(targetSupportLengths.size(0) == batchSize, "targetSupportLengths batch size mismatch"); TORCH_CHECK(targetSupportLengths.size(1) == numDraftTokens, "targetSupportLengths numDraftTokens mismatch"); + TORCH_CHECK(targetSupportLengths.is_cuda(), "targetSupportLengths must be a CUDA tensor when non-empty"); + TORCH_CHECK(targetSupportLengths.device() == candidates.device(), + "targetSupportLengths must be on the same device as candidates"); } TORCH_CHECK((targetSupportIndices.numel() == 0) == (targetSupportLengths.numel() == 0), "targetSupportIndices and targetSupportLengths must both be empty or both be non-empty"); @@ -264,12 +280,24 @@ void verify_dynamic_tree_rejection_out_op(th::Tensor& candidates, th::Tensor& dr TORCH_CHECK(retrieveNextToken.size(1) == numDraftTokens, "retrieveNextToken size mismatch"); TORCH_CHECK(retrieveNextSibling.size(0) == batchSize, "retrieveNextSibling batch size mismatch"); TORCH_CHECK(retrieveNextSibling.size(1) == numDraftTokens, "retrieveNextSibling size mismatch"); + TORCH_CHECK(draftProbs.device() == candidates.device(), "draftProbs must be on the same device as candidates"); + TORCH_CHECK(targetProbs.device() == candidates.device(), "targetProbs must be on the same device as candidates"); + TORCH_CHECK( + draftProbIndices.device() == candidates.device(), "draftProbIndices must be on the same device as candidates"); + TORCH_CHECK(retrieveNextToken.device() == candidates.device(), + "retrieveNextToken must be on the same device as candidates"); + TORCH_CHECK(retrieveNextSibling.device() == candidates.device(), + "retrieveNextSibling must be on the same device as candidates"); TORCH_CHECK(acceptIndex.scalar_type() == torch::kInt64, "acceptIndex must be int64 tensor"); TORCH_CHECK(acceptTokenNum.scalar_type() == torch::kInt64, "acceptTokenNum must be int64 tensor"); TORCH_CHECK(acceptToken.scalar_type() == torch::kInt64, "acceptToken must be int64 tensor"); TORCH_CHECK(acceptIndex.size(0) >= batchSize && acceptIndex.size(1) >= numSpecStep, "acceptIndex buffer too small"); TORCH_CHECK(acceptTokenNum.size(0) >= batchSize, "acceptTokenNum buffer too small"); TORCH_CHECK(acceptToken.size(0) >= batchSize && acceptToken.size(1) >= numSpecStep, "acceptToken buffer too small"); + TORCH_CHECK(acceptIndex.device() == candidates.device(), "acceptIndex must be on the same device as candidates"); + TORCH_CHECK( + acceptTokenNum.device() == candidates.device(), "acceptTokenNum must be on the same device as candidates"); + TORCH_CHECK(acceptToken.device() == candidates.device(), "acceptToken must be on the same device as candidates"); TORCH_CHECK(seed.scalar_type() == torch::kInt64, "seed must be int64 tensor"); TORCH_CHECK(offset.scalar_type() == torch::kInt64, "offset must be int64 tensor"); TORCH_CHECK(seed.numel() >= 1, "seed tensor must have at least one element"); diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index f7b84d31ccca..129415dc27d9 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -457,6 +457,13 @@ def __init__( cuda_graph_batch_sizes, self.batch_size, self.max_num_tokens, self.original_max_total_draft_tokens, self._cuda_graph_padding_enabled) if cuda_graph_batch_sizes else [] + dynamic_tree_warmup_max_batch_size = self._get_dynamic_tree_warmup_max_batch_size( + ) + if dynamic_tree_warmup_max_batch_size is not None: + self._cuda_graph_batch_sizes = [ + bs for bs in self._cuda_graph_batch_sizes + if bs <= dynamic_tree_warmup_max_batch_size + ] self._max_cuda_graph_batch_size = (self._cuda_graph_batch_sizes[-1] if self._cuda_graph_batch_sizes else 0) diff --git a/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py b/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py index 849cd87b3e3a..f48c7856522f 100644 --- a/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py +++ b/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py @@ -267,7 +267,11 @@ def verify_dynamic_tree_rejection_from_logits_out( num_draft_prob_rows = draft_logits_tree.shape[0] // num_gens target_vocab_size = target_logits_tree.shape[-1] - top_k_max = int(top_k.max().item()) if top_k is not None else 0 + if top_k is None: + top_k_max = 0 + else: + enabled_top_k = top_k[(top_k > 0) & (top_k < target_vocab_size)] + top_k_max = int(enabled_top_k.max().item()) if enabled_top_k.numel() > 0 else 0 try: torch.cuda.nvtx.range_push("verify_dynamic_tree_rejection_from_logits_out") diff --git a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py index 32fb46d822b4..fa83f28f4440 100644 --- a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py @@ -819,19 +819,19 @@ def _sample_and_accept_dynamic_tree( skip_top_p = getattr(spec_metadata, "skip_top_p", False) skip_temperature = getattr(spec_metadata, "skip_temperature", False) - if spec_metadata.temperatures is None: + if spec_metadata.request_temperatures is None: temps = torch.ones(num_gens, dtype=torch.float32, device=device) skip_temperature = True else: - temps = spec_metadata.temperatures[gen_slice] + temps = spec_metadata.request_temperatures[gen_slice] top_ks = None - if not skip_top_k and spec_metadata.top_ks is not None: - top_ks = spec_metadata.top_ks[gen_slice] + if not skip_top_k and spec_metadata.request_top_ks is not None: + top_ks = spec_metadata.request_top_ks[gen_slice] top_ps = None - if not skip_top_p and spec_metadata.top_ps is not None: - top_ps = spec_metadata.top_ps[gen_slice] + if not skip_top_p and spec_metadata.request_top_ps is not None: + top_ps = spec_metadata.request_top_ps[gen_slice] # Lazily initialize seed/offset tensors on correct device if self.seed is None: diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 509bd01c72cc..1d2cc4dead7a 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -650,12 +650,18 @@ def _normalize_request_sampling_params( dtype=torch.float32, pin_memory=prefer_pinned()), non_blocking=True) - self.request_top_ks[:len(request_top_ks)].copy_(torch.tensor( - request_top_ks, dtype=torch.int32, pin_memory=prefer_pinned()), - non_blocking=True) - self.request_top_ps[:len(request_top_ps)].copy_(torch.tensor( - request_top_ps, dtype=torch.float32, pin_memory=prefer_pinned()), - non_blocking=True) + self.request_top_ks[:len(request_top_ks)].copy_( + torch.tensor(request_top_ks, + dtype=torch.int32, + pin_memory=prefer_pinned()), + non_blocking=True, + ) + self.request_top_ps[:len(request_top_ps)].copy_( + torch.tensor(request_top_ps, + dtype=torch.float32, + pin_memory=prefer_pinned()), + non_blocking=True, + ) self.skip_temperature = not temperature_enabled self.skip_top_k = not top_k_enabled self.skip_top_p = not top_p_enabled @@ -1010,8 +1016,9 @@ def _sample_and_accept_draft_tokens_rejection( if logits.dim() == 1: logits = logits.unsqueeze(0) + runtime_draft_len = draft_tokens.shape[1] draft_vocab_size = draft_probs.shape[-1] - num_target_tokens = batch_size * (self.max_draft_len + 1) + num_target_tokens = batch_size * (runtime_draft_len + 1) temperatures = spec_metadata.temperatures[:num_target_tokens] top_ks = spec_metadata.top_ks[:num_target_tokens] @@ -1021,14 +1028,16 @@ def _sample_and_accept_draft_tokens_rejection( temperatures, top_ks, top_ps) target_probs = target_probs_flat.reshape(batch_size, - self.max_draft_len + 1, + runtime_draft_len + 1, vocab_size) - runtime_draft_len = draft_probs.shape[1] + 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) if draft_vocab_size != vocab_size: full_draft_probs = torch.zeros( - (batch_size, self.max_draft_len, vocab_size), + (batch_size, runtime_draft_len, vocab_size), dtype=torch.float32, device=device) if d2t is not None: @@ -1106,15 +1115,15 @@ def _compute_and_store_draft_probs( draft_logits_flat = draft_logits.transpose(0, 1).reshape(-1, vocab_size) num_draft_tokens = batch_size * draft_tokens_per_request - if spec_metadata.temperatures is not None: - draft_temps = spec_metadata.temperatures[:batch_size].repeat_interleave( + if spec_metadata.request_temperatures is not None: + draft_temps = spec_metadata.request_temperatures[:batch_size].repeat_interleave( draft_tokens_per_request) - draft_top_ks = spec_metadata.top_ks[:batch_size].repeat_interleave( + draft_top_ks = spec_metadata.request_top_ks[:batch_size].repeat_interleave( draft_tokens_per_request - ) if spec_metadata.top_ks is not None else None - draft_top_ps = spec_metadata.top_ps[:batch_size].repeat_interleave( + ) if spec_metadata.request_top_ks is not None else None + draft_top_ps = spec_metadata.request_top_ps[:batch_size].repeat_interleave( draft_tokens_per_request - ) if spec_metadata.top_ps is not None else None + ) if spec_metadata.request_top_ps is not None else None else: draft_temps = torch.ones(num_draft_tokens, device=device) draft_top_ks = None diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index cf1814cdbf6b..63e982424b39 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -1257,7 +1257,8 @@ def forward( draft_step=i) if spec_metadata.use_rejection_sampling: - draft_logits_list.append(logits[last_tokens_idx].clone()) + draft_logits_list.append( + logits[:gather_ids.numel()].clone()) if self.model_config.mapping.enable_attention_dp and \ getattr(self.model_config.mapping, 'enable_lm_head_tp_in_adp', False): diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index be9aa97c08ea..93c40979a7c3 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -974,6 +974,16 @@ def validate_draft_len_schedule_and_sort(cls, v, info): return dict(sorted(v.items(), key=lambda x: x[0])) return v + @model_validator(mode='after') + def validate_rejection_sampling_config(self): + if self.use_rejection_sampling and getattr(self, 'sa_config', + None) is not None: + raise ValueError( + "use_rejection_sampling is incompatible with sa_config " + "because SA enhancement may override the proposed draft tokens." + ) + return self + @model_validator(mode='after') # 1. Validate that max_concurrency and draft_len_schedule are mutually exclusive. # 2. If max_concurrency is set, translate it to the corresponding draft_len_schedule. @@ -4128,6 +4138,20 @@ def validate_speculative_config(self): exclude={"decoding_type"}) self.speculative_config = Eagle3DecodingConfig(**eagle_data) + if self.speculative_config.use_rejection_sampling: + if self.backend != "pytorch": + raise ValueError( + "use_rejection_sampling is only supported on the " + "PyTorch backend.") + if not isinstance( + self.speculative_config, + (DraftTargetDecodingConfig, Eagle3DecodingConfig, + MTPDecodingConfig, PARDDecodingConfig)): + raise ValueError( + "use_rejection_sampling is only supported for " + "PyTorch one-model speculative decoding paths " + "(Draft_Target, Eagle3, MTP, and PARD).") + if isinstance(self.speculative_config, PARDDecodingConfig): assert self.speculative_config.max_draft_len > 0, "PARD max_draft_len must be > 0" From ffa2736d52fcd9184f315588509353dfdca559d9 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 22 Apr 2026 05:59:53 -0700 Subject: [PATCH 06/17] [None][docs] add docstrings for speculative rejection updates Signed-off-by: ZhaoyangWang --- tensorrt_llm/_torch/speculative/dynamic_tree_ops.py | 2 ++ tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py | 1 + tensorrt_llm/_torch/speculative/interface.py | 2 ++ tensorrt_llm/_torch/speculative/mtp.py | 1 + tensorrt_llm/llmapi/llm_args.py | 3 +++ tests/integration/defs/accuracy/test_llm_api_pytorch.py | 1 + 6 files changed, 10 insertions(+) diff --git a/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py b/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py index f48c7856522f..40dac26978cd 100644 --- a/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py +++ b/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py @@ -51,6 +51,7 @@ def __init__( max_batch_size: int, device: torch.device, ): + """Allocate reusable CUDA buffers for dynamic-tree build/verify ops.""" self.K = dynamic_tree_max_topK self.depth = max_draft_len @@ -89,6 +90,7 @@ def _get_rejection_rng_tensor( buffer: torch.Tensor, name: str, ) -> torch.Tensor: + """Normalize a rejection RNG input to a one-element int64 CUDA tensor.""" if isinstance(value, int): buffer.fill_(value) return buffer diff --git a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py index fa83f28f4440..9302e5dac490 100644 --- a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py @@ -146,6 +146,7 @@ class Eagle3OneModelDynamicTreeWorker(Eagle3OneModelWorker): def __init__( self, spec_config: "EagleDecodingConfig", mapping, use_separate_draft_kv_cache: bool = False ): + """Initialize dynamic-tree specific buffers and helper ops.""" super().__init__(spec_config, mapping, use_separate_draft_kv_cache) assert self.use_dynamic_tree, ( "Eagle3OneModelDynamicTreeWorker requires use_dynamic_tree=True" diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 1d2cc4dead7a..3dd96b140664 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -543,6 +543,7 @@ def populate_sampling_params_for_one_model( DISABLE_TOPP_VAL = 1.0 def _first_or_none(values): + """Return the first sampling parameter value when present.""" return values[0] if values is not None and len(values) > 0 else None def _normalize_request_sampling_params( @@ -551,6 +552,7 @@ def _normalize_request_sampling_params( top_k: Optional[int], top_p: Optional[float], ) -> tuple[float, int, float, bool, bool, bool]: + """Convert request sampling params into normalized per-request scalars.""" is_greedy = SamplingParams.params_imply_greedy_decoding( temperature=temperature, top_k=top_k, diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 63e982424b39..57cb07914dd8 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -1151,6 +1151,7 @@ def forward( draft_model, resource_manager=None, ): + """Run one-model MTP-EAGLE drafting and collect logits for rejection sampling.""" batch_size = attn_metadata.num_seqs num_contexts = attn_metadata.num_contexts diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 93c40979a7c3..1123ac63b53a 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -976,6 +976,7 @@ def validate_draft_len_schedule_and_sort(cls, v, info): @model_validator(mode='after') def validate_rejection_sampling_config(self): + """Reject SA-enhanced configurations that invalidate rejection sampling.""" if self.use_rejection_sampling and getattr(self, 'sa_config', None) is not None: raise ValueError( @@ -3440,6 +3441,7 @@ def validate_build_config_remaining(self): @model_validator(mode="after") def validate_speculative_config(self): + """Validate speculative decoding config against backend-specific restrictions.""" if self.speculative_config: if not self.speculative_config.supports_backend(self.backend): raise ValueError( @@ -4120,6 +4122,7 @@ def set_model_format(self): @model_validator(mode="after") def validate_speculative_config(self): + """Validate speculative decoding config against backend-specific restrictions.""" if self.speculative_config: if not self.speculative_config.supports_backend(self.backend): raise ValueError( diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 4a2bb5a8d37b..92596151314e 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -192,6 +192,7 @@ def test_dummy_load_format(self): ids=["no_dynamic_tree", "dynamic_tree"]) def test_eagle3_rejection_dynamic_tree_smoke(self, use_dynamic_tree, mocker): + """Smoke-test one-model Eagle3 rejection sampling with both tree modes.""" mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", 128) eagle_model_dir = f"{llm_models_root()}/EAGLE3-LLaMA3.1-Instruct-8B" From 72dd59f95d0bd2845c447c891a654824900f7dc8 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 22 Apr 2026 19:43:45 -0700 Subject: [PATCH 07/17] [None][fix] narrow rejection sampling support to Eagle3 Signed-off-by: ZhaoyangWang --- .../_torch/speculative/draft_target.py | 3 +- .../_torch/speculative/dynamic_tree_ops.py | 132 +++---- .../_torch/speculative/eagle3_dynamic_tree.py | 338 ++++++++---------- tensorrt_llm/_torch/speculative/mtp.py | 16 +- tensorrt_llm/_torch/speculative/pard.py | 9 +- tensorrt_llm/_torch/speculative/utils.py | 6 - tensorrt_llm/llmapi/llm_args.py | 15 +- 7 files changed, 202 insertions(+), 317 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/draft_target.py b/tensorrt_llm/_torch/speculative/draft_target.py index d14394452095..c026b6d5b290 100644 --- a/tensorrt_llm/_torch/speculative/draft_target.py +++ b/tensorrt_llm/_torch/speculative/draft_target.py @@ -63,7 +63,6 @@ def __post_init__(self): def prepare(self): """Prepare the metadata before model forward.""" - super().prepare() assert self.request_ids is not None # Update batch indices num_seqs = len(self.request_ids) @@ -315,7 +314,7 @@ def sample_and_accept_draft_tokens( else: draft_tokens = spec_metadata.draft_tokens.reshape(num_gens, self.max_draft_len) - return self._accept_draft_tokens( + return self._sample_and_accept_draft_tokens_base( logits, draft_tokens, num_contexts, batch_size, spec_metadata ) diff --git a/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py b/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py index 40dac26978cd..cb62c8e1b015 100644 --- a/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py +++ b/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py @@ -209,23 +209,19 @@ def verify_dynamic_tree_greedy_out( tree_valid = torch.ones(num_gens, dtype=torch.bool, device=candidates.device) try: - torch.cuda.nvtx.range_push("verify_dynamic_tree_greedy_out") - try: - torch.ops.trtllm.verify_dynamic_tree_greedy_out_op( - candidates, - retrieve_index, - retrieve_next_token, - retrieve_next_sibling, - target_predict, - predicts, - accept_index, - accept_token_num, - accept_token, - tree_valid, - num_spec_step, - ) - finally: - torch.cuda.nvtx.range_pop() + torch.ops.trtllm.verify_dynamic_tree_greedy_out_op( + candidates, + retrieve_index, + retrieve_next_token, + retrieve_next_sibling, + target_predict, + predicts, + accept_index, + accept_token_num, + accept_token, + tree_valid, + num_spec_step, + ) except Exception as e: raise RuntimeError( f"verify_dynamic_tree_greedy_out_op failed: {e}\n" @@ -276,70 +272,48 @@ def verify_dynamic_tree_rejection_from_logits_out( top_k_max = int(enabled_top_k.max().item()) if enabled_top_k.numel() > 0 else 0 try: - torch.cuda.nvtx.range_push("verify_dynamic_tree_rejection_from_logits_out") - try: - torch.cuda.nvtx.range_push( - "trtllm::compute_draft_probs_for_dynamic_tree_rejection_op" - ) - try: - draft_probs_tree = ( - torch.ops.trtllm.compute_draft_probs_for_dynamic_tree_rejection_op( - draft_logits_tree, - temperatures, - num_draft_prob_rows, - target_vocab_size, - top_k, - top_p, - skip_temperature, - d2t=d2t, - top_k_max=top_k_max, - ) - ) - finally: - torch.cuda.nvtx.range_pop() + draft_probs_tree = torch.ops.trtllm.compute_draft_probs_for_dynamic_tree_rejection_op( + draft_logits_tree, + temperatures, + num_draft_prob_rows, + target_vocab_size, + top_k, + top_p, + skip_temperature, + d2t=d2t, + top_k_max=top_k_max, + ) - torch.cuda.nvtx.range_push( - "trtllm::compute_target_probs_for_dynamic_tree_rejection_op" - ) - try: - ( - target_probs_tree, - target_support_indices, - target_support_lengths, - ) = torch.ops.trtllm.compute_target_probs_for_dynamic_tree_rejection_op( - target_logits_tree, - temperatures, - num_draft_tokens, - top_k, - top_p, - skip_temperature, - top_k_max=top_k_max, - ) - finally: - torch.cuda.nvtx.range_pop() + ( + target_probs_tree, + target_support_indices, + target_support_lengths, + ) = torch.ops.trtllm.compute_target_probs_for_dynamic_tree_rejection_op( + target_logits_tree, + temperatures, + num_draft_tokens, + top_k, + top_p, + skip_temperature, + top_k_max=top_k_max, + ) - torch.cuda.nvtx.range_push("trtllm::verify_dynamic_tree_rejection_out_op") - try: - torch.ops.trtllm.verify_dynamic_tree_rejection_out_op( - candidates, - draft_probs_tree, - target_probs_tree, - target_support_indices, - target_support_lengths, - draft_prob_indices, - retrieve_next_token, - retrieve_next_sibling, - accept_index, - accept_tok_num, - accept_token, - num_spec_step, - seed_tensor, - offset_tensor, - ) - finally: - torch.cuda.nvtx.range_pop() - finally: - torch.cuda.nvtx.range_pop() + torch.ops.trtllm.verify_dynamic_tree_rejection_out_op( + candidates, + draft_probs_tree, + target_probs_tree, + target_support_indices, + target_support_lengths, + draft_prob_indices, + retrieve_next_token, + retrieve_next_sibling, + accept_index, + accept_tok_num, + accept_token, + num_spec_step, + seed_tensor, + offset_tensor, + ) except Exception as e: raise RuntimeError( f"dynamic tree rejection op chain failed: {e}\n" diff --git a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py index 9302e5dac490..645b2da0b852 100644 --- a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py @@ -392,6 +392,10 @@ def _relocate_kv_eagerly(self, attn_metadata, batch_size): if cache_mgr is None: return + assert len(set(cache_mgr.num_kv_heads_per_layer)) == 1, ( + "update_kv_cache_draft_token_location_2d requires uniform num_kv_heads across all layers, " + f"but got {cache_mgr.num_kv_heads_per_layer}" + ) torch.ops.tensorrt_llm.update_kv_cache_draft_token_location_2d( self._accepted_draft_indices_tensor[:batch_size], self._num_accepted_tokens_buf[:batch_size], @@ -735,208 +739,152 @@ def _sample_and_accept_dynamic_tree( self, logits, attn_metadata, spec_metadata, batch_size, num_contexts, num_gens ): """Dynamic tree verification using CUDA kernel.""" - torch.cuda.nvtx.range_push("eagle3_dynamic_tree::_sample_and_accept_dynamic_tree") - try: - if num_gens > self._max_batch_size: - raise RuntimeError( - f"Dynamic tree batch size {num_gens} exceeds configured " - f"max_batch_size {self._max_batch_size}" + if num_gens > self._max_batch_size: + raise RuntimeError( + f"Dynamic tree batch size {num_gens} exceeds configured " + f"max_batch_size {self._max_batch_size}" + ) + N = self.tokens_per_gen_step + max_path_len = self._max_path_len + + # Reset output buffers + self._accepted_tokens_buf[:batch_size].zero_() + accepted_tokens = self._accepted_tokens_buf[:batch_size, :max_path_len] + self._num_accepted_tokens_buf[:batch_size].fill_(1) + num_accepted_tokens = self._num_accepted_tokens_buf[:batch_size] + self._accepted_draft_indices_tensor[:batch_size].fill_(-1) + + num_flat_tokens = logits.shape[0] + torch.argmax(logits, dim=-1, out=self._target_tokens_buf[:num_flat_tokens]) + target_tokens = self._target_tokens_buf[:num_flat_tokens] + + # Context requests: accept sampled token + accepted_tokens[:num_contexts, 0] = target_tokens[:num_contexts].to(torch.int32) + + # Generation requests: tree verification + if num_gens > 0: + spec_tree_manager = self.spec_tree_manager + + target_predict = self._target_predict_buf[:num_gens] + target_predict[:] = target_tokens[num_contexts:].reshape(num_gens, N) + + if spec_tree_manager is None: + # CUDA graph warmup: accept only the first token per request + num_accepted_tokens[num_contexts:batch_size] = 1 + accepted_tokens[num_contexts:batch_size, 0] = target_predict[:, 0].to(torch.int32) + self._accepted_draft_indices_tensor[num_contexts:batch_size] = -1 + return accepted_tokens, num_accepted_tokens + + candidates = self._candidates_buf[:num_gens] + candidates[:, 1:] = spec_metadata.draft_tokens.reshape(num_gens, N - 1).to(torch.int64) + candidates[:, 0] = target_predict[:, 0] + + # Slots for gen rows: real py_seq_slot vs dummy; dummy -> slot_has_tree False + gen_slot_ids = spec_tree_manager._all_slot_ids_buf[ + num_contexts : num_contexts + num_gens + ] + tree_valid = spec_tree_manager.slot_has_tree[gen_slot_ids] + + if self._can_use_rejection_sampling(spec_metadata): + vocab_size = logits.shape[-1] + num_ctx_tokens = logits.shape[0] - num_gens * N + device = logits.device + + draft_logits_tree = self._get_unique_draft_logits(num_gens) + draft_prob_indices = self._build_draft_prob_indices(num_gens) + target_logits_tree = logits[num_ctx_tokens:].reshape(-1, vocab_size) + gen_slice = slice(num_contexts, num_contexts + num_gens) + skip_top_k = getattr(spec_metadata, "skip_top_k", False) + skip_top_p = getattr(spec_metadata, "skip_top_p", False) + skip_temperature = getattr(spec_metadata, "skip_temperature", False) + + if spec_metadata.request_temperatures is None: + temps = torch.ones(num_gens, dtype=torch.float32, device=device) + skip_temperature = True + else: + temps = spec_metadata.request_temperatures[gen_slice] + + top_ks = None + if not skip_top_k and spec_metadata.request_top_ks is not None: + top_ks = spec_metadata.request_top_ks[gen_slice] + + top_ps = None + if not skip_top_p and spec_metadata.request_top_ps is not None: + top_ps = spec_metadata.request_top_ps[gen_slice] + + # Lazily initialize seed/offset tensors on correct device + 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) + # Use in-place operations for CUDA graph compatibility + self.seed.add_(1).remainder_(2**31) + + _, accept_index, accept_token_num, accept_token = ( + self.tree_ops_converter.verify_dynamic_tree_rejection_from_logits_out( + candidates, + draft_logits_tree, + target_logits_tree, + draft_prob_indices, + spec_tree_manager.retrieve_next_token[:num_gens], + spec_tree_manager.retrieve_next_sibling[:num_gens], + temps, + top_ks, + top_ps, + skip_temperature, + num_gens, + self._max_path_len, + seed=self.seed, + offset=self.offset, + d2t=self._d2t, + ) ) - N = self.tokens_per_gen_step - max_path_len = self._max_path_len - torch.cuda.nvtx.range_push( - "eagle3_dynamic_tree::_sample_and_accept_dynamic_tree::prepare_outputs" - ) - try: - # Reset output buffers - self._accepted_tokens_buf[:batch_size].zero_() - accepted_tokens = self._accepted_tokens_buf[:batch_size, :max_path_len] - self._num_accepted_tokens_buf[:batch_size].fill_(1) - num_accepted_tokens = self._num_accepted_tokens_buf[:batch_size] - self._accepted_draft_indices_tensor[:batch_size].fill_(-1) - finally: - torch.cuda.nvtx.range_pop() - - torch.cuda.nvtx.range_push( - "eagle3_dynamic_tree::_sample_and_accept_dynamic_tree::compute_target_tokens" - ) - try: - num_flat_tokens = logits.shape[0] - torch.argmax(logits, dim=-1, out=self._target_tokens_buf[:num_flat_tokens]) - target_tokens = self._target_tokens_buf[:num_flat_tokens] - - # Context requests: accept sampled token - accepted_tokens[:num_contexts, 0] = target_tokens[:num_contexts].to(torch.int32) - finally: - torch.cuda.nvtx.range_pop() - - # Generation requests: tree verification - if num_gens > 0: - torch.cuda.nvtx.range_push( - "eagle3_dynamic_tree::_sample_and_accept_dynamic_tree::generation_verify" + self._finalize_dynamic_tree_verify_outputs( + accept_index=accept_index, + accept_token_num=accept_token_num, + accept_token=accept_token, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_contexts=num_contexts, + batch_size=batch_size, + num_gens=num_gens, + max_path_len=max_path_len, ) - try: - spec_tree_manager = self.spec_tree_manager - - target_predict = self._target_predict_buf[:num_gens] - target_predict[:] = target_tokens[num_contexts:].reshape(num_gens, N) - - if spec_tree_manager is None: - # CUDA graph warmup: accept only the first token per request - num_accepted_tokens[num_contexts:batch_size] = 1 - accepted_tokens[num_contexts:batch_size, 0] = target_predict[:, 0].to( - torch.int32 - ) - self._accepted_draft_indices_tensor[num_contexts:batch_size] = -1 - return accepted_tokens, num_accepted_tokens - - candidates = self._candidates_buf[:num_gens] - candidates[:, 1:] = spec_metadata.draft_tokens.reshape(num_gens, N - 1).to( - torch.int64 - ) - candidates[:, 0] = target_predict[:, 0] - - # Slots for gen rows: real py_seq_slot vs dummy; dummy -> slot_has_tree False. - gen_slot_ids = spec_tree_manager._all_slot_ids_buf[ - num_contexts : num_contexts + num_gens - ] - tree_valid = spec_tree_manager.slot_has_tree[gen_slot_ids] - - if self._can_use_rejection_sampling(spec_metadata): - torch.cuda.nvtx.range_push( - "eagle3_dynamic_tree::_sample_and_accept_dynamic_tree::rejection_verify" - ) - try: - vocab_size = logits.shape[-1] - num_ctx_tokens = logits.shape[0] - num_gens * N - device = logits.device - - draft_logits_tree = self._get_unique_draft_logits(num_gens) - draft_prob_indices = self._build_draft_prob_indices(num_gens) - target_logits_tree = logits[num_ctx_tokens:].reshape(-1, vocab_size) - gen_slice = slice(num_contexts, num_contexts + num_gens) - skip_top_k = getattr(spec_metadata, "skip_top_k", False) - skip_top_p = getattr(spec_metadata, "skip_top_p", False) - skip_temperature = getattr(spec_metadata, "skip_temperature", False) - - if spec_metadata.request_temperatures is None: - temps = torch.ones(num_gens, dtype=torch.float32, device=device) - skip_temperature = True - else: - temps = spec_metadata.request_temperatures[gen_slice] - - top_ks = None - if not skip_top_k and spec_metadata.request_top_ks is not None: - top_ks = spec_metadata.request_top_ks[gen_slice] - - top_ps = None - if not skip_top_p and spec_metadata.request_top_ps is not None: - top_ps = spec_metadata.request_top_ps[gen_slice] - - # Lazily initialize seed/offset tensors on correct device - 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) - # Use in-place operations for CUDA graph compatibility - self.seed.add_(1).remainder_(2**31) - - _, accept_index, accept_token_num, accept_token = ( - self.tree_ops_converter.verify_dynamic_tree_rejection_from_logits_out( - candidates, - draft_logits_tree, - target_logits_tree, - draft_prob_indices, - spec_tree_manager.retrieve_next_token[:num_gens], - spec_tree_manager.retrieve_next_sibling[:num_gens], - temps, - top_ks, - top_ps, - skip_temperature, - num_gens, - self._max_path_len, - seed=self.seed, - offset=self.offset, - d2t=self._d2t, - ) - ) - finally: - torch.cuda.nvtx.range_pop() - - torch.cuda.nvtx.range_push( - "eagle3_dynamic_tree::_sample_and_accept_dynamic_tree::finalize_rejection" - ) - try: - self._finalize_dynamic_tree_verify_outputs( - accept_index=accept_index, - accept_token_num=accept_token_num, - accept_token=accept_token, - accepted_tokens=accepted_tokens, - num_accepted_tokens=num_accepted_tokens, - num_contexts=num_contexts, - batch_size=batch_size, - num_gens=num_gens, - max_path_len=max_path_len, - ) - num_accepted_tokens = self._apply_force_accepted_tokens( - num_accepted_tokens, num_contexts, self.max_draft_len - ) - return accepted_tokens, num_accepted_tokens - finally: - torch.cuda.nvtx.range_pop() - - torch.cuda.nvtx.range_push( - "eagle3_dynamic_tree::_sample_and_accept_dynamic_tree::greedy_verify" - ) - try: - _, accept_index, accept_token_num, accept_token = ( - self.tree_ops_converter.verify_dynamic_tree_greedy_out( - candidates, - spec_tree_manager.retrieve_index[:num_gens], - spec_tree_manager.retrieve_next_token[:num_gens], - spec_tree_manager.retrieve_next_sibling[:num_gens], - target_predict, - num_gens, - self._max_path_len, - tree_valid=tree_valid, - ) - ) - finally: - torch.cuda.nvtx.range_pop() - - torch.cuda.nvtx.range_push( - "eagle3_dynamic_tree::_sample_and_accept_dynamic_tree::finalize_greedy" - ) - try: - self._finalize_dynamic_tree_verify_outputs( - accept_index=accept_index, - accept_token_num=accept_token_num, - accept_token=accept_token, - accepted_tokens=accepted_tokens, - num_accepted_tokens=num_accepted_tokens, - num_contexts=num_contexts, - batch_size=batch_size, - num_gens=num_gens, - max_path_len=max_path_len, - ) - finally: - torch.cuda.nvtx.range_pop() - finally: - torch.cuda.nvtx.range_pop() - - torch.cuda.nvtx.range_push( - "eagle3_dynamic_tree::_sample_and_accept_dynamic_tree::apply_force_accepted_tokens" - ) - try: num_accepted_tokens = self._apply_force_accepted_tokens( num_accepted_tokens, num_contexts, self.max_draft_len ) - finally: - torch.cuda.nvtx.range_pop() + return accepted_tokens, num_accepted_tokens + + _, accept_index, accept_token_num, accept_token = ( + self.tree_ops_converter.verify_dynamic_tree_greedy_out( + candidates, + spec_tree_manager.retrieve_index[:num_gens], + spec_tree_manager.retrieve_next_token[:num_gens], + spec_tree_manager.retrieve_next_sibling[:num_gens], + target_predict, + num_gens, + self._max_path_len, + tree_valid=tree_valid, + ) + ) + + self._finalize_dynamic_tree_verify_outputs( + accept_index=accept_index, + accept_token_num=accept_token_num, + accept_token=accept_token, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_contexts=num_contexts, + batch_size=batch_size, + num_gens=num_gens, + max_path_len=max_path_len, + ) + + num_accepted_tokens = self._apply_force_accepted_tokens( + num_accepted_tokens, num_contexts, self.max_draft_len + ) - return accepted_tokens, num_accepted_tokens - finally: - torch.cuda.nvtx.range_pop() + return accepted_tokens, num_accepted_tokens def _can_use_rejection_sampling(self, spec_metadata) -> bool: """Check if rejection sampling can be used for dynamic tree verification. diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 57cb07914dd8..c310cf57816c 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -172,7 +172,6 @@ def all_rank_num_seqs(self, value: List[int]): self.subseq_all_rank_num_tokens = value def prepare(self): - super().prepare() assert self.request_ids is not None num_seqs = len(self.request_ids) # update batch indices @@ -842,7 +841,7 @@ def sample_and_accept_draft_tokens( num_gens, mtp_num_modules) # Use base implementation for strict acceptance - accepted_tokens, num_accepted_tokens = self._accept_draft_tokens( + accepted_tokens, num_accepted_tokens = self._sample_and_accept_draft_tokens_base( logits, draft_tokens, num_contexts, batch_size, spec_metadata) @@ -1151,7 +1150,6 @@ def forward( draft_model, resource_manager=None, ): - """Run one-model MTP-EAGLE drafting and collect logits for rejection sampling.""" batch_size = attn_metadata.num_seqs num_contexts = attn_metadata.num_contexts @@ -1193,7 +1191,6 @@ def forward( # Predict draft tokens next_draft_tokens = [] - draft_logits_list = [] with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): for i in range(self.mtp_num_modules): if i == 0: @@ -1257,10 +1254,6 @@ def forward( self.guided_decoder.execute_draft_batch(logits, draft_step=i) - if spec_metadata.use_rejection_sampling: - draft_logits_list.append( - logits[:gather_ids.numel()].clone()) - if self.model_config.mapping.enable_attention_dp and \ getattr(self.model_config.mapping, 'enable_lm_head_tp_in_adp', False): mapping_lm_head_tp = draft_model.mtp_layers[ @@ -1336,13 +1329,6 @@ def forward( stacked[num_contexts:] = gen_draft_tokens next_draft_tokens = [stacked[:, i] for i in range(stacked.shape[1])] - if spec_metadata.use_rejection_sampling and draft_logits_list: - d2t_inner = getattr(draft_model, "model", draft_model) - d2t_param = getattr(d2t_inner, "d2t", None) - spec_metadata.d2t = d2t_param.data if d2t_param is not None else None - self._compute_and_store_draft_probs(draft_logits_list, - spec_metadata, batch_size) - next_draft_tokens, next_new_tokens = self._prepare_next_tokens( next_draft_tokens, accepted_tokens, spec_metadata, batch_size, num_accepted_tokens) diff --git a/tensorrt_llm/_torch/speculative/pard.py b/tensorrt_llm/_torch/speculative/pard.py index 53c6a0e7cb09..628de48022d7 100644 --- a/tensorrt_llm/_torch/speculative/pard.py +++ b/tensorrt_llm/_torch/speculative/pard.py @@ -36,7 +36,6 @@ def __post_init__(self): self.is_spec_dec_dynamic_tree = False def prepare(self): - super().prepare() assert self.request_ids is not None num_seqs = len(self.request_ids) @@ -211,7 +210,7 @@ def forward( else: logits_for_accept = logits - accepted_tokens, num_accepted_tokens = self._accept_draft_tokens( + accepted_tokens, num_accepted_tokens = self._sample_and_accept_draft_tokens_base( logits_for_accept, draft_tokens, num_contexts, batch_size, spec_metadata ) @@ -298,12 +297,6 @@ def forward( gen_draft_tokens ) - if spec_metadata.use_rejection_sampling: - draft_logits_list = [gen_logits[:, i, :] for i in range(self.max_draft_len)] - d2t_param = getattr(draft_model.model, "d2t", None) - spec_metadata.d2t = d2t_param.data if d2t_param is not None else None - self._compute_and_store_draft_probs(draft_logits_list, spec_metadata, num_gens) - # Pad from (num_gens, K) to (num_gens, 2K-1). if K > 1: pad = torch.zeros((num_gens, K - 1), dtype=torch.int32, device="cuda") diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 0be688b2030b..68ef347b0533 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -52,8 +52,6 @@ def get_spec_metadata(spec_config, max_num_requests=max_num_requests, mtp_hidden_states_manager=spec_resource_manager, allow_advanced_sampling=spec_config.allow_advanced_sampling, - use_rejection_sampling=use_rejection_sampling, - vocab_size=vocab_size, ) if spec_config.spec_dec_mode.is_mtp_eagle(): return Eagle3SpecMetadata( @@ -113,8 +111,6 @@ def get_spec_metadata(spec_config, spec_dec_mode=spec_config.spec_dec_mode, max_num_requests=max_num_requests, allow_advanced_sampling=spec_config.allow_advanced_sampling, - use_rejection_sampling=use_rejection_sampling, - vocab_size=vocab_size, spec_resource_manager=spec_resource_manager, ) if spec_config.spec_dec_mode.is_dflash(): @@ -138,8 +134,6 @@ def get_spec_metadata(spec_config, max_num_requests=max_num_requests, max_num_tokens=max_num_tokens, allow_advanced_sampling=spec_config.allow_advanced_sampling, - use_rejection_sampling=use_rejection_sampling, - vocab_size=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 1123ac63b53a..b6ec5a475b5f 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3441,7 +3441,6 @@ def validate_build_config_remaining(self): @model_validator(mode="after") def validate_speculative_config(self): - """Validate speculative decoding config against backend-specific restrictions.""" if self.speculative_config: if not self.speculative_config.supports_backend(self.backend): raise ValueError( @@ -4122,7 +4121,6 @@ def set_model_format(self): @model_validator(mode="after") def validate_speculative_config(self): - """Validate speculative decoding config against backend-specific restrictions.""" if self.speculative_config: if not self.speculative_config.supports_backend(self.backend): raise ValueError( @@ -4142,18 +4140,11 @@ def validate_speculative_config(self): self.speculative_config = Eagle3DecodingConfig(**eagle_data) if self.speculative_config.use_rejection_sampling: - if self.backend != "pytorch": - raise ValueError( - "use_rejection_sampling is only supported on the " - "PyTorch backend.") - if not isinstance( - self.speculative_config, - (DraftTargetDecodingConfig, Eagle3DecodingConfig, - MTPDecodingConfig, PARDDecodingConfig)): + if not isinstance(self.speculative_config, + Eagle3DecodingConfig): raise ValueError( "use_rejection_sampling is only supported for " - "PyTorch one-model speculative decoding paths " - "(Draft_Target, Eagle3, MTP, and PARD).") + "PyTorch Eagle3 one-model speculative decoding paths.") if isinstance(self.speculative_config, PARDDecodingConfig): assert self.speculative_config.max_draft_len > 0, "PARD max_draft_len must be > 0" From 3fa7548f0ad7912a23cf7f8c922a4af86ea3a8fe Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 29 Apr 2026 03:22:59 -0700 Subject: [PATCH 08/17] [None][fix] fix dynamic tree rejection sampling ordering bias and precision Signed-off-by: ZhaoyangWang --- .../speculativeDecoding/dynamicTreeKernels.cu | 255 ++++++++++++------ .../speculativeDecoding/dynamicTreeKernels.h | 8 +- cpp/tensorrt_llm/thop/dynamicTreeOp.cpp | 16 +- .../_torch/speculative/dynamic_tree_ops.py | 9 +- .../_torch/speculative/eagle3_dynamic_tree.py | 5 +- tensorrt_llm/_torch/speculative/interface.py | 8 +- .../_torch/speculative/spec_tree_manager.py | 2 + 7 files changed, 209 insertions(+), 94 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu index 54fb366f54b7..ba6205e214f9 100644 --- a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu +++ b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu @@ -305,6 +305,11 @@ namespace { constexpr double kGreedyTempThreshold = 1e-4; +bool isTopPEnabled(torch::optional const& topP) +{ + return topP.has_value() && topP->defined(); +} + torch::Tensor computeSoftmaxForProbOp(torch::Tensor logits) { TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA tensor"); @@ -343,7 +348,7 @@ torch::Tensor applyTopKTopPForProbOp(torch::Tensor logits, torch::optional const& topP, int32_t kMax) { bool const hasTopK = topK.has_value() && topK->defined(); - bool const hasTopP = topP.has_value() && topP->defined(); + bool const hasTopP = isTopPEnabled(topP); if (!hasTopK && !hasTopP) { @@ -459,6 +464,7 @@ torch::Tensor computeProbsFromLogits(torch::Tensor const& logits, torch::Tensor = (skipTemperature ? logits : logits.div(safeTemperatures.unsqueeze(1))).contiguous().to(torch::kFloat32); bool const hasTopK = topK.has_value() && topK->defined(); + bool const hasTopP = isTopPEnabled(topP); int64_t const vocabSize = scaledLogits.size(1); int64_t const nRows = scaledLogits.size(0); @@ -467,11 +473,12 @@ torch::Tensor computeProbsFromLogits(torch::Tensor const& logits, torch::Tensor { // Two-stage CUDA top-k/top-p masking (mirrors invokeBatchTopKSampling). maskedLogits = torch::empty_like(scaledLogits); + auto topKForKernel = topK->to(torch::kInt32).contiguous(); + auto topPForKernel = hasTopP ? topP->to(torch::kFloat32).contiguous() : torch::Tensor(); auto stream = at::cuda::getCurrentCUDAStream(scaledLogits.device().index()); invokeTopKTopPMaskingForProbs(scaledLogits.data_ptr(), maskedLogits.data_ptr(), - topK->to(torch::kInt32).contiguous().data_ptr(), - (topP.has_value() && topP->defined()) ? topP->to(torch::kFloat32).contiguous().data_ptr() : nullptr, - kMax, static_cast(nRows), static_cast(vocabSize), stream); + topKForKernel.data_ptr(), hasTopP ? topPForKernel.data_ptr() : nullptr, kMax, + static_cast(nRows), static_cast(vocabSize), stream); } else { @@ -528,9 +535,8 @@ torch::Tensor computeDraftProbsForDynamicTreeRejection(torch::Tensor const& draf auto draftTopK = topK.has_value() && topK->defined() ? torch::optional(topK->repeat_interleave(numDraftProbRows)) : torch::optional(); - auto draftTopP = topP.has_value() && topP->defined() - ? torch::optional(topP->repeat_interleave(numDraftProbRows)) - : torch::optional(); + auto draftTopP = isTopPEnabled(topP) ? torch::optional(topP->repeat_interleave(numDraftProbRows)) + : torch::optional(); auto draftProbs = computeProbsFromLogits(draftLogits, draftTemps, draftTopK, draftTopP, skipTemperature, kMax) .reshape({batchSize, numDraftProbRows, draftVocabSize}); @@ -595,12 +601,11 @@ std::tuple computeTargetProbsForDyn auto targetTopK = topK.has_value() && topK->defined() ? torch::optional(topK->repeat_interleave(numDraftTokens)) : torch::optional(); - auto targetTopP = topP.has_value() && topP->defined() - ? torch::optional(topP->repeat_interleave(numDraftTokens)) - : torch::optional(); + auto targetTopP = isTopPEnabled(topP) ? torch::optional(topP->repeat_interleave(numDraftTokens)) + : torch::optional(); bool const hasTopK = targetTopK.has_value() && targetTopK->defined(); - bool const hasTopP = targetTopP.has_value() && targetTopP->defined(); + bool const hasTopP = isTopPEnabled(targetTopP); bool const hasFiltering = hasTopK || hasTopP; torch::Tensor effectiveTargetTopK; bool hasDisabledTopKRows = false; @@ -637,10 +642,11 @@ std::tuple computeTargetProbsForDyn { // Fast two-stage CUDA path for masked logits. maskedTargetLogits = torch::empty_like(scaledTargetLogits); + auto topKForKernel = effectiveTargetTopK.to(torch::kInt32).contiguous(); + auto topPForKernel = hasTopP ? targetTopP->to(torch::kFloat32).contiguous() : torch::Tensor(); auto stream = at::cuda::getCurrentCUDAStream(scaledTargetLogits.device().index()); invokeTopKTopPMaskingForProbs(scaledTargetLogits.data_ptr(), maskedTargetLogits.data_ptr(), - effectiveTargetTopK.to(torch::kInt32).contiguous().data_ptr(), - hasTopP ? targetTopP->to(torch::kFloat32).contiguous().data_ptr() : nullptr, kMax, + topKForKernel.data_ptr(), hasTopP ? topPForKernel.data_ptr() : nullptr, kMax, static_cast(nRows), static_cast(targetVocabSize), stream); // Extract support indices: the finite positions after masking (at most kMax per row). @@ -1140,43 +1146,68 @@ void invokeVerifyDynamicTreeGreedy(int64_t* predicts, int64_t* acceptIndex, int6 #include +/// Map curand_uniform (0, 1] to [0, 1) so that cumulative-sum sampling +/// never falls off the end of a probability distribution due to float32 +/// rounding. 1.0 is mapped to 0.0 (probability mass epsilon). +__device__ __forceinline__ float curand_uniform_open_right(curandStatePhilox4_32_10_t& state) +{ + float u = curand_uniform(&state); // (0, 1] + return u < 1.0f ? u : 0.0f; // [0, 1) +} + __device__ int64_t sampleFromDistribution(curandStatePhilox4_32_10_t& state, float const* probs, uint32_t vocabSize) { - float r = curand_uniform(&state); + float r = curand_uniform_open_right(state); // [0, 1) float cumsum = 0.0f; - int64_t sampledTok = static_cast(vocabSize) - 1; // fallback: last vocab token + int64_t sampledTok = 0; for (uint32_t v = 0; v < vocabSize; ++v) { cumsum += probs[v]; - if (r <= cumsum) + if (r < cumsum) { sampledTok = static_cast(v); - break; + return sampledTok; } } - return sampledTok; + // Float32 cumsum may not reach 1.0 for large vocabs. + // Fall back to the last token with positive probability. + for (int64_t v = static_cast(vocabSize) - 1; v >= 0; --v) + { + if (probs[v] > 0.0f) + { + return v; + } + } + return static_cast(vocabSize) - 1; } __device__ int64_t sampleFromIndexedDistribution(curandStatePhilox4_32_10_t& state, float const* probs, int32_t const* supportIndices, uint32_t supportSize, uint32_t vocabSize) { - float r = curand_uniform(&state); + float r = curand_uniform_open_right(state); // [0, 1) float cumsum = 0.0f; - int64_t sampledTok = static_cast(vocabSize) - 1; // fallback: last vocab token + int64_t sampledTok = static_cast(vocabSize) - 1; for (uint32_t i = 0; i < supportSize; ++i) { int32_t const tok = supportIndices[i]; cumsum += probs[tok]; - if (r <= cumsum) + if (r < cumsum) { - sampledTok = static_cast(tok); - break; + return static_cast(tok); } } + // Fallback: last support token with positive probability. + for (int64_t i = static_cast(supportSize) - 1; i >= 0; --i) + { + if (probs[supportIndices[i]] > 0.0f) + { + return static_cast(supportIndices[i]); + } + } return sampledTok; } @@ -1192,6 +1223,7 @@ __device__ int64_t sampleFromIndexedDistribution(curandStatePhilox4_32_10_t& sta //! \param draftProbIndices [in] tree position -> draftProbs row [bs, numDraftTokens], root unused. int32. //! \param retrieveNextToken [in] first-child pointer [bs, numDraftTokens], -1=none. int32. //! \param retrieveNextSibling [in] next-sibling pointer [bs, numDraftTokens], -1=none. int32. +//! \param treeValid [in] per-request tree validity flag [bs]. bool. //! \param batchSize batch size. //! \param numDraftProbRows unique draft-prob rows per request. //! \param maxTargetSupportSize support-array width. Zero when targetSupportIndices is null. @@ -1205,9 +1237,9 @@ template __global__ void verifyDynamicTreeRejectionKernel(int64_t* acceptIndex, int64_t* acceptTokenNum, int64_t* acceptToken, int64_t const* candidates, float const* draftProbs, float const* targetProbs, int32_t const* targetSupportIndices, int32_t const* targetSupportLengths, int32_t const* draftProbIndices, int32_t const* retrieveNextToken, - int32_t const* retrieveNextSibling, uint32_t batchSize, uint32_t numDraftProbRows, uint32_t maxTargetSupportSize, - uint32_t numSpeculativeTokens, uint32_t numDraftTokens, uint32_t vocabSize, int64_t const* seed, - int64_t const* offset) + int32_t const* retrieveNextSibling, bool const* treeValid, uint32_t batchSize, uint32_t numDraftProbRows, + uint32_t maxTargetSupportSize, uint32_t numSpeculativeTokens, uint32_t numDraftTokens, uint32_t vocabSize, + int64_t const* seed, int64_t const* offset) { uint32_t bx = blockIdx.x; int32_t const tid = static_cast(threadIdx.x); @@ -1227,16 +1259,22 @@ __global__ void verifyDynamicTreeRejectionKernel(int64_t* acceptIndex, int64_t* __shared__ int32_t sLastAcceptedLocalIdx; __shared__ uint32_t sNumAcceptedTokens; - __shared__ int32_t sCurIndex; __shared__ int32_t sFirstChild; __shared__ bool sHasTerminalToken; - __shared__ bool sAcceptedSibling; __shared__ float sDiffSum; __shared__ float sTargetMass; __shared__ float sPrefixBase; __shared__ int32_t sWinnerIndex; __shared__ int64_t sSampledToken; + // Independent sibling testing: collect all accepted siblings at each depth, + // then select one weighted by target probability to remove ordering bias. + static constexpr int32_t kMaxSiblingsPerDepth = 64; + __shared__ int32_t sAccSibIdx[kMaxSiblingsPerDepth]; + __shared__ int64_t sAccSibTok[kMaxSiblingsPerDepth]; + __shared__ float sAccSibQProb[kMaxSiblingsPerDepth]; + __shared__ int32_t sNumAccSiblings; + uint32_t batchOffset = bx * numDraftTokens; curandStatePhilox4_32_10_t state; @@ -1246,6 +1284,33 @@ __global__ void verifyDynamicTreeRejectionKernel(int64_t* acceptIndex, int64_t* static_cast(seed[0]), static_cast(bx), static_cast(offset[0]), &state); } __syncthreads(); + bool const hasCompactTargetSupport = targetSupportIndices != nullptr && targetSupportLengths != nullptr; + + // First-gen or dummy request: no valid tree exists yet. Sample directly + // from the target distribution at the root and skip tree traversal. + if (treeValid != nullptr && !treeValid[bx]) + { + if (tid == 0) + { + float const* tProbs = targetProbs + static_cast(bx) * numDraftTokens * vocabSize; + int64_t sampledToken; + if (hasCompactTargetSupport) + { + uint32_t const supportOffset = static_cast(bx) * numDraftTokens * maxTargetSupportSize; + uint32_t const supportSize = static_cast(targetSupportLengths[batchOffset]); + sampledToken = sampleFromIndexedDistribution( + state, tProbs, targetSupportIndices + supportOffset, supportSize, vocabSize); + } + else + { + sampledToken = sampleFromDistribution(state, tProbs, vocabSize); + } + acceptIndex[bx * numSpeculativeTokens] = 0; + acceptTokenNum[bx] = 0; + acceptToken[bx * numSpeculativeTokens] = sampledToken; + } + return; + } // Root (depth 0): initialize path state at tree position 0. // @@ -1278,33 +1343,59 @@ __global__ void verifyDynamicTreeRejectionKernel(int64_t* acceptIndex, int64_t* sLastAcceptedLocalIdx = 0; acceptIndex[bx * numSpeculativeTokens] = sLastAcceptedLocalIdx; sNumAcceptedTokens = 0; - sCurIndex = 0; sHasTerminalToken = false; } __syncthreads(); - bool const hasCompactTargetSupport = targetSupportIndices != nullptr && targetSupportLengths != nullptr; - for (uint32_t j = 1; j < numSpeculativeTokens; ++j) { - // Advance to the first child of the last accepted node. - // Continuing the example above: - // j = 1, curIndex = 0 (E) -> firstChild = F1 - // j = 2, curIndex = 1 (F1) -> firstChild = G1 + // Get first child of the last accepted node. if (tid == 0) { sFirstChild = retrieveNextToken[batchOffset + sLastAcceptedLocalIdx]; - sCurIndex = sFirstChild; - sAcceptedSibling = false; } __syncthreads(); - while (sCurIndex != -1 && !sAcceptedSibling) + // Leaf node: no children at this depth. + // Emit bonus token from the target distribution at the last accepted position. + if (sFirstChild == -1) { if (tid == 0) { - int32_t const draftLocalIdx = sCurIndex; // retrieveIndex is identity: draftLocalIdx == curIndex - int64_t const draftTokenId = candidates[batchOffset + sCurIndex]; - int32_t const draftProbRow = draftProbIndices[batchOffset + sCurIndex]; + float const* tProbs = targetProbs + + (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * vocabSize; + if (hasCompactTargetSupport) + { + uint32_t const supportOffset + = (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * maxTargetSupportSize; + uint32_t const supportSize + = static_cast(targetSupportLengths[batchOffset + sLastAcceptedLocalIdx]); + acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] = sampleFromIndexedDistribution( + state, tProbs, targetSupportIndices + supportOffset, supportSize, vocabSize); + } + else + { + acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] + = sampleFromDistribution(state, tProbs, vocabSize); + } + sHasTerminalToken = true; + } + __syncthreads(); + break; + } + + // ---- Phase 1: Independently test ALL siblings with Bernoulli rejection ---- + // Unlike sequential testing which accepts the first passing sibling + // (creating ordering bias), we test every sibling and collect all that + // pass. This removes the systematic preference for earlier siblings in + // the linked list and produces output closer to the target distribution. + if (tid == 0) + { + sNumAccSiblings = 0; + int32_t childIdx = sFirstChild; + while (childIdx != -1 && sNumAccSiblings < kMaxSiblingsPerDepth) + { + int64_t const draftTokenId = candidates[batchOffset + childIdx]; + int32_t const draftProbRow = draftProbIndices[batchOffset + childIdx]; float const pDraft = draftProbs[(static_cast(bx) * numDraftProbRows + draftProbRow) * vocabSize + draftTokenId]; @@ -1313,56 +1404,60 @@ __global__ void verifyDynamicTreeRejectionKernel(int64_t* acceptIndex, int64_t* + draftTokenId]; float const acceptProb = fminf(1.0f, pTarget / (pDraft + 1e-10f)); - float const u = curand_uniform(&state); + float const u = curand_uniform_open_right(state); if (u < acceptProb) { - acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] = draftTokenId; - ++sNumAcceptedTokens; - acceptIndex[bx * numSpeculativeTokens + sNumAcceptedTokens] = draftLocalIdx; - sLastAcceptedLocalIdx = draftLocalIdx; - sAcceptedSibling = true; - } - else - { - sCurIndex = retrieveNextSibling[batchOffset + sCurIndex]; + int32_t const idx = sNumAccSiblings++; + sAccSibIdx[idx] = childIdx; + sAccSibTok[idx] = draftTokenId; + sAccSibQProb[idx] = pTarget; } + childIdx = retrieveNextSibling[batchOffset + childIdx]; } - __syncthreads(); } + __syncthreads(); - if (sCurIndex == -1) + // ---- Phase 2: Select among accepted siblings or emit correction ---- + if (sNumAccSiblings > 0) { - // All siblings exhausted. Two sub-cases: - // (a) firstChild == -1: leaf node, no draft tokens at this depth. - // Emit the final bonus token sampled from q(.|lastAcceptedLocalIdx). - // (b) firstChild != -1: every sibling was rejected -> sample correction token - // from relu(q - p) at firstChild's position to restore the target distribution. - if (sFirstChild == -1) + // At least one sibling accepted. Select one weighted by target + // probability q(ci) so that the choice follows the target distribution. + if (tid == 0) { - if (tid == 0) + int32_t selectedIdx = 0; + if (sNumAccSiblings > 1) { - float const* tProbs = targetProbs - + (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * vocabSize; - if (hasCompactTargetSupport) + float totalQ = 0.0f; + for (int32_t i = 0; i < sNumAccSiblings; ++i) { - uint32_t const supportOffset - = (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) - * maxTargetSupportSize; - uint32_t const supportSize - = static_cast(targetSupportLengths[batchOffset + sLastAcceptedLocalIdx]); - acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] = sampleFromIndexedDistribution( - state, tProbs, targetSupportIndices + supportOffset, supportSize, vocabSize); + totalQ += sAccSibQProb[i]; } - else + float const r = curand_uniform_open_right(state) * totalQ; + float cumsum = 0.0f; + selectedIdx = sNumAccSiblings - 1; // fallback + for (int32_t i = 0; i < sNumAccSiblings; ++i) { - acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] - = sampleFromDistribution(state, tProbs, vocabSize); + cumsum += sAccSibQProb[i]; + if (r < cumsum) + { + selectedIdx = i; + break; + } } - sHasTerminalToken = true; } + + int32_t const childIdx = sAccSibIdx[selectedIdx]; + acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] = sAccSibTok[selectedIdx]; + ++sNumAcceptedTokens; + acceptIndex[bx * numSpeculativeTokens + sNumAcceptedTokens] = childIdx; + sLastAcceptedLocalIdx = childIdx; } - else + __syncthreads(); + } + else + { + // All siblings rejected -> sample correction token from relu(q - p). { int32_t const draftProbRow = draftProbIndices[batchOffset + sFirstChild]; float const* tProbs @@ -1400,7 +1495,7 @@ __global__ void verifyDynamicTreeRejectionKernel(int64_t* acceptIndex, int64_t* if (useDiff) { - float const r = curand_uniform(&state); + float const r = curand_uniform_open_right(state); float cumsum = 0.0f; for (uint32_t i = 0; i < targetSupportSize; ++i) { @@ -1445,7 +1540,7 @@ __global__ void verifyDynamicTreeRejectionKernel(int64_t* acceptIndex, int64_t* sSampledToken = static_cast(vocabSize) - 1; if (diffSum > 1e-10f) { - sTargetMass = curand_uniform(&state) * diffSum; + sTargetMass = curand_uniform_open_right(state) * diffSum; } else { @@ -1541,7 +1636,7 @@ __global__ void verifyDynamicTreeRejectionKernel(int64_t* acceptIndex, int64_t* void invokeVerifyDynamicTreeRejection(int64_t* acceptIndex, int64_t* acceptTokenNum, int64_t* acceptToken, int64_t const* candidates, float const* draftProbs, float const* targetProbs, int32_t const* targetSupportIndices, int32_t const* targetSupportLengths, int32_t const* draftProbIndices, int32_t const* retrieveNextToken, - int32_t const* retrieveNextSibling, SizeType32 batchSize, SizeType32 numDraftProbRows, + int32_t const* retrieveNextSibling, bool const* treeValid, SizeType32 batchSize, SizeType32 numDraftProbRows, SizeType32 maxTargetSupportSize, SizeType32 numDraftTokens, SizeType32 numSpecStep, SizeType32 vocabSize, int64_t const* seed, int64_t const* offset, cudaStream_t stream) { @@ -1551,8 +1646,8 @@ void invokeVerifyDynamicTreeRejection(int64_t* acceptIndex, int64_t* acceptToken verifyDynamicTreeRejectionKernel<<>>(acceptIndex, acceptTokenNum, acceptToken, candidates, draftProbs, targetProbs, targetSupportIndices, targetSupportLengths, - draftProbIndices, retrieveNextToken, retrieveNextSibling, batchSize, numDraftProbRows, maxTargetSupportSize, - numSpecStep, numDraftTokens, vocabSize, seed, offset); + draftProbIndices, retrieveNextToken, retrieveNextSibling, treeValid, batchSize, numDraftProbRows, + maxTargetSupportSize, numSpecStep, numDraftTokens, vocabSize, seed, offset); sync_check_cuda_error(stream); } diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h index 78eecfdcbe96..48e788e91e3e 100644 --- a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h +++ b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h @@ -126,6 +126,7 @@ void invokeVerifyDynamicTreeGreedy(int64_t* predicts, int64_t* acceptIndex, int6 //! corresponding row in draftProbs. Root is unused. //! \param retrieveNextToken [batchSize, numDraftTokens] int32 first-child pointer, -1=none. //! \param retrieveNextSibling [batchSize, numDraftTokens] int32 next-sibling pointer, -1=none. +//! \param treeValid [batchSize] bool; false means no valid tree exists for this request. //! \param batchSize runtime::SizeType32. //! \param numDraftProbRows runtime::SizeType32. Number of unique draft-prob rows per request. //! \param maxTargetSupportSize runtime::SizeType32. Third dim of targetSupportIndices. Can be zero. @@ -138,9 +139,10 @@ void invokeVerifyDynamicTreeGreedy(int64_t* predicts, int64_t* acceptIndex, int6 void invokeVerifyDynamicTreeRejection(int64_t* acceptIndex, int64_t* acceptTokenNum, int64_t* acceptToken, int64_t const* candidates, float const* draftProbs, float const* targetProbs, int32_t const* targetSupportIndices, int32_t const* targetSupportLengths, int32_t const* draftProbIndices, int32_t const* retrieveNextToken, - int32_t const* retrieveNextSibling, runtime::SizeType32 batchSize, runtime::SizeType32 numDraftProbRows, - runtime::SizeType32 maxTargetSupportSize, runtime::SizeType32 numDraftTokens, runtime::SizeType32 numSpecStep, - runtime::SizeType32 vocabSize, int64_t const* seed, int64_t const* offset, cudaStream_t stream); + int32_t const* retrieveNextSibling, bool const* treeValid, runtime::SizeType32 batchSize, + runtime::SizeType32 numDraftProbRows, runtime::SizeType32 maxTargetSupportSize, runtime::SizeType32 numDraftTokens, + runtime::SizeType32 numSpecStep, runtime::SizeType32 vocabSize, int64_t const* seed, int64_t const* offset, + cudaStream_t stream); //! \brief Compute draft probabilities for dynamic-tree rejection sampling from logits. //! \param draftLogits [batchSize * numDraftProbRows, draftVocabSize], on GPU. diff --git a/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp b/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp index e795a04c96a1..2790f06cf0d6 100644 --- a/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp +++ b/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp @@ -215,8 +215,8 @@ void verify_dynamic_tree_greedy_out_op(th::Tensor& candidates, th::Tensor& retri //! Accepts draft tokens by rejection sampling at each depth using pre-computed probabilities. void verify_dynamic_tree_rejection_out_op(th::Tensor& candidates, th::Tensor& draftProbs, th::Tensor& targetProbs, th::Tensor& targetSupportIndices, th::Tensor& targetSupportLengths, th::Tensor& draftProbIndices, - th::Tensor& retrieveNextToken, th::Tensor& retrieveNextSibling, th::Tensor& acceptIndex, th::Tensor& acceptTokenNum, - th::Tensor& acceptToken, int64_t numSpecStep, th::Tensor& seed, th::Tensor& offset) + th::Tensor& retrieveNextToken, th::Tensor& retrieveNextSibling, th::Tensor& treeValid, th::Tensor& acceptIndex, + th::Tensor& acceptTokenNum, th::Tensor& acceptToken, int64_t numSpecStep, th::Tensor& seed, th::Tensor& offset) { TORCH_CHECK(candidates.dim() == 2, "candidates must be 2D tensor"); TORCH_CHECK(draftProbs.dim() == 3, "draftProbs must be 3D tensor"); @@ -228,18 +228,21 @@ void verify_dynamic_tree_rejection_out_op(th::Tensor& candidates, th::Tensor& dr TORCH_CHECK(draftProbIndices.dim() == 2, "draftProbIndices must be 2D tensor"); TORCH_CHECK(retrieveNextToken.dim() == 2, "retrieveNextToken must be 2D tensor"); TORCH_CHECK(retrieveNextSibling.dim() == 2, "retrieveNextSibling must be 2D tensor"); + TORCH_CHECK(treeValid.dim() == 1, "treeValid must be 1D tensor"); TORCH_CHECK(candidates.scalar_type() == torch::kInt64, "candidates must be int64 tensor"); TORCH_CHECK(draftProbs.scalar_type() == torch::kFloat32, "draftProbs must be float32 tensor"); TORCH_CHECK(targetProbs.scalar_type() == torch::kFloat32, "targetProbs must be float32 tensor"); TORCH_CHECK(targetSupportIndices.scalar_type() == torch::kInt32, "targetSupportIndices must be int32 tensor"); TORCH_CHECK(targetSupportLengths.scalar_type() == torch::kInt32, "targetSupportLengths must be int32 tensor"); TORCH_CHECK(draftProbIndices.scalar_type() == torch::kInt32, "draftProbIndices must be int32 tensor"); + TORCH_CHECK(treeValid.scalar_type() == torch::kBool, "treeValid must be bool tensor"); TORCH_CHECK(candidates.is_cuda(), "candidates must be a CUDA tensor"); TORCH_CHECK(draftProbs.is_cuda(), "draftProbs must be a CUDA tensor"); TORCH_CHECK(targetProbs.is_cuda(), "targetProbs must be a CUDA tensor"); TORCH_CHECK(draftProbIndices.is_cuda(), "draftProbIndices must be a CUDA tensor"); TORCH_CHECK(retrieveNextToken.is_cuda(), "retrieveNextToken must be a CUDA tensor"); TORCH_CHECK(retrieveNextSibling.is_cuda(), "retrieveNextSibling must be a CUDA tensor"); + TORCH_CHECK(treeValid.is_cuda(), "treeValid must be a CUDA tensor"); TORCH_CHECK(acceptIndex.is_cuda(), "acceptIndex must be a CUDA tensor"); TORCH_CHECK(acceptTokenNum.is_cuda(), "acceptTokenNum must be a CUDA tensor"); TORCH_CHECK(acceptToken.is_cuda(), "acceptToken must be a CUDA tensor"); @@ -280,6 +283,7 @@ void verify_dynamic_tree_rejection_out_op(th::Tensor& candidates, th::Tensor& dr TORCH_CHECK(retrieveNextToken.size(1) == numDraftTokens, "retrieveNextToken size mismatch"); TORCH_CHECK(retrieveNextSibling.size(0) == batchSize, "retrieveNextSibling batch size mismatch"); TORCH_CHECK(retrieveNextSibling.size(1) == numDraftTokens, "retrieveNextSibling size mismatch"); + TORCH_CHECK(treeValid.size(0) >= batchSize, "treeValid buffer too small"); TORCH_CHECK(draftProbs.device() == candidates.device(), "draftProbs must be on the same device as candidates"); TORCH_CHECK(targetProbs.device() == candidates.device(), "targetProbs must be on the same device as candidates"); TORCH_CHECK( @@ -288,6 +292,7 @@ void verify_dynamic_tree_rejection_out_op(th::Tensor& candidates, th::Tensor& dr "retrieveNextToken must be on the same device as candidates"); TORCH_CHECK(retrieveNextSibling.device() == candidates.device(), "retrieveNextSibling must be on the same device as candidates"); + TORCH_CHECK(treeValid.device() == candidates.device(), "treeValid must be on the same device as candidates"); TORCH_CHECK(acceptIndex.scalar_type() == torch::kInt64, "acceptIndex must be int64 tensor"); TORCH_CHECK(acceptTokenNum.scalar_type() == torch::kInt64, "acceptTokenNum must be int64 tensor"); TORCH_CHECK(acceptToken.scalar_type() == torch::kInt64, "acceptToken must be int64 tensor"); @@ -319,8 +324,9 @@ void verify_dynamic_tree_rejection_out_op(th::Tensor& candidates, th::Tensor& dr targetSupportIndices.numel() > 0 ? targetSupportIndices.data_ptr() : nullptr, targetSupportLengths.numel() > 0 ? targetSupportLengths.data_ptr() : nullptr, draftProbIndices.data_ptr(), retrieveNextToken.data_ptr(), - retrieveNextSibling.data_ptr(), batchSize, numDraftProbRows, maxTargetSupportSize, numDraftTokens, - numSpecStep, vocabSize, seed.data_ptr(), offset.data_ptr(), stream); + retrieveNextSibling.data_ptr(), treeValid.data_ptr(), batchSize, numDraftProbRows, + maxTargetSupportSize, numDraftTokens, numSpecStep, vocabSize, seed.data_ptr(), + offset.data_ptr(), stream); } th::Tensor compute_draft_probs_for_dynamic_tree_rejection_op(th::Tensor draftLogits, th::Tensor temperatures, @@ -430,7 +436,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) "verify_dynamic_tree_rejection_out_op(" "Tensor candidates, Tensor draftProbs, Tensor targetProbs, Tensor targetSupportIndices, " "Tensor targetSupportLengths, Tensor draftProbIndices, " - "Tensor retrieveNextToken, Tensor retrieveNextSibling, " + "Tensor retrieveNextToken, Tensor retrieveNextSibling, Tensor treeValid, " "Tensor(a!) acceptIndex, Tensor(b!) acceptTokenNum, Tensor(c!) acceptToken, " "int numSpecStep, Tensor seed, Tensor offset) -> ()"); } diff --git a/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py b/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py index cb62c8e1b015..9aa6f7c09b10 100644 --- a/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py +++ b/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py @@ -239,6 +239,7 @@ def verify_dynamic_tree_rejection_from_logits_out( draft_prob_indices: torch.Tensor, retrieve_next_token: torch.Tensor, retrieve_next_sibling: torch.Tensor, + tree_valid: torch.Tensor, temperatures: torch.Tensor, top_k: torch.Tensor | None, top_p: torch.Tensor | None, @@ -254,7 +255,8 @@ def verify_dynamic_tree_rejection_from_logits_out( This path keeps draft/target logits as inputs, computes unique draft and target probabilities with separate CUDA ops, then runs the tree rejection kernel as a third CUDA op. `draft_prob_indices` maps each - tree position to its shared draft-prob row. + tree position to its shared draft-prob row. `tree_valid` guards + first-gen and dummy requests that do not have a usable tree yet. """ accept_index = self._rej_accept_index_buf[:num_gens] accept_token = self._rej_accept_token_buf[:num_gens] @@ -265,6 +267,10 @@ def verify_dynamic_tree_rejection_from_logits_out( num_draft_prob_rows = draft_logits_tree.shape[0] // num_gens target_vocab_size = target_logits_tree.shape[-1] + if tree_valid is None: + tree_valid = torch.ones(num_gens, dtype=torch.bool, device=candidates.device) + tree_valid = tree_valid.contiguous() + if top_k is None: top_k_max = 0 else: @@ -307,6 +313,7 @@ def verify_dynamic_tree_rejection_from_logits_out( draft_prob_indices, retrieve_next_token, retrieve_next_sibling, + tree_valid, accept_index, accept_tok_num, accept_token, diff --git a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py index 645b2da0b852..0b61a70dfc11 100644 --- a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py @@ -710,8 +710,10 @@ def _forward_draft_loop( # Rejection sampling needs to map each final draft token back to its history # buffer index to retrieve the corresponding draft logits. Greedy verification # doesn't need this mapping since it only compares token IDs. + # Only save the gen rows (skip context rows) so that + # _build_draft_prob_indices can read them back at [:num_gens]. if spec_metadata.use_rejection_sampling: - self._topk_score_indices_buf[:batch_size].copy_(topk_score_indices) + self._topk_score_indices_buf[:num_gens].copy_(topk_score_indices[num_contexts:]) if spec_tree_manager is not None: # Gen-only; spec-dec batch is [:num_gens] (not num_contexts-offset). @@ -827,6 +829,7 @@ def _sample_and_accept_dynamic_tree( draft_prob_indices, spec_tree_manager.retrieve_next_token[:num_gens], spec_tree_manager.retrieve_next_sibling[:num_gens], + tree_valid, temps, top_ks, top_ps, diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 3dd96b140664..3ebd41d27e85 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -1024,7 +1024,7 @@ def _sample_and_accept_draft_tokens_rejection( temperatures = spec_metadata.temperatures[:num_target_tokens] top_ks = spec_metadata.top_ks[:num_target_tokens] - top_ps = spec_metadata.top_ps[:num_target_tokens] + top_ps = None if spec_metadata.skip_top_p else spec_metadata.top_ps[:num_target_tokens] target_probs_flat = compute_probs_from_logits(logits.clone(), temperatures, top_ks, @@ -1123,9 +1123,9 @@ def _compute_and_store_draft_probs( draft_top_ks = spec_metadata.request_top_ks[:batch_size].repeat_interleave( draft_tokens_per_request ) if spec_metadata.request_top_ks is not None else None - draft_top_ps = spec_metadata.request_top_ps[:batch_size].repeat_interleave( - draft_tokens_per_request - ) if spec_metadata.request_top_ps is not None else None + draft_top_ps = (spec_metadata.request_top_ps[:batch_size].repeat_interleave( + draft_tokens_per_request) if not spec_metadata.skip_top_p + and spec_metadata.request_top_ps is not None else None) else: draft_temps = torch.ones(num_draft_tokens, device=device) draft_top_ks = None diff --git a/tensorrt_llm/_torch/speculative/spec_tree_manager.py b/tensorrt_llm/_torch/speculative/spec_tree_manager.py index d6e7cd826d06..80d9228e6582 100644 --- a/tensorrt_llm/_torch/speculative/spec_tree_manager.py +++ b/tensorrt_llm/_torch/speculative/spec_tree_manager.py @@ -234,6 +234,8 @@ def mark_tree_valid(self, slot_ids, count): return # index_fill_: graph-capture-safe (no fancy indexing). self.slot_has_tree.index_fill_(0, slot_ids[:count], True) + # CUDA graph padding may use the dummy slot; it must never carry a real tree. + self.slot_has_tree.narrow(0, self._dummy_slot_id, 1).fill_(False) def mark_tree_invalid(self, slot_id): """Clear validity when a slot is freed.""" From c66164747a8498bb5427dae0f850424a2e23b3b3 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 29 Apr 2026 03:23:32 -0700 Subject: [PATCH 09/17] Fix dynamic tree issue Signed-off-by: ZhaoyangWang --- tensorrt_llm/_torch/speculative/dynamic_tree_ops.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py b/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py index 9aa6f7c09b10..e8dfd9ee7ec2 100644 --- a/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py +++ b/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py @@ -146,6 +146,14 @@ def build_dynamic_tree( # Non-packed path ignores this (bool [bs,N,N] still has a last dim). num_int32_per_row = tree_mask.shape[-1] + # The CUDA builder writes only active tree links/bits. Clear the reused + # work buffers first so stale slot/tree state cannot leak into this tree. + tree_mask[:bs].zero_() + positions[:bs].zero_() + retrieve_index[:bs].zero_() + retrieve_next_token[:bs].fill_(-1) + retrieve_next_sibling[:bs].fill_(-1) + # Call CUDA kernel in-place try: torch.ops.trtllm.build_dynamic_tree_op( From a385342cb7c240f23204bbdcb0ba50f849095a28 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 29 Apr 2026 19:14:47 -0700 Subject: [PATCH 10/17] [None][fix] address rejection sampling QA review feedback Signed-off-by: ZhaoyangWang --- tensorrt_llm/_torch/speculative/dynamic_tree_ops.py | 13 ++++++++++--- tensorrt_llm/llmapi/llm_args.py | 4 ++-- .../integration/test_lists/qa/llm_function_core.txt | 2 ++ 3 files changed, 14 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py b/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py index e8dfd9ee7ec2..d0126d534f22 100644 --- a/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py +++ b/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py @@ -98,11 +98,11 @@ def _get_rejection_rng_tensor( raise TypeError(f"{name} must be an int or torch.Tensor, got {type(value)!r}") if value.dtype != torch.int64: raise TypeError(f"{name} must be int64 tensor, got {value.dtype}") - if value.numel() < 1: - raise ValueError(f"{name} tensor must have at least one element") + if value.numel() != 1: + raise ValueError(f"{name} tensor must have exactly one element, got {value.numel()}") if value.device != buffer.device: raise ValueError(f"{name} tensor must be on device {buffer.device}, got {value.device}") - return value.reshape(-1)[:1] + return value.reshape(-1) def build_dynamic_tree( self, @@ -272,6 +272,13 @@ def verify_dynamic_tree_rejection_from_logits_out( seed_tensor = self._get_rejection_rng_tensor(seed, self._rej_seed_buf, "seed") offset_tensor = self._get_rejection_rng_tensor(offset, self._rej_offset_buf, "offset") num_draft_tokens = candidates.shape[1] + if num_gens <= 0: + raise ValueError(f"num_gens must be positive, got {num_gens}") + if draft_logits_tree.shape[0] % num_gens != 0: + raise ValueError( + f"draft_logits_tree rows ({draft_logits_tree.shape[0]}) must be divisible by " + f"num_gens ({num_gens})" + ) num_draft_prob_rows = draft_logits_tree.shape[0] // num_gens target_vocab_size = target_logits_tree.shape[-1] diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index b6ec5a475b5f..b9d20e67b85f 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -919,8 +919,8 @@ class DecodingBaseConfig(StrictBaseModel): status="prototype", description= "If true, enables rejection sampling for one-model speculative decoding paths. " - "This is intended for non-greedy sampling configurations on the PyTorch backend." - ) + "This is intended for non-greedy sampling configurations on the PyTorch backend. " + "The non-dynamic-tree one-model path requires FlashInfer.") # If set, drafting is allowed to use chain drafter. _allow_chain_drafter: bool = PrivateAttr(True) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 0044105fb351..4d4a6a6daa73 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -585,6 +585,8 @@ accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_dummy_load_forma accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3[sampler_async_worker=False-eagle3_one_model=False-overlap_scheduler=False] accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3[sampler_async_worker=False-eagle3_one_model=True-overlap_scheduler=True] accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3[sampler_async_worker=True-eagle3_one_model=True-overlap_scheduler=True] +accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3_rejection_dynamic_tree_smoke[no_dynamic_tree] +accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3_rejection_dynamic_tree_smoke[dynamic_tree] accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3_sa accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3_sa_global_pool accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8[fp8kv=False-attn_backend=FLASHINFER-torch_compile=False] From 26c4fa496f6cebd3f3462fee332fe962302cf079 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 6 May 2026 19:50:30 -0700 Subject: [PATCH 11/17] [None][perf] optimize dynamic tree rejection sampling Signed-off-by: ZhaoyangWang --- .../speculativeDecoding/dynamicTreeKernels.cu | 880 +++++++++++++++--- .../speculativeDecoding/dynamicTreeKernels.h | 4 +- cpp/tensorrt_llm/thop/dynamicTreeOp.cpp | 20 +- .../_torch/speculative/dynamic_tree_ops.py | 3 + .../_torch/speculative/eagle3_dynamic_tree.py | 8 + tensorrt_llm/_torch/speculative/interface.py | 18 +- 6 files changed, 766 insertions(+), 167 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu index ba6205e214f9..8b94d9dfb559 100644 --- a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu +++ b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu @@ -27,9 +27,12 @@ #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/reduceKernelUtils.cuh" +#include "tensorrt_llm/common/vec_dtypes.cuh" #include "tensorrt_llm/kernels/decodingCommon.h" #include #include +#include +#include #include #include TRTLLM_NAMESPACE_BEGIN @@ -307,7 +310,7 @@ constexpr double kGreedyTempThreshold = 1e-4; bool isTopPEnabled(torch::optional const& topP) { - return topP.has_value() && topP->defined(); + return topP.has_value() && topP->defined() && topP->lt(1.0).any().item(); } torch::Tensor computeSoftmaxForProbOp(torch::Tensor logits) @@ -334,6 +337,110 @@ torch::Tensor computeSoftmaxForProbOp(torch::Tensor logits) return probs; } +struct DraftProbMaxFloatOp +{ + __device__ __forceinline__ float operator()(float a, float b) const + { + return a > b ? a : b; + } +}; + +template +__global__ void computeDraftProbsSkipAllKernel(float const* draftLogits, int32_t const* d2t, float* draftProbs, + int32_t nRows, int32_t draftVocabSize, int32_t targetVocabSize) +{ + int32_t const rowId = static_cast(blockIdx.x); + int32_t const tid = static_cast(threadIdx.x); + if (rowId >= nRows) + { + return; + } + + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage tempStorage; + __shared__ float sMaxLogit; + __shared__ float sExpSum; + + float const* rowLogits = draftLogits + static_cast(rowId) * draftVocabSize; + float* rowProbs = draftProbs + static_cast(rowId) * targetVocabSize; + + for (int32_t v = tid; v < targetVocabSize; v += BLOCK_SIZE) + { + rowProbs[v] = 0.0f; + } + + float localMax = -FLT_MAX; + for (int32_t v = tid; v < draftVocabSize; v += BLOCK_SIZE) + { + localMax = fmaxf(localMax, rowLogits[v]); + } + + float const blockMax = BlockReduce(tempStorage).Reduce(localMax, DraftProbMaxFloatOp{}); + if (tid == 0) + { + sMaxLogit = blockMax; + } + __syncthreads(); + + float localSum = 0.0f; + for (int32_t v = tid; v < draftVocabSize; v += BLOCK_SIZE) + { + localSum += __expf(rowLogits[v] - sMaxLogit); + } + + float const blockSum = BlockReduce(tempStorage).Sum(localSum); + if (tid == 0) + { + constexpr float kFloatSoftmaxEpsilon = 1e-6f; + sExpSum = blockSum + kFloatSoftmaxEpsilon; + } + __syncthreads(); + + for (int32_t v = tid; v < draftVocabSize; v += BLOCK_SIZE) + { + int64_t const targetIdx + = d2t != nullptr ? static_cast(v) + static_cast(d2t[v]) : static_cast(v); + if (targetIdx >= 0 && targetIdx < targetVocabSize) + { + rowProbs[targetIdx] = __expf(rowLogits[v] - sMaxLogit) / sExpSum; + } + } +} + +torch::Tensor computeDraftProbsSkipAllForDynamicTreeRejection(torch::Tensor const& draftLogits, int64_t batchSize, + SizeType32 const numDraftProbRows, SizeType32 const targetVocabSize, torch::optional const& d2t) +{ + auto const draftVocabSize = draftLogits.size(1); + bool const hasD2T = d2t.has_value() && d2t->defined(); + + auto draftLogitsFloat = draftLogits.contiguous().to(torch::kFloat32); + if (!hasD2T && draftVocabSize == targetVocabSize) + { + return computeSoftmaxForProbOp(draftLogitsFloat).reshape({batchSize, numDraftProbRows, targetVocabSize}); + } + + auto fullDraftProbs = torch::empty({draftLogitsFloat.size(0), targetVocabSize}, + torch::TensorOptions().dtype(torch::kFloat32).device(draftLogitsFloat.device())); + torch::Tensor d2tInt; + int32_t const* d2tPtr = nullptr; + if (hasD2T) + { + d2tInt = d2t->contiguous().to(torch::kInt32); + d2tPtr = d2tInt.data_ptr(); + } + + constexpr int32_t kBlockSize = 1024; + dim3 grid(draftLogitsFloat.size(0)); + dim3 block(kBlockSize); + auto stream = at::cuda::getCurrentCUDAStream(draftLogitsFloat.device().index()); + computeDraftProbsSkipAllKernel<<>>(draftLogitsFloat.data_ptr(), d2tPtr, + fullDraftProbs.data_ptr(), static_cast(draftLogitsFloat.size(0)), + static_cast(draftVocabSize), static_cast(targetVocabSize)); + sync_check_cuda_error(stream); + + return fullDraftProbs.reshape({batchSize, numDraftProbRows, targetVocabSize}); +} + // Fast path for top-K (and optional top-P) filtering using torch::topk instead of a // full vocab-size sort. kMax must be provided as a CPU integer (the caller computes it // via topK.max().item() on the Python side). When kMax == 0 or kMax >= vocabSize the @@ -347,15 +454,16 @@ torch::Tensor computeSoftmaxForProbOp(torch::Tensor logits) torch::Tensor applyTopKTopPForProbOp(torch::Tensor logits, torch::optional const& topK, torch::optional const& topP, int32_t kMax) { - bool const hasTopK = topK.has_value() && topK->defined(); - bool const hasTopP = isTopPEnabled(topP); + int64_t const vocabSize = logits.size(1); + bool const hasTopK = topK.has_value() && topK->defined() + && torch::logical_and(topK->gt(0), topK->lt(vocabSize)).any().item(); + bool const hasTopP = topP.has_value() && topP->defined() && topP->lt(1.0).any().item(); if (!hasTopK && !hasTopP) { return logits; } - int64_t const vocabSize = logits.size(1); torch::Tensor effectiveTopK; bool hasDisabledTopKRows = false; if (hasTopK) @@ -463,10 +571,11 @@ torch::Tensor computeProbsFromLogits(torch::Tensor const& logits, torch::Tensor auto scaledLogits = (skipTemperature ? logits : logits.div(safeTemperatures.unsqueeze(1))).contiguous().to(torch::kFloat32); - bool const hasTopK = topK.has_value() && topK->defined(); - bool const hasTopP = isTopPEnabled(topP); int64_t const vocabSize = scaledLogits.size(1); int64_t const nRows = scaledLogits.size(0); + bool const hasTopK = topK.has_value() && topK->defined() + && torch::logical_and(topK->gt(0), topK->lt(vocabSize)).any().item(); + bool const hasTopP = topP.has_value() && topP->defined() && topP->lt(1.0).any().item(); torch::Tensor maskedLogits; if (hasTopK && kMax > 0 && kMax < vocabSize) @@ -496,7 +605,7 @@ torch::Tensor computeProbsFromLogits(torch::Tensor const& logits, torch::Tensor torch::Tensor computeDraftProbsForDynamicTreeRejection(torch::Tensor const& draftLogits, torch::Tensor const& temperatures, SizeType32 const numDraftProbRows, torch::optional const& topK, torch::optional const& topP, SizeType32 const targetVocabSize, bool skipTemperature, - torch::optional const& d2t, SizeType32 const kMax) + torch::optional const& d2t, SizeType32 const kMax, bool skipAllSamplingParams) { TORCH_CHECK(draftLogits.is_cuda(), "draftLogits must be a CUDA tensor"); TORCH_CHECK(temperatures.is_cuda(), "temperatures must be a CUDA tensor"); @@ -531,6 +640,12 @@ torch::Tensor computeDraftProbsForDynamicTreeRejection(torch::Tensor const& draf TORCH_CHECK(d2t->size(0) >= draftVocabSize, "d2t size mismatch"); } + if (skipAllSamplingParams) + { + return computeDraftProbsSkipAllForDynamicTreeRejection( + draftLogits, batchSize, numDraftProbRows, targetVocabSize, d2t); + } + auto draftTemps = temperatures.repeat_interleave(numDraftProbRows); auto draftTopK = topK.has_value() && topK->defined() ? torch::optional(topK->repeat_interleave(numDraftProbRows)) @@ -568,7 +683,7 @@ torch::Tensor computeDraftProbsForDynamicTreeRejection(torch::Tensor const& draf std::tuple computeTargetProbsForDynamicTreeRejection( torch::Tensor const& targetLogits, torch::Tensor const& temperatures, SizeType32 const numDraftTokens, torch::optional const& topK, torch::optional const& topP, bool skipTemperature, - SizeType32 const kMax) + SizeType32 const kMax, bool skipAllSamplingParams) { TORCH_CHECK(targetLogits.is_cuda(), "targetLogits must be a CUDA tensor"); TORCH_CHECK(temperatures.is_cuda(), "temperatures must be a CUDA tensor"); @@ -597,6 +712,17 @@ std::tuple computeTargetProbsForDyn TORCH_CHECK(topP->size(0) == batchSize, "top_p size mismatch"); } + if (skipAllSamplingParams) + { + auto targetSupportIndices + = torch::empty({0}, torch::TensorOptions().dtype(torch::kInt32).device(targetLogits.device())); + auto targetSupportLengths + = torch::empty({0}, torch::TensorOptions().dtype(torch::kInt32).device(targetLogits.device())); + auto targetProbs = computeSoftmaxForProbOp(targetLogits); + return std::make_tuple(targetProbs.reshape({batchSize, numDraftTokens, targetVocabSize}), targetSupportIndices, + targetSupportLengths); + } + auto targetTemps = temperatures.repeat_interleave(numDraftTokens); auto targetTopK = topK.has_value() && topK->defined() ? torch::optional(topK->repeat_interleave(numDraftTokens)) @@ -1152,7 +1278,7 @@ void invokeVerifyDynamicTreeGreedy(int64_t* predicts, int64_t* acceptIndex, int6 __device__ __forceinline__ float curand_uniform_open_right(curandStatePhilox4_32_10_t& state) { float u = curand_uniform(&state); // (0, 1] - return u < 1.0f ? u : 0.0f; // [0, 1) + return u < 1.0f ? u : 0.0f; // [0, 1) } __device__ int64_t sampleFromDistribution(curandStatePhilox4_32_10_t& state, float const* probs, uint32_t vocabSize) @@ -1211,6 +1337,37 @@ __device__ int64_t sampleFromIndexedDistribution(curandStatePhilox4_32_10_t& sta return sampledTok; } +struct MinInt32Op +{ + __device__ __forceinline__ int32_t operator()(int32_t a, int32_t b) const + { + return a < b ? a : b; + } +}; + +struct MaxInt32Op +{ + __device__ __forceinline__ int32_t operator()(int32_t a, int32_t b) const + { + return a > b ? a : b; + } +}; + +struct MaxFloatOp +{ + __device__ __forceinline__ float operator()(float a, float b) const + { + return a > b ? a : b; + } +}; + +struct SoftmaxStats +{ + float maxVal; + float sumVal; + int32_t argmax; +}; + //! \param acceptIndex [out] accepted path as tree positions [bs, numSpecStep]. int64. //! \param acceptTokenNum [out] number of accepted draft tokens (excl. root) [bs]. int64. //! \param acceptToken [out] emitted token ids [bs, numSpecStep]. int64. @@ -1233,27 +1390,31 @@ __device__ int64_t sampleFromIndexedDistribution(curandStatePhilox4_32_10_t& sta //! \param vocabSize vocabulary size. //! \param seed [1] int64 on GPU. Philox RNG seed. //! \param offset [1] int64 on GPU. Philox RNG offset. -template +template __global__ void verifyDynamicTreeRejectionKernel(int64_t* acceptIndex, int64_t* acceptTokenNum, int64_t* acceptToken, - int64_t const* candidates, float const* draftProbs, float const* targetProbs, int32_t const* targetSupportIndices, + int64_t const* candidates, float const* draftInputs, float const* targetInputs, int32_t const* targetSupportIndices, int32_t const* targetSupportLengths, int32_t const* draftProbIndices, int32_t const* retrieveNextToken, int32_t const* retrieveNextSibling, bool const* treeValid, uint32_t batchSize, uint32_t numDraftProbRows, uint32_t maxTargetSupportSize, uint32_t numSpeculativeTokens, uint32_t numDraftTokens, uint32_t vocabSize, - int64_t const* seed, int64_t const* offset) + uint32_t draftVocabSize, int32_t const* targetToDraft, int64_t const* seed, int64_t const* offset, + float const* temperatures) { uint32_t bx = blockIdx.x; int32_t const tid = static_cast(threadIdx.x); + constexpr uint32_t kVecSize = 4; if (bx >= batchSize) { return; } using BlockReduce = cub::BlockReduce; + using BlockReduceInt = cub::BlockReduce; using BlockScan = cub::BlockScan; __shared__ union { typename BlockReduce::TempStorage reduce; + typename BlockReduceInt::TempStorage reduceInt; typename BlockScan::TempStorage scan; } tempStorage; @@ -1265,14 +1426,15 @@ __global__ void verifyDynamicTreeRejectionKernel(int64_t* acceptIndex, int64_t* __shared__ float sTargetMass; __shared__ float sPrefixBase; __shared__ int32_t sWinnerIndex; + __shared__ int32_t sLastValidIndex; __shared__ int64_t sSampledToken; + __shared__ float sLogitsMax; + __shared__ float sLogitsSum; + __shared__ int32_t sLogitsArgmax; - // Independent sibling testing: collect all accepted siblings at each depth, - // then select one weighted by target probability to remove ordering bias. - static constexpr int32_t kMaxSiblingsPerDepth = 64; - __shared__ int32_t sAccSibIdx[kMaxSiblingsPerDepth]; - __shared__ int64_t sAccSibTok[kMaxSiblingsPerDepth]; - __shared__ float sAccSibQProb[kMaxSiblingsPerDepth]; + // The first sibling that passes the rejection test at the current depth. + __shared__ int32_t sAccSibIdx; + __shared__ int64_t sAccSibTok; __shared__ int32_t sNumAccSiblings; uint32_t batchOffset = bx * numDraftTokens; @@ -1285,29 +1447,497 @@ __global__ void verifyDynamicTreeRejectionKernel(int64_t* acceptIndex, int64_t* } __syncthreads(); bool const hasCompactTargetSupport = targetSupportIndices != nullptr && targetSupportLengths != nullptr; + bool const isGreedyRequest + = USE_LOGITS && temperatures != nullptr && temperatures[bx] <= static_cast(kGreedyTempThreshold); + float const* draftProbs = draftInputs; + float const* targetProbs = targetInputs; + uint32_t const draftRowStride = USE_LOGITS ? draftVocabSize : vocabSize; + + auto canVectorizeLoad = [&](float const* probs, uint32_t rowSize) -> bool + { + constexpr uint32_t kLoadAlignmentBytes = kVecSize * sizeof(float); + return (rowSize % kVecSize == 0) && (reinterpret_cast(probs) % kLoadAlignmentBytes == 0); + }; + + auto loadProbVec = [&](float const* probs, uint32_t base, bool useVectorizedLoads, uint32_t rowSize) + { + flashinfer::vec_t probVec; + probVec.fill(0.0f); + if (useVectorizedLoads && base + kVecSize <= rowSize) + { + probVec.cast_load(probs + base); + } + else + { +#pragma unroll + for (uint32_t j = 0; j < kVecSize; ++j) + { + uint32_t const v = base + j; + if (v < rowSize) + { + probVec[j] = probs[v]; + } + } + } + return probVec; + }; + + auto computeLogitsStats = [&](float const* logitsRow, uint32_t rowSize) -> SoftmaxStats + { + float threadMax = -FLT_MAX; + for (uint32_t v = static_cast(tid); v < rowSize; v += BLOCK_SIZE) + { + threadMax = fmaxf(threadMax, logitsRow[v]); + } + + float const blockMax = BlockReduce(tempStorage.reduce).Reduce(threadMax, MaxFloatOp{}); + if (tid == 0) + { + sLogitsMax = blockMax; + } + __syncthreads(); + + float threadSum = 0.0f; + int32_t localArgmax = static_cast(rowSize); + for (uint32_t v = static_cast(tid); v < rowSize; v += BLOCK_SIZE) + { + float const logit = logitsRow[v]; + threadSum += __expf(logit - sLogitsMax); + if (logit == sLogitsMax && localArgmax == static_cast(rowSize)) + { + localArgmax = static_cast(v); + } + } + + float const blockSum = BlockReduce(tempStorage.reduce).Sum(threadSum); + __syncthreads(); + int32_t const blockArgmax = BlockReduceInt(tempStorage.reduceInt).Reduce(localArgmax, MinInt32Op{}); + if (tid == 0) + { + sLogitsSum = blockSum; + sLogitsArgmax = blockArgmax; + } + __syncthreads(); + + return SoftmaxStats{sLogitsMax, sLogitsSum, sLogitsArgmax}; + }; + + auto probFromLogits = [&](float const* logitsRow, uint32_t tokenId, uint32_t rowSize, SoftmaxStats stats) -> float + { + if (tokenId >= rowSize) + { + return 0.0f; + } + if (isGreedyRequest) + { + return tokenId == static_cast(stats.argmax) ? 1.0f : 0.0f; + } + constexpr float kFloatSoftmaxEpsilon = 1e-6f; + return __expf(logitsRow[tokenId] - stats.maxVal) / (stats.sumVal + kFloatSoftmaxEpsilon); + }; + + auto targetTokenToDraftToken = [&](uint32_t targetTokenId) -> int32_t + { + if (targetTokenId >= vocabSize) + { + return -1; + } + if (targetToDraft != nullptr) + { + return targetToDraft[targetTokenId]; + } + return targetTokenId < draftVocabSize ? static_cast(targetTokenId) : -1; + }; + + auto draftProbFromTargetToken + = [&](float const* draftLogitsRow, uint32_t targetTokenId, SoftmaxStats stats) -> float + { + int32_t const draftTokenId = targetTokenToDraftToken(targetTokenId); + if (draftTokenId < 0) + { + return 0.0f; + } + return probFromLogits(draftLogitsRow, static_cast(draftTokenId), draftVocabSize, stats); + }; + + auto sampleProbTile = [&](float(&value)[kVecSize], uint32_t base) -> bool + { + float const tileSum = BlockReduce(tempStorage.reduce).template Sum(value); + if (tid == 0) + { + sDiffSum = tileSum; + } + __syncthreads(); + + int32_t localLastValid = -1; +#pragma unroll + for (uint32_t j = 0; j < kVecSize; ++j) + { + uint32_t const v = base + j; + if (v < vocabSize && value[j] > 0.0f) + { + localLastValid = static_cast(v); + } + } + int32_t const blockLastValid = BlockReduceInt(tempStorage.reduceInt).Reduce(localLastValid, MaxInt32Op{}); + if (tid == 0 && blockLastValid >= 0) + { + sLastValidIndex = blockLastValid; + } + __syncthreads(); + + if (sPrefixBase + sDiffSum > sTargetMass) + { + float inclusive[kVecSize]; + BlockScan(tempStorage.scan).template InclusiveSum(value, inclusive); + __syncthreads(); + + int32_t localWinner = static_cast(vocabSize); +#pragma unroll + for (uint32_t j = 0; j < kVecSize; ++j) + { + uint32_t const v = base + j; + if (v < vocabSize && value[j] > 0.0f && sPrefixBase + inclusive[j] > sTargetMass) + { + localWinner = static_cast(v); + break; + } + } + + int32_t const blockWinner = BlockReduceInt(tempStorage.reduceInt).Reduce(localWinner, MinInt32Op{}); + if (tid == 0) + { + sWinnerIndex = blockWinner; + sSampledToken = blockWinner < static_cast(vocabSize) ? static_cast(blockWinner) + : static_cast(vocabSize) - 1; + } + __syncthreads(); + return true; + } + + if (tid == 0) + { + sPrefixBase += sDiffSum; + } + __syncthreads(); + return false; + }; + + auto sampleTargetFullVocab = [&](float const* tProbs) + { + bool const useVectorizedLoads = canVectorizeLoad(tProbs, vocabSize); + uint32_t const numIters = (vocabSize + BLOCK_SIZE * kVecSize - 1) / (BLOCK_SIZE * kVecSize); + if (tid == 0) + { + sPrefixBase = 0.0f; + sWinnerIndex = static_cast(vocabSize); + sLastValidIndex = -1; + sSampledToken = static_cast(vocabSize) - 1; + sTargetMass = curand_uniform_open_right(state); + } + __syncthreads(); + +#pragma unroll 2 + for (uint32_t i = 0; i < numIters; ++i) + { + uint32_t const base = (i * BLOCK_SIZE + static_cast(tid)) * kVecSize; + auto const qVec = loadProbVec(tProbs, base, useVectorizedLoads, vocabSize); + float value[kVecSize]; +#pragma unroll + for (uint32_t j = 0; j < kVecSize; ++j) + { + value[j] = qVec[j]; + } + + if (sampleProbTile(value, base)) + { + break; + } + } + + if (tid == 0 && sWinnerIndex >= static_cast(vocabSize) && sLastValidIndex >= 0) + { + sSampledToken = static_cast(sLastValidIndex); + } + __syncthreads(); + }; + + auto sampleResidualFullVocab = [&](float const* tProbs, float const* dProbs) + { + bool const useVectorizedTargetLoads = canVectorizeLoad(tProbs, vocabSize); + bool const useVectorizedDraftLoads = canVectorizeLoad(dProbs, vocabSize); + uint32_t const numIters = (vocabSize + BLOCK_SIZE * kVecSize - 1) / (BLOCK_SIZE * kVecSize); + float totalSum = 0.0f; +#pragma unroll 2 + for (uint32_t i = 0; i < numIters; ++i) + { + uint32_t const base = (i * BLOCK_SIZE + static_cast(tid)) * kVecSize; + auto const qVec = loadProbVec(tProbs, base, useVectorizedTargetLoads, vocabSize); + auto const pVec = loadProbVec(dProbs, base, useVectorizedDraftLoads, vocabSize); + float value[kVecSize]; +#pragma unroll + for (uint32_t j = 0; j < kVecSize; ++j) + { + value[j] = fmaxf(qVec[j] - pVec[j], 0.0f); + } + totalSum += BlockReduce(tempStorage.reduce).template Sum(value); + __syncthreads(); + } + + if (tid == 0) + { + sDiffSum = totalSum; + sPrefixBase = 0.0f; + sWinnerIndex = static_cast(vocabSize); + sLastValidIndex = -1; + sSampledToken = static_cast(vocabSize) - 1; + sTargetMass = totalSum > 1e-10f ? curand_uniform_open_right(state) * totalSum : 0.0f; + if (totalSum <= 1e-10f) + { + sSampledToken = sampleFromDistribution(state, tProbs, vocabSize); + } + } + __syncthreads(); + + if (sDiffSum <= 1e-10f) + { + return; + } + +#pragma unroll 2 + for (uint32_t i = 0; i < numIters; ++i) + { + uint32_t const base = (i * BLOCK_SIZE + static_cast(tid)) * kVecSize; + auto const qVec = loadProbVec(tProbs, base, useVectorizedTargetLoads, vocabSize); + auto const pVec = loadProbVec(dProbs, base, useVectorizedDraftLoads, vocabSize); + float value[kVecSize]; +#pragma unroll + for (uint32_t j = 0; j < kVecSize; ++j) + { + value[j] = fmaxf(qVec[j] - pVec[j], 0.0f); + } + + if (sampleProbTile(value, base)) + { + break; + } + } + + if (tid == 0 && sWinnerIndex >= static_cast(vocabSize) && sLastValidIndex >= 0) + { + sSampledToken = static_cast(sLastValidIndex); + } + __syncthreads(); + }; + + auto sampleTargetLogitsFullVocabWithStats = [&](float const* tLogits, SoftmaxStats targetStats) + { + if (tid == 0) + { + constexpr float kFloatSoftmaxEpsilon = 1e-6f; + sPrefixBase = 0.0f; + sWinnerIndex = static_cast(vocabSize); + sLastValidIndex = -1; + sSampledToken = static_cast(vocabSize) - 1; + sTargetMass = isGreedyRequest + ? 0.0f + : curand_uniform_open_right(state) * (targetStats.sumVal + kFloatSoftmaxEpsilon); + if (isGreedyRequest) + { + sSampledToken = static_cast(targetStats.argmax); + } + } + __syncthreads(); + + if (isGreedyRequest) + { + return; + } + + bool const useVectorizedLoads = canVectorizeLoad(tLogits, vocabSize); + uint32_t const numIters = (vocabSize + BLOCK_SIZE * kVecSize - 1) / (BLOCK_SIZE * kVecSize); +#pragma unroll 2 + for (uint32_t i = 0; i < numIters; ++i) + { + uint32_t const base = (i * BLOCK_SIZE + static_cast(tid)) * kVecSize; + auto const qVec = loadProbVec(tLogits, base, useVectorizedLoads, vocabSize); + float value[kVecSize]; +#pragma unroll + for (uint32_t j = 0; j < kVecSize; ++j) + { + uint32_t const v = base + j; + value[j] = v < vocabSize ? __expf(qVec[j] - targetStats.maxVal) : 0.0f; + } + + if (sampleProbTile(value, base)) + { + break; + } + } + + if (tid == 0 && sWinnerIndex >= static_cast(vocabSize) && sLastValidIndex >= 0) + { + sSampledToken = static_cast(sLastValidIndex); + } + __syncthreads(); + }; + + auto sampleTargetLogitsFullVocab = [&](float const* tLogits) + { + auto const targetStats = computeLogitsStats(tLogits, vocabSize); + sampleTargetLogitsFullVocabWithStats(tLogits, targetStats); + }; + + auto sampleResidualLogitsFullVocab + = [&](float const* tLogits, float const* dLogits, SoftmaxStats targetStats, SoftmaxStats draftStats) + { + if (isGreedyRequest) + { + if (tid == 0) + { + sSampledToken = static_cast(targetStats.argmax); + } + __syncthreads(); + return; + } + + bool const useVectorizedTargetLoads = canVectorizeLoad(tLogits, vocabSize); + bool const useVectorizedDraftLoads = targetToDraft == nullptr && canVectorizeLoad(dLogits, draftVocabSize); + uint32_t const numIters = (vocabSize + BLOCK_SIZE * kVecSize - 1) / (BLOCK_SIZE * kVecSize); + constexpr float kFloatSoftmaxEpsilon = 1e-6f; + float totalSum = 0.0f; +#pragma unroll 2 + for (uint32_t i = 0; i < numIters; ++i) + { + uint32_t const base = (i * BLOCK_SIZE + static_cast(tid)) * kVecSize; + auto const qVec = loadProbVec(tLogits, base, useVectorizedTargetLoads, vocabSize); + flashinfer::vec_t pVec; + pVec.fill(0.0f); + if (targetToDraft == nullptr) + { + pVec = loadProbVec(dLogits, base, useVectorizedDraftLoads, draftVocabSize); + } + float value[kVecSize]; +#pragma unroll + for (uint32_t j = 0; j < kVecSize; ++j) + { + uint32_t const v = base + j; + if (v < vocabSize) + { + float const q = __expf(qVec[j] - targetStats.maxVal) / (targetStats.sumVal + kFloatSoftmaxEpsilon); + int32_t const draftTokenId = targetTokenToDraftToken(v); + float p = 0.0f; + if (draftTokenId >= 0 && draftTokenId < static_cast(draftVocabSize)) + { + float const draftLogit = targetToDraft == nullptr ? pVec[j] : dLogits[draftTokenId]; + p = __expf(draftLogit - draftStats.maxVal) / (draftStats.sumVal + kFloatSoftmaxEpsilon); + } + value[j] = fmaxf(q - p, 0.0f); + } + else + { + value[j] = 0.0f; + } + } + totalSum += BlockReduce(tempStorage.reduce).template Sum(value); + __syncthreads(); + } + + if (tid == 0) + { + sDiffSum = totalSum; + sPrefixBase = 0.0f; + sWinnerIndex = static_cast(vocabSize); + sLastValidIndex = -1; + sSampledToken = static_cast(vocabSize) - 1; + sTargetMass = totalSum > 1e-10f ? curand_uniform_open_right(state) * totalSum : 0.0f; + } + __syncthreads(); + + if (sDiffSum <= 1e-10f) + { + sampleTargetLogitsFullVocabWithStats(tLogits, targetStats); + return; + } + +#pragma unroll 2 + for (uint32_t i = 0; i < numIters; ++i) + { + uint32_t const base = (i * BLOCK_SIZE + static_cast(tid)) * kVecSize; + auto const qVec = loadProbVec(tLogits, base, useVectorizedTargetLoads, vocabSize); + flashinfer::vec_t pVec; + pVec.fill(0.0f); + if (targetToDraft == nullptr) + { + pVec = loadProbVec(dLogits, base, useVectorizedDraftLoads, draftVocabSize); + } + float value[kVecSize]; +#pragma unroll + for (uint32_t j = 0; j < kVecSize; ++j) + { + uint32_t const v = base + j; + if (v < vocabSize) + { + float const q = __expf(qVec[j] - targetStats.maxVal) / (targetStats.sumVal + kFloatSoftmaxEpsilon); + int32_t const draftTokenId = targetTokenToDraftToken(v); + float p = 0.0f; + if (draftTokenId >= 0 && draftTokenId < static_cast(draftVocabSize)) + { + float const draftLogit = targetToDraft == nullptr ? pVec[j] : dLogits[draftTokenId]; + p = __expf(draftLogit - draftStats.maxVal) / (draftStats.sumVal + kFloatSoftmaxEpsilon); + } + value[j] = fmaxf(q - p, 0.0f); + } + else + { + value[j] = 0.0f; + } + } + + if (sampleProbTile(value, base)) + { + break; + } + } + + if (tid == 0 && sWinnerIndex >= static_cast(vocabSize) && sLastValidIndex >= 0) + { + sSampledToken = static_cast(sLastValidIndex); + } + __syncthreads(); + }; // First-gen or dummy request: no valid tree exists yet. Sample directly // from the target distribution at the root and skip tree traversal. if (treeValid != nullptr && !treeValid[bx]) { - if (tid == 0) + float const* tProbs = targetProbs + static_cast(bx) * numDraftTokens * vocabSize; + if (hasCompactTargetSupport) { - float const* tProbs = targetProbs + static_cast(bx) * numDraftTokens * vocabSize; - int64_t sampledToken; - if (hasCompactTargetSupport) + if (tid == 0) { uint32_t const supportOffset = static_cast(bx) * numDraftTokens * maxTargetSupportSize; uint32_t const supportSize = static_cast(targetSupportLengths[batchOffset]); - sampledToken = sampleFromIndexedDistribution( + sSampledToken = sampleFromIndexedDistribution( state, tProbs, targetSupportIndices + supportOffset, supportSize, vocabSize); } + } + else + { + if constexpr (USE_LOGITS) + { + sampleTargetLogitsFullVocab(tProbs); + } else { - sampledToken = sampleFromDistribution(state, tProbs, vocabSize); + sampleTargetFullVocab(tProbs); } + } + if (tid == 0) + { acceptIndex[bx * numSpeculativeTokens] = 0; acceptTokenNum[bx] = 0; - acceptToken[bx * numSpeculativeTokens] = sampledToken; + acceptToken[bx * numSpeculativeTokens] = sSampledToken; } return; } @@ -1359,48 +1989,71 @@ __global__ void verifyDynamicTreeRejectionKernel(int64_t* acceptIndex, int64_t* // Emit bonus token from the target distribution at the last accepted position. if (sFirstChild == -1) { - if (tid == 0) + float const* tProbs + = targetProbs + (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * vocabSize; + if (hasCompactTargetSupport) { - float const* tProbs = targetProbs - + (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * vocabSize; - if (hasCompactTargetSupport) + if (tid == 0) { uint32_t const supportOffset = (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * maxTargetSupportSize; uint32_t const supportSize = static_cast(targetSupportLengths[batchOffset + sLastAcceptedLocalIdx]); - acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] = sampleFromIndexedDistribution( + sSampledToken = sampleFromIndexedDistribution( state, tProbs, targetSupportIndices + supportOffset, supportSize, vocabSize); } + } + else + { + if constexpr (USE_LOGITS) + { + sampleTargetLogitsFullVocab(tProbs); + } else { - acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] - = sampleFromDistribution(state, tProbs, vocabSize); + sampleTargetFullVocab(tProbs); } + } + if (tid == 0) + { + acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] = sSampledToken; sHasTerminalToken = true; } __syncthreads(); break; } - // ---- Phase 1: Independently test ALL siblings with Bernoulli rejection ---- - // Unlike sequential testing which accepts the first passing sibling - // (creating ordering bias), we test every sibling and collect all that - // pass. This removes the systematic preference for earlier siblings in - // the linked list and produces output closer to the target distribution. + // Test siblings in linked-list order. Once a sibling passes the + // Bernoulli rejection test, accept it immediately and skip the rest. + int32_t const firstDraftProbRow = draftProbIndices[batchOffset + sFirstChild]; + float const* siblingTargetRow + = targetProbs + (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * vocabSize; + float const* siblingDraftRow + = draftProbs + (static_cast(bx) * numDraftProbRows + firstDraftProbRow) * draftRowStride; + SoftmaxStats siblingTargetStats{}; + SoftmaxStats siblingDraftStats{}; + if constexpr (USE_LOGITS) + { + siblingTargetStats = computeLogitsStats(siblingTargetRow, vocabSize); + siblingDraftStats = computeLogitsStats(siblingDraftRow, draftVocabSize); + } + if (tid == 0) { sNumAccSiblings = 0; int32_t childIdx = sFirstChild; - while (childIdx != -1 && sNumAccSiblings < kMaxSiblingsPerDepth) + while (childIdx != -1) { int64_t const draftTokenId = candidates[batchOffset + childIdx]; int32_t const draftProbRow = draftProbIndices[batchOffset + childIdx]; - float const pDraft - = draftProbs[(static_cast(bx) * numDraftProbRows + draftProbRow) * vocabSize + uint32_t const tokenId = static_cast(draftTokenId); + float const pDraft = USE_LOGITS + ? draftProbFromTargetToken(siblingDraftRow, tokenId, siblingDraftStats) + : draftProbs[(static_cast(bx) * numDraftProbRows + draftProbRow) * vocabSize + draftTokenId]; - float const pTarget - = targetProbs[(static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * vocabSize + float const pTarget = USE_LOGITS + ? probFromLogits(siblingTargetRow, tokenId, vocabSize, siblingTargetStats) + : targetProbs[(static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * vocabSize + draftTokenId]; float const acceptProb = fminf(1.0f, pTarget / (pDraft + 1e-10f)); @@ -1408,47 +2061,23 @@ __global__ void verifyDynamicTreeRejectionKernel(int64_t* acceptIndex, int64_t* if (u < acceptProb) { - int32_t const idx = sNumAccSiblings++; - sAccSibIdx[idx] = childIdx; - sAccSibTok[idx] = draftTokenId; - sAccSibQProb[idx] = pTarget; + sAccSibIdx = childIdx; + sAccSibTok = draftTokenId; + sNumAccSiblings = 1; + break; } childIdx = retrieveNextSibling[batchOffset + childIdx]; } } __syncthreads(); - // ---- Phase 2: Select among accepted siblings or emit correction ---- + // Select the first accepted sibling or emit correction when all siblings reject. if (sNumAccSiblings > 0) { - // At least one sibling accepted. Select one weighted by target - // probability q(ci) so that the choice follows the target distribution. if (tid == 0) { - int32_t selectedIdx = 0; - if (sNumAccSiblings > 1) - { - float totalQ = 0.0f; - for (int32_t i = 0; i < sNumAccSiblings; ++i) - { - totalQ += sAccSibQProb[i]; - } - float const r = curand_uniform_open_right(state) * totalQ; - float cumsum = 0.0f; - selectedIdx = sNumAccSiblings - 1; // fallback - for (int32_t i = 0; i < sNumAccSiblings; ++i) - { - cumsum += sAccSibQProb[i]; - if (r < cumsum) - { - selectedIdx = i; - break; - } - } - } - - int32_t const childIdx = sAccSibIdx[selectedIdx]; - acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] = sAccSibTok[selectedIdx]; + int32_t const childIdx = sAccSibIdx; + acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] = sAccSibTok; ++sNumAcceptedTokens; acceptIndex[bx * numSpeculativeTokens + sNumAcceptedTokens] = childIdx; sLastAcceptedLocalIdx = childIdx; @@ -1463,7 +2092,7 @@ __global__ void verifyDynamicTreeRejectionKernel(int64_t* acceptIndex, int64_t* float const* tProbs = targetProbs + (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * vocabSize; float const* dProbs - = draftProbs + (static_cast(bx) * numDraftProbRows + draftProbRow) * vocabSize; + = draftProbs + (static_cast(bx) * numDraftProbRows + draftProbRow) * draftRowStride; int32_t const* tProbIndices = nullptr; uint32_t targetSupportSize = vocabSize; if (hasCompactTargetSupport) @@ -1521,73 +2150,13 @@ __global__ void verifyDynamicTreeRejectionKernel(int64_t* acceptIndex, int64_t* } else { - float threadDiffSum = 0.0f; - for (uint32_t v = static_cast(tid); v < targetSupportSize; v += BLOCK_SIZE) + if constexpr (USE_LOGITS) { - float const diff = tProbs[v] - dProbs[v]; - if (diff > 0.0f) - { - threadDiffSum += diff; - } + sampleResidualLogitsFullVocab(tProbs, dProbs, siblingTargetStats, siblingDraftStats); } - - float const diffSum = BlockReduce(tempStorage.reduce).Sum(threadDiffSum); - if (tid == 0) + else { - sDiffSum = diffSum; - sPrefixBase = 0.0f; - sWinnerIndex = static_cast(targetSupportSize); - sSampledToken = static_cast(vocabSize) - 1; - if (diffSum > 1e-10f) - { - sTargetMass = curand_uniform_open_right(state) * diffSum; - } - else - { - sSampledToken = sampleFromDistribution(state, tProbs, vocabSize); - } - } - __syncthreads(); - - if (sDiffSum > 1e-10f) - { - for (uint32_t tileStart = 0; tileStart < targetSupportSize; tileStart += BLOCK_SIZE) - { - float value = 0.0f; - uint32_t const v = tileStart + static_cast(tid); - if (v < targetSupportSize) - { - float const diff = tProbs[v] - dProbs[v]; - value = diff > 0.0f ? diff : 0.0f; - } - - float inclusive = 0.0f; - float tileSum = 0.0f; - BlockScan(tempStorage.scan).InclusiveSum(value, inclusive, tileSum); - float const threshold = sTargetMass; - float const prefixBase = sPrefixBase; - - if (value > 0.0f && prefixBase + inclusive >= threshold) - { - atomicMin(&sWinnerIndex, static_cast(v)); - } - __syncthreads(); - - if (tid == 0) - { - if (sWinnerIndex < static_cast(targetSupportSize)) - { - sSampledToken = static_cast(sWinnerIndex); - } - sPrefixBase += tileSum; - } - __syncthreads(); - - if (sWinnerIndex < static_cast(targetSupportSize)) - { - break; - } - } + sampleResidualFullVocab(tProbs, dProbs); } if (tid == 0) @@ -1606,25 +2175,35 @@ __global__ void verifyDynamicTreeRejectionKernel(int64_t* acceptIndex, int64_t* { // Reached max speculative depth while continuing to accept the draft path. // Emit the final bonus token from the last accepted position. - if (tid == 0) + float const* tProbs + = targetProbs + (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * vocabSize; + if (hasCompactTargetSupport) { - float const* tProbs - = targetProbs + (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * vocabSize; - if (hasCompactTargetSupport) + if (tid == 0) { uint32_t const supportOffset = (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * maxTargetSupportSize; uint32_t const supportSize = static_cast(targetSupportLengths[batchOffset + sLastAcceptedLocalIdx]); - acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] = sampleFromIndexedDistribution( + sSampledToken = sampleFromIndexedDistribution( state, tProbs, targetSupportIndices + supportOffset, supportSize, vocabSize); } + } + else + { + if constexpr (USE_LOGITS) + { + sampleTargetLogitsFullVocab(tProbs); + } else { - acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] - = sampleFromDistribution(state, tProbs, vocabSize); + sampleTargetFullVocab(tProbs); } } + if (tid == 0) + { + acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] = sSampledToken; + } } if (tid == 0) @@ -1640,14 +2219,15 @@ void invokeVerifyDynamicTreeRejection(int64_t* acceptIndex, int64_t* acceptToken SizeType32 maxTargetSupportSize, SizeType32 numDraftTokens, SizeType32 numSpecStep, SizeType32 vocabSize, int64_t const* seed, int64_t const* offset, cudaStream_t stream) { - constexpr int32_t kVerifyDynamicTreeRejectionBlockSize = 128; + constexpr int32_t kVerifyDynamicTreeRejectionBlockSize = 1024; dim3 grid(batchSize); dim3 block(kVerifyDynamicTreeRejectionBlockSize); - verifyDynamicTreeRejectionKernel<<>>(acceptIndex, - acceptTokenNum, acceptToken, candidates, draftProbs, targetProbs, targetSupportIndices, targetSupportLengths, - draftProbIndices, retrieveNextToken, retrieveNextSibling, treeValid, batchSize, numDraftProbRows, - maxTargetSupportSize, numSpecStep, numDraftTokens, vocabSize, seed, offset); + verifyDynamicTreeRejectionKernel<<>>( + acceptIndex, acceptTokenNum, acceptToken, candidates, draftProbs, targetProbs, targetSupportIndices, + targetSupportLengths, draftProbIndices, retrieveNextToken, retrieveNextSibling, treeValid, batchSize, + numDraftProbRows, maxTargetSupportSize, numSpecStep, numDraftTokens, vocabSize, vocabSize, + /*targetToDraft=*/nullptr, seed, offset, nullptr); sync_check_cuda_error(stream); } diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h index 48e788e91e3e..b7adbdc2f573 100644 --- a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h +++ b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h @@ -156,7 +156,7 @@ void invokeVerifyDynamicTreeRejection(int64_t* acceptIndex, int64_t* acceptToken torch::Tensor computeDraftProbsForDynamicTreeRejection(torch::Tensor const& draftLogits, torch::Tensor const& temperatures, runtime::SizeType32 numDraftProbRows, torch::optional const& topK, torch::optional const& topP, runtime::SizeType32 targetVocabSize, bool skipTemperature, - torch::optional const& d2t, runtime::SizeType32 kMax = 0); + torch::optional const& d2t, runtime::SizeType32 kMax = 0, bool skipAllSamplingParams = false); //! \brief Compute target probabilities for dynamic-tree rejection sampling from logits. //! \param targetLogits [batchSize * numDraftTokens, targetVocabSize], on GPU. @@ -175,7 +175,7 @@ torch::Tensor computeDraftProbsForDynamicTreeRejection(torch::Tensor const& draf std::tuple computeTargetProbsForDynamicTreeRejection( torch::Tensor const& targetLogits, torch::Tensor const& temperatures, runtime::SizeType32 numDraftTokens, torch::optional const& topK, torch::optional const& topP, bool skipTemperature, - runtime::SizeType32 kMax = 0); + runtime::SizeType32 kMax = 0, bool skipAllSamplingParams = false); } // namespace kernels::speculative_decoding diff --git a/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp b/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp index 2790f06cf0d6..7f0bd16b93da 100644 --- a/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp +++ b/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp @@ -35,10 +35,10 @@ void invokeBuildDraftProbIndices(int64_t const* topkScoreIndices, int32_t* draft th::Tensor computeDraftProbsForDynamicTreeRejection(th::Tensor const& draftLogits, th::Tensor const& temperatures, runtime::SizeType32 numDraftProbRows, th::optional const& topK, th::optional const& topP, runtime::SizeType32 targetVocabSize, bool skipTemperature, th::optional const& d2t, - runtime::SizeType32 kMax); + runtime::SizeType32 kMax, bool skipAllSamplingParams); std::tuple computeTargetProbsForDynamicTreeRejection(th::Tensor const& targetLogits, th::Tensor const& temperatures, runtime::SizeType32 numDraftTokens, th::optional const& topK, - th::optional const& topP, bool skipTemperature, runtime::SizeType32 kMax); + th::optional const& topP, bool skipTemperature, runtime::SizeType32 kMax, bool skipAllSamplingParams); } // namespace kernels::speculative_decoding namespace torch_ext @@ -331,18 +331,18 @@ void verify_dynamic_tree_rejection_out_op(th::Tensor& candidates, th::Tensor& dr th::Tensor compute_draft_probs_for_dynamic_tree_rejection_op(th::Tensor draftLogits, th::Tensor temperatures, int64_t numDraftProbRows, int64_t targetVocabSize, th::optional topK, th::optional topP, - bool skipTemperature, th::optional d2t, int64_t topKMax) + bool skipTemperature, th::optional d2t, int64_t topKMax, bool skipAllSamplingParams) { - return tk::computeDraftProbsForDynamicTreeRejection( - draftLogits, temperatures, numDraftProbRows, topK, topP, targetVocabSize, skipTemperature, d2t, topKMax); + return tk::computeDraftProbsForDynamicTreeRejection(draftLogits, temperatures, numDraftProbRows, topK, topP, + targetVocabSize, skipTemperature, d2t, topKMax, skipAllSamplingParams); } std::tuple compute_target_probs_for_dynamic_tree_rejection_op( th::Tensor targetLogits, th::Tensor temperatures, int64_t numDraftTokens, th::optional topK, - th::optional topP, bool skipTemperature, int64_t topKMax) + th::optional topP, bool skipTemperature, int64_t topKMax, bool skipAllSamplingParams) { return tk::computeTargetProbsForDynamicTreeRejection( - targetLogits, temperatures, numDraftTokens, topK, topP, skipTemperature, topKMax); + targetLogits, temperatures, numDraftTokens, topK, topP, skipTemperature, topKMax, skipAllSamplingParams); } th::Tensor compute_probs_from_logits_op(th::Tensor logits, th::Tensor temperatures, th::optional topK, @@ -452,7 +452,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) "compute_draft_probs_for_dynamic_tree_rejection_op(" "Tensor draftLogits, Tensor temperatures, int numDraftProbRows, int targetVocabSize, " "Tensor? top_k=None, Tensor? top_p=None, bool skip_temperature=False, Tensor? d2t=None, " - "int top_k_max=0) -> Tensor"); + "int top_k_max=0, bool skip_all_sampling_params=False) -> Tensor"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) @@ -466,8 +466,8 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) m.def( "compute_target_probs_for_dynamic_tree_rejection_op(" "Tensor targetLogits, Tensor temperatures, int numDraftTokens, " - "Tensor? top_k=None, Tensor? top_p=None, bool skip_temperature=False, int top_k_max=0) -> (Tensor, Tensor, " - "Tensor)"); + "Tensor? top_k=None, Tensor? top_p=None, bool skip_temperature=False, int top_k_max=0, " + "bool skip_all_sampling_params=False) -> (Tensor, Tensor, Tensor)"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) diff --git a/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py b/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py index d0126d534f22..77edcb7c1cb1 100644 --- a/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py +++ b/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py @@ -257,6 +257,7 @@ def verify_dynamic_tree_rejection_from_logits_out( seed: int | torch.Tensor = 0, offset: int | torch.Tensor = 0, d2t: torch.Tensor | None = None, + skip_all_sampling_params: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Tree-aware rejection sampling from logits (three CUDA ops). @@ -303,6 +304,7 @@ def verify_dynamic_tree_rejection_from_logits_out( skip_temperature, d2t=d2t, top_k_max=top_k_max, + skip_all_sampling_params=skip_all_sampling_params, ) ( @@ -317,6 +319,7 @@ def verify_dynamic_tree_rejection_from_logits_out( top_p, skip_temperature, top_k_max=top_k_max, + skip_all_sampling_params=skip_all_sampling_params, ) torch.ops.trtllm.verify_dynamic_tree_rejection_out_op( diff --git a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py index 0b61a70dfc11..2644a1b5dbc1 100644 --- a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py @@ -814,6 +814,13 @@ def _sample_and_accept_dynamic_tree( if not skip_top_p and spec_metadata.request_top_ps is not None: top_ps = spec_metadata.request_top_ps[gen_slice] + skip_all_sampling_params = ( + skip_temperature + and skip_top_k + and skip_top_p + and not getattr(spec_metadata, "has_greedy_requests", False) + ) + # Lazily initialize seed/offset tensors on correct device if self.seed is None: self.seed = torch.tensor([0], dtype=torch.int64, device=device) @@ -839,6 +846,7 @@ def _sample_and_accept_dynamic_tree( seed=self.seed, offset=self.offset, d2t=self._d2t, + skip_all_sampling_params=skip_all_sampling_params, ) ) diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 3ebd41d27e85..d623ce0eeb61 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -448,6 +448,7 @@ class SpecMetadata: skip_temperature: bool = False skip_top_k: bool = False skip_top_p: bool = False + has_greedy_requests: bool = False # Sampling parameters indexed per request. request_temperatures: Optional[torch.Tensor] = None request_top_ks: Optional[torch.Tensor] = None @@ -534,6 +535,7 @@ def populate_sampling_params_for_one_model( request_top_ps = [] top_k_enabled = False top_p_enabled = False + has_greedy_requests = False temperature_enabled = False # Need to use a very small value for temperature when disabled to avoid division by 0 @@ -551,7 +553,7 @@ def _normalize_request_sampling_params( temperature: Optional[float], top_k: Optional[int], top_p: Optional[float], - ) -> tuple[float, int, float, bool, bool, bool]: + ) -> tuple[float, int, float, bool, bool, bool, bool]: """Convert request sampling params into normalized per-request scalars.""" is_greedy = SamplingParams.params_imply_greedy_decoding( temperature=temperature, @@ -578,6 +580,7 @@ def _normalize_request_sampling_params( use_temperature, use_top_k, use_top_p, + is_greedy, ) for request in requests: @@ -596,6 +599,7 @@ def _normalize_request_sampling_params( use_temperature, use_top_k, use_top_p, + is_greedy, ) = _normalize_request_sampling_params( temperature=temp_val, top_k=tk_val, @@ -605,6 +609,7 @@ 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 request_temperatures.append(temp_val) request_top_ks.append(tk_val) @@ -667,6 +672,7 @@ def _normalize_request_sampling_params( 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 class SpecWorkerBase(nn.Module, ABC): @@ -1024,7 +1030,8 @@ def _sample_and_accept_draft_tokens_rejection( temperatures = spec_metadata.temperatures[:num_target_tokens] top_ks = spec_metadata.top_ks[:num_target_tokens] - top_ps = None if spec_metadata.skip_top_p else spec_metadata.top_ps[:num_target_tokens] + top_ps = None if spec_metadata.skip_top_p else spec_metadata.top_ps[: + num_target_tokens] target_probs_flat = compute_probs_from_logits(logits.clone(), temperatures, top_ks, @@ -1123,9 +1130,10 @@ def _compute_and_store_draft_probs( draft_top_ks = spec_metadata.request_top_ks[:batch_size].repeat_interleave( draft_tokens_per_request ) if spec_metadata.request_top_ks is not None else None - draft_top_ps = (spec_metadata.request_top_ps[:batch_size].repeat_interleave( - draft_tokens_per_request) if not spec_metadata.skip_top_p - and spec_metadata.request_top_ps is not None else None) + draft_top_ps = ( + spec_metadata.request_top_ps[:batch_size].repeat_interleave( + draft_tokens_per_request) if not spec_metadata.skip_top_p + and spec_metadata.request_top_ps is not None else None) else: draft_temps = torch.ones(num_draft_tokens, device=device) draft_top_ks = None From cf9e9f89958add10e32148b0e54ed5608c841b94 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 6 May 2026 20:35:00 -0700 Subject: [PATCH 12/17] [None][fix] remove stale disaggregated B200 test entries Signed-off-by: ZhaoyangWang --- tests/integration/test_lists/test-db/l0_dgx_b200.yml | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index dea6041d46e0..d15b1c7a782f 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -293,9 +293,6 @@ l0_dgx_b200: - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-two_model-no_overlap_scheduler] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3_rejection_dynamic_tree_smoke[no_dynamic_tree] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3_rejection_dynamic_tree_smoke[dynamic_tree] - - disaggregated/test_disaggregated.py::test_disaggregated_benchmark_on_diff_backends[llama-v3-8b-hf] - - disaggregated/test_disaggregated.py::test_disaggregated_benchmark_on_diff_backends[DeepSeek-V3-Lite-bf16] - - disaggregated/test_disaggregated.py::test_disaggregated_benchmark_on_diff_backends[llama-3.1-8b-instruct-hf-fp8] - unittest/_torch/multi_gpu_modeling -k "deepseek" - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[pp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTEDSL-mtp_nextn=2-ep4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] From 0310fea44c788f2831257cd7be747ad3193750b2 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Thu, 7 May 2026 00:00:03 -0700 Subject: [PATCH 13/17] [None][test] add single-GPU Eagle3 rejection QA smoke test Signed-off-by: ZhaoyangWang --- tests/integration/defs/accuracy/test_llm_api_pytorch.py | 8 ++++---- tests/integration/test_lists/test-db/l0_dgx_b200.yml | 2 -- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 92596151314e..f518ca05bd15 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -187,7 +187,6 @@ def test_dummy_load_format(self): task = MMLU(self.MODEL_NAME) task.evaluate(llm, is_integration_test=True) - @pytest.mark.skip_less_device(4) @pytest.mark.parametrize("use_dynamic_tree", [False, True], ids=["no_dynamic_tree", "dynamic_tree"]) def test_eagle3_rejection_dynamic_tree_smoke(self, use_dynamic_tree, @@ -203,17 +202,18 @@ def test_eagle3_rejection_dynamic_tree_smoke(self, use_dynamic_tree, allow_advanced_sampling=True, use_rejection_sampling=True, ) + max_batch_size = 1 if use_dynamic_tree: spec_config_kwargs.update( use_dynamic_tree=True, dynamic_tree_max_topK=4, max_total_draft_tokens=16, - max_batch_size=4, + max_batch_size=max_batch_size, ) llm = LLM( self.MODEL_PATH, - tensor_parallel_size=4, + tensor_parallel_size=1, pipeline_parallel_size=1, attn_backend="TRTLLM", disable_overlap_scheduler=True, @@ -221,7 +221,7 @@ def test_eagle3_rejection_dynamic_tree_smoke(self, use_dynamic_tree, kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.4, dtype="auto"), max_seq_len=4096, - max_batch_size=4, + max_batch_size=max_batch_size, speculative_config=Eagle3DecodingConfig(**spec_config_kwargs), ) diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index d15b1c7a782f..ab3f8c18d5d8 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -291,8 +291,6 @@ l0_dgx_b200: - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-two_model-overlap_scheduler] - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-two_model-no_overlap_scheduler] - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-two_model-no_overlap_scheduler] - - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3_rejection_dynamic_tree_smoke[no_dynamic_tree] - - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3_rejection_dynamic_tree_smoke[dynamic_tree] - unittest/_torch/multi_gpu_modeling -k "deepseek" - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[pp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTEDSL-mtp_nextn=2-ep4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] From 26b1b3a5e25f84a25da5fe36751ea7a8a2746001 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Fri, 8 May 2026 02:29:01 -0700 Subject: [PATCH 14/17] [None][refactor] make Eagle3 dynamic-tree max_batch_size internal Eagle3DecodingConfig.max_batch_size only ever existed to give Eagle3OneModelDynamicTreeWorker.__init__ access to the global max_batch_size for sizing its persistent, batch-indexed CUDA buffers (draft_tokens_buffer, history_*_buffer, tree_mask_buffer, etc.). Those buffers are indexed by batch_idx at runtime with no bounds check, so this value MUST equal the global max_batch_size; while it was user-settable, a smaller value silently passed Pydantic validation but would OOB during warmup and real generation. Convert it to a PrivateAttr (_max_batch_size), following the same pattern as _allow_chain_drafter / _allow_separate_draft_kv_cache on DecodingBaseConfig. py_executor_creator is now the single writer, populating it from the global max_batch_size. Because Eagle3DecodingConfig inherits from StrictBaseModel (extra="forbid"), users who try to set it now get an explicit ValidationError instead of a silently-accepted, OOB-prone configuration. With the invariant guaranteed structurally, drop the three warmup/CUDA-graph clamps and the _get_dynamic_tree_warmup_max_batch_size helper in model_engine.py, and remove the redundant max_batch_size= kwarg from the dynamic-tree test in test_eagle3.py. Signed-off-by: ZhaoyangWang --- .../_torch/pyexecutor/model_engine.py | 26 ------------------- .../_torch/pyexecutor/py_executor_creator.py | 10 ++++--- .../_torch/speculative/eagle3_dynamic_tree.py | 14 +++++----- tensorrt_llm/llmapi/llm_args.py | 14 +++++++--- .../_torch/speculative/test_eagle3.py | 1 - 5 files changed, 25 insertions(+), 40 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 129415dc27d9..b75c60424537 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -457,13 +457,6 @@ def __init__( cuda_graph_batch_sizes, self.batch_size, self.max_num_tokens, self.original_max_total_draft_tokens, self._cuda_graph_padding_enabled) if cuda_graph_batch_sizes else [] - dynamic_tree_warmup_max_batch_size = self._get_dynamic_tree_warmup_max_batch_size( - ) - if dynamic_tree_warmup_max_batch_size is not None: - self._cuda_graph_batch_sizes = [ - bs for bs in self._cuda_graph_batch_sizes - if bs <= dynamic_tree_warmup_max_batch_size - ] self._max_cuda_graph_batch_size = (self._cuda_graph_batch_sizes[-1] if self._cuda_graph_batch_sizes else 0) @@ -755,11 +748,6 @@ def _get_max_shape_warmup_requests( max_batch_size = min( self.batch_size, curr_max_num_tokens // (1 + self.max_total_draft_tokens) // self.max_beam_width) - dynamic_tree_warmup_max_batch_size = self._get_dynamic_tree_warmup_max_batch_size( - ) - if dynamic_tree_warmup_max_batch_size is not None: - max_batch_size = min(max_batch_size, - dynamic_tree_warmup_max_batch_size) warmup_requests_configs = [ (curr_max_num_tokens, 0), # max_num_tokens, pure context @@ -789,14 +777,6 @@ def _get_full_general_warmup_requests( # Deduplicate the warmup_configs while keeping the order. return list(dict.fromkeys(warmup_configs)) - def _get_dynamic_tree_warmup_max_batch_size(self) -> Optional[int]: - """Return the dynamic-tree worker batch capacity for warmup batch clamping.""" - if (self.spec_config is None or self.is_draft_model - or not self.spec_config.spec_dec_mode.use_one_engine() - or not getattr(self.spec_config, "use_dynamic_tree", False)): - return None - return getattr(self.spec_config, "max_batch_size", None) - @with_warmup_flag @warmup_with_kv_cache_cleanup def warmup(self, resource_manager: ResourceManager) -> None: @@ -1050,12 +1030,6 @@ def _capture_generation_cuda_graphs(self, # Determine which graphs to capture graphs_to_capture = self._get_graphs_to_capture(cuda_graph_batch_sizes, spec_resource_manager) - dynamic_tree_warmup_max_batch_size = self._get_dynamic_tree_warmup_max_batch_size( - ) - if dynamic_tree_warmup_max_batch_size is not None: - graphs_to_capture = [(bs, draft_len) - for bs, draft_len in graphs_to_capture - if bs <= dynamic_tree_warmup_max_batch_size] graphs_to_capture = sorted(graphs_to_capture, reverse=True) # Create CUDA graphs for short and long sequences separately for sparse attention. # self.max_seq_len is the global max sequence length. For Helix CP each diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 522157840f55..fa33e0ce8643 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -446,9 +446,13 @@ def create_py_executor( has_draft_model_engine = spec_config.spec_dec_mode.has_draft_model() has_spec_drafter = spec_config.spec_dec_mode.has_spec_drafter() - if hasattr(spec_config, - 'max_batch_size') and spec_config.max_batch_size is None: - spec_config.max_batch_size = max_batch_size + # Eagle3DecodingConfig._max_batch_size is internally managed: the + # dynamic-tree worker pre-allocates batch-indexed CUDA buffers sized + # by this value, and runtime indexes them with no bounds check. It + # MUST equal the global max_batch_size to avoid OOB; we populate it + # here as the single source of truth. + if hasattr(spec_config, '_max_batch_size'): + spec_config._max_batch_size = max_batch_size # WAR for https://nvbugs/5807902 # Disable separate draft KV cache in disaggregated mode diff --git a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py index 2644a1b5dbc1..4fad34a74131 100644 --- a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py @@ -157,12 +157,14 @@ def __init__( self.K = spec_config.dynamic_tree_max_topK self.max_total_draft_tokens = spec_config.tokens_per_gen_step - 1 self.tokens_per_gen_step = spec_config.tokens_per_gen_step - if spec_config.max_batch_size is None: - raise ValueError( - "Eagle3OneModelDynamicTreeWorker requires max_batch_size to be set " - "on Eagle3DecodingConfig when use_dynamic_tree=True." - ) - self._max_batch_size = spec_config.max_batch_size + # Eagle3DecodingConfig._max_batch_size is auto-populated by + # py_executor_creator from the global max_batch_size, so it should + # always be set by the time we get here. + assert spec_config._max_batch_size is not None, ( + "Eagle3DecodingConfig._max_batch_size was not populated; " + "py_executor_creator should have set it from the global max_batch_size." + ) + self._max_batch_size = spec_config._max_batch_size K = self.K max_draft_len = spec_config.max_draft_len diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index b9d20e67b85f..151db668ad80 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1367,10 +1367,16 @@ class SAEnhancerConfig(StrictBaseModel): class Eagle3DecodingConfig(EagleDecodingConfig): decoding_type: Literal["Eagle3"] = "Eagle3" - max_batch_size: Optional[int] = Field( - default=None, - description="Max batch size for pre-allocating dynamic tree buffers. " - "Required when use_dynamic_tree=True.") + # Backs the dynamic-tree worker's pre-allocated, batch-indexed CUDA buffers + # (draft_tokens_buffer, history_*_buffer, tree_mask_buffer, etc. in + # Eagle3OneModelDynamicTreeWorker.__init__). This MUST equal the global + # max_batch_size: the worker indexes those buffers with batch_idx in + # [0, global_max_batch_size) at runtime with no bounds check, so any value + # smaller than the global will OOB during warmup or generation as soon as + # batch_idx exceeds this capacity (illegal memory access). It is therefore + # exposed as a PrivateAttr -- not a user-tunable knob -- and is + # auto-populated by py_executor_creator from the global max_batch_size. + _max_batch_size: Optional[int] = PrivateAttr(default=None) sa_config: Optional[SAEnhancerConfig] = Field( default=None, diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index 2ff2f6f5e868..ba97bff32422 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -994,7 +994,6 @@ def test_llama_eagle3_dynamic_tree(use_cuda_graph: bool, use_dynamic_tree=True, dynamic_tree_max_topK=dynamic_tree_max_topK, max_total_draft_tokens=max_total_draft_tokens, - max_batch_size=max_batch_size, ) # Create the LLM instance From 71a8887501548f14a8cc35f259ce4d8d1d10eea5 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Tue, 12 May 2026 03:16:06 -0700 Subject: [PATCH 15/17] [TRTLLM-11540][fix] drop max_batch_size kwarg from Eagle3DecodingConfig call sites Eagle3DecodingConfig.max_batch_size was moved to a PrivateAttr (_max_batch_size) auto-populated by py_executor_creator from the global LLM max_batch_size. Three call sites (quickstart example, integration smoke test, unit test) and one doc example still passed it as a constructor kwarg, which triggers pydantic extra_forbidden on every CI run since pipeline #37119. Signed-off-by: ZhaoyangWang --- docs/source/features/speculative-decoding.md | 4 ++-- examples/llm-api/quickstart_advanced.py | 3 +-- tests/integration/defs/accuracy/test_llm_api_pytorch.py | 1 - tests/unittest/_torch/speculative/test_eagle3.py | 1 - 4 files changed, 3 insertions(+), 6 deletions(-) diff --git a/docs/source/features/speculative-decoding.md b/docs/source/features/speculative-decoding.md index ce0465dfa294..9dbfff337ab4 100644 --- a/docs/source/features/speculative-decoding.md +++ b/docs/source/features/speculative-decoding.md @@ -59,7 +59,8 @@ To enable dynamic tree mode, set `use_dynamic_tree=True` on the `Eagle3DecodingC * `use_dynamic_tree` (`bool`): Enables dynamic tree draft generation. Mutually exclusive with `eagle_choices` (static tree). * `dynamic_tree_max_topK` (`int`): Maximum number of tokens to expand per node at each draft layer. * `max_total_draft_tokens` (`int`, optional): Total draft token budget for the tree. Must satisfy `max_draft_len <= max_total_draft_tokens <= dynamic_tree_max_topK * max_draft_len`. Defaults to `dynamic_tree_max_topK * max_draft_len` if not set. -* `max_batch_size` (`int`): Required when `use_dynamic_tree=True` for pre-allocating dynamic tree CUDA buffers. + +When `use_dynamic_tree=True`, the dynamic tree CUDA buffers are pre-allocated based on the `LLM`'s `max_batch_size`; it is propagated internally and must not be passed to `Eagle3DecodingConfig` directly. ```python from tensorrt_llm.llmapi import Eagle3DecodingConfig @@ -70,7 +71,6 @@ speculative_config = Eagle3DecodingConfig( use_dynamic_tree=True, dynamic_tree_max_topK=10, max_total_draft_tokens=60, - max_batch_size=4, ) llm = LLM("/path/to/target_model", speculative_config=speculative_config) diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py index 8d66de1c11ba..10901d87c520 100644 --- a/examples/llm-api/quickstart_advanced.py +++ b/examples/llm-api/quickstart_advanced.py @@ -296,8 +296,7 @@ def setup_llm(args, **kwargs): dynamic_tree_max_topK=args.dynamic_tree_max_topK, allow_advanced_sampling=args.allow_advanced_sampling, eagle3_model_arch=args.eagle3_model_arch, - max_total_draft_tokens=args.max_total_draft_tokens, - max_batch_size=args.max_batch_size) + max_total_draft_tokens=args.max_total_draft_tokens) elif spec_decode_algo == "DFLASH": spec_config = DFlashDecodingConfig( max_draft_len=args.spec_decode_max_draft_len, diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index f518ca05bd15..13bf1bb2a9be 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -208,7 +208,6 @@ def test_eagle3_rejection_dynamic_tree_smoke(self, use_dynamic_tree, use_dynamic_tree=True, dynamic_tree_max_topK=4, max_total_draft_tokens=16, - max_batch_size=max_batch_size, ) llm = LLM( diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index ba97bff32422..20c188578d41 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -872,7 +872,6 @@ def test_llama_eagle3_rejection_sampling_modes(use_dynamic_tree: bool, use_dynamic_tree=True, dynamic_tree_max_topK=dynamic_tree_max_top_k, max_total_draft_tokens=max_total_draft_tokens, - max_batch_size=max_batch_size, ) llm_spec = LLM(**llm_common_config, From 38a011cf572a83addb290b2146778020105a8836 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Tue, 12 May 2026 19:26:56 -0700 Subject: [PATCH 16/17] [None][fix] make Eagle3 one-model rejection sampling CUDA-graph-safe Eliminate host-device syncs (`.item()`) and missing fake-impl registrations introduced by the dynamic-tree rejection sampling rewrite: - Pass `top_ks=None` whenever `spec_metadata.skip_top_k` is set (mirroring the existing `skip_top_p` handling) so the C++ ops can short-circuit on a host-only optional check. - Replace the unconditional `topK->gt(0).any().item()` / `topP->lt(1.0).any().item()` probes in `computeProbsFromLogits` and `applyTopKTopPForProbOp` with host-side `has_value()` checks; the per-row `effectiveTopK` formula already handles disabled rows. The fast kernel-path probe is deferred into the `kMax > 0` branch (used only by the non-graph-captured dynamic-tree caller). - Register `torch.library.register_fake` impls for the three new ops (`compute_probs_from_logits_op`, `compute_draft_probs_for_dynamic_tree_rejection_op`, `compute_target_probs_for_dynamic_tree_rejection_op`) so `test_register_fake` passes. Fixes the CI failures `test_llama_eagle3_rejection_sampling_modes[True-no_dynamic_tree]` (cudaErrorStreamCaptureUnsupported during warmup graph capture) and `test_register_fake` on PR #12588. Signed-off-by: ZhaoyangWang --- .../speculativeDecoding/dynamicTreeKernels.cu | 44 +++++++++++---- .../_torch/custom_ops/cpp_custom_ops.py | 54 +++++++++++++++++++ tensorrt_llm/_torch/speculative/interface.py | 13 +++-- 3 files changed, 97 insertions(+), 14 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu index 8b94d9dfb559..7a5e922d6ad1 100644 --- a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu +++ b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu @@ -455,9 +455,12 @@ torch::Tensor applyTopKTopPForProbOp(torch::Tensor logits, torch::optional const& topP, int32_t kMax) { int64_t const vocabSize = logits.size(1); - bool const hasTopK = topK.has_value() && topK->defined() - && torch::logical_and(topK->gt(0), topK->lt(vocabSize)).any().item(); - bool const hasTopP = topP.has_value() && topP->defined() && topP->lt(1.0).any().item(); + // Host-only checks: the caller is expected to pass nullopt when filtering is fully + // disabled (see SpecMetadata.skip_top_k / skip_top_p). Probing the tensor contents + // via `.item()` here would force a host-device sync and break CUDA graph + // capture; the per-row `effectiveTopK` formula below already handles disabled rows. + bool const hasTopK = topK.has_value() && topK->defined(); + bool const hasTopP = topP.has_value() && topP->defined(); if (!hasTopK && !hasTopP) { @@ -465,12 +468,22 @@ torch::Tensor applyTopKTopPForProbOp(torch::Tensor logits, torch::optionalto(torch::kLong); effectiveTopK = torch::where(topKLong > 0, topKLong, torch::full_like(topKLong, vocabSize)).clamp_max(vocabSize); + } + + // Fast path uses `topk(kMax)` which is unsafe when any row has effective top-k > kMax + // (i.e. disabled rows expand to the full vocab). Detecting this requires a tensor + // reduction + `.item()`, which is incompatible with CUDA graph capture. Only + // probe when the caller explicitly opted into the fast path via kMax > 0 (today only + // the dynamic-tree caller, which is not graph-captured). + bool hasDisabledTopKRows = false; + if (hasTopK && kMax > 0 && kMax < vocabSize) + { + auto topKLong = topK->to(torch::kLong); hasDisabledTopKRows = topKLong.le(0).any().item(); } @@ -573,20 +586,31 @@ torch::Tensor computeProbsFromLogits(torch::Tensor const& logits, torch::Tensor int64_t const vocabSize = scaledLogits.size(1); int64_t const nRows = scaledLogits.size(0); - bool const hasTopK = topK.has_value() && topK->defined() - && torch::logical_and(topK->gt(0), topK->lt(vocabSize)).any().item(); - bool const hasTopP = topP.has_value() && topP->defined() && topP->lt(1.0).any().item(); + // Host-only presence checks; see comment in applyTopKTopPForProbOp() for why we + // avoid probing tensor contents (would sync and break CUDA graph capture). + bool const hasTopKPresence = topK.has_value() && topK->defined(); + bool const hasTopPPresence = topP.has_value() && topP->defined(); + + // The kernel path produces -inf for rows whose top_k value is 0, so it is only + // safe when every row has an active top_k filter. Determining that requires a + // host-device sync, so only probe when the caller has opted into the kernel + // path (kMax > 0). The kMax > 0 callers (dynamic-tree) are not graph-captured. + bool useKernelPath = false; + if (hasTopKPresence && kMax > 0 && kMax < vocabSize) + { + useKernelPath = torch::logical_and(topK->gt(0), topK->lt(vocabSize)).any().item(); + } torch::Tensor maskedLogits; - if (hasTopK && kMax > 0 && kMax < vocabSize) + if (useKernelPath) { // Two-stage CUDA top-k/top-p masking (mirrors invokeBatchTopKSampling). maskedLogits = torch::empty_like(scaledLogits); auto topKForKernel = topK->to(torch::kInt32).contiguous(); - auto topPForKernel = hasTopP ? topP->to(torch::kFloat32).contiguous() : torch::Tensor(); + auto topPForKernel = hasTopPPresence ? topP->to(torch::kFloat32).contiguous() : torch::Tensor(); auto stream = at::cuda::getCurrentCUDAStream(scaledLogits.device().index()); invokeTopKTopPMaskingForProbs(scaledLogits.data_ptr(), maskedLogits.data_ptr(), - topKForKernel.data_ptr(), hasTopP ? topPForKernel.data_ptr() : nullptr, kMax, + topKForKernel.data_ptr(), hasTopPPresence ? topPForKernel.data_ptr() : nullptr, kMax, static_cast(nRows), static_cast(vocabSize), stream); } else diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 75f4f9d84465..6ffb227560d0 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -1190,3 +1190,57 @@ def _(like: torch.Tensor, out_shape = shape if shape is not None else list(like.shape) dtype = out_dtype if out_dtype is not None else like.dtype return like.new_empty(out_shape, dtype=dtype), output_buffer_kind + + @torch.library.register_fake("trtllm::compute_probs_from_logits_op") + def _(logits: torch.Tensor, + temperatures: torch.Tensor, + top_k: Optional[torch.Tensor] = None, + top_p: Optional[torch.Tensor] = None, + skip_temperature: bool = False) -> torch.Tensor: + return logits.new_empty(list(logits.shape), dtype=torch.float32) + + @torch.library.register_fake( + "trtllm::compute_draft_probs_for_dynamic_tree_rejection_op") + def _(draftLogits: torch.Tensor, + temperatures: torch.Tensor, + numDraftProbRows: int, + targetVocabSize: int, + top_k: Optional[torch.Tensor] = None, + top_p: Optional[torch.Tensor] = None, + skip_temperature: bool = False, + d2t: Optional[torch.Tensor] = None, + top_k_max: int = 0, + skip_all_sampling_params: bool = False) -> torch.Tensor: + batch_size = temperatures.shape[0] + return draftLogits.new_empty( + (batch_size, numDraftProbRows, targetVocabSize), + dtype=torch.float32) + + @torch.library.register_fake( + "trtllm::compute_target_probs_for_dynamic_tree_rejection_op") + def _( + targetLogits: torch.Tensor, + temperatures: torch.Tensor, + numDraftTokens: int, + top_k: Optional[torch.Tensor] = None, + top_p: Optional[torch.Tensor] = None, + skip_temperature: bool = False, + top_k_max: int = 0, + skip_all_sampling_params: bool = False + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + batch_size = temperatures.shape[0] + target_vocab_size = targetLogits.shape[-1] + target_probs = targetLogits.new_empty( + (batch_size, numDraftTokens, target_vocab_size), + dtype=torch.float32) + has_filtering = (top_k is not None) or (top_p is not None) + if skip_all_sampling_params or not has_filtering: + support_indices = targetLogits.new_empty((0, ), dtype=torch.int32) + support_lengths = targetLogits.new_empty((0, ), dtype=torch.int32) + else: + support_dim = top_k_max if top_k_max > 0 else target_vocab_size + support_indices = targetLogits.new_empty( + (batch_size, numDraftTokens, support_dim), dtype=torch.int32) + support_lengths = targetLogits.new_empty( + (batch_size, numDraftTokens), dtype=torch.int32) + return target_probs, support_indices, support_lengths diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index d623ce0eeb61..c7a8358124cf 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -1029,7 +1029,11 @@ def _sample_and_accept_draft_tokens_rejection( num_target_tokens = batch_size * (runtime_draft_len + 1) temperatures = spec_metadata.temperatures[:num_target_tokens] - top_ks = spec_metadata.top_ks[:num_target_tokens] + # Pass None instead of an all-disabled tensor so the C++ op can short-circuit + # on a host-side check rather than a `.item()` sync, which would break + # CUDA graph capture. + top_ks = None if spec_metadata.skip_top_k else spec_metadata.top_ks[: + num_target_tokens] top_ps = None if spec_metadata.skip_top_p else spec_metadata.top_ps[: num_target_tokens] @@ -1127,9 +1131,10 @@ def _compute_and_store_draft_probs( if spec_metadata.request_temperatures is not None: draft_temps = spec_metadata.request_temperatures[:batch_size].repeat_interleave( draft_tokens_per_request) - draft_top_ks = spec_metadata.request_top_ks[:batch_size].repeat_interleave( - draft_tokens_per_request - ) if spec_metadata.request_top_ks is not None else None + draft_top_ks = ( + spec_metadata.request_top_ks[:batch_size].repeat_interleave( + draft_tokens_per_request) if not spec_metadata.skip_top_k + and spec_metadata.request_top_ks is not None else None) draft_top_ps = ( spec_metadata.request_top_ps[:batch_size].repeat_interleave( draft_tokens_per_request) if not spec_metadata.skip_top_p From 8801e35df8e894aba13c6836215fa14f9650c41e Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Thu, 14 May 2026 00:02:49 -0700 Subject: [PATCH 17/17] [None][fix] address Eagle3 rejection sampling review feedback Signed-off-by: ZhaoyangWang --- docs/source/features/speculative-decoding.md | 2 +- tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py | 15 +++++++++++++++ .../_torch/speculative/one_model_sampler.py | 2 +- 3 files changed, 17 insertions(+), 2 deletions(-) diff --git a/docs/source/features/speculative-decoding.md b/docs/source/features/speculative-decoding.md index 9dbfff337ab4..097edc9e2e6e 100644 --- a/docs/source/features/speculative-decoding.md +++ b/docs/source/features/speculative-decoding.md @@ -60,7 +60,7 @@ To enable dynamic tree mode, set `use_dynamic_tree=True` on the `Eagle3DecodingC * `dynamic_tree_max_topK` (`int`): Maximum number of tokens to expand per node at each draft layer. * `max_total_draft_tokens` (`int`, optional): Total draft token budget for the tree. Must satisfy `max_draft_len <= max_total_draft_tokens <= dynamic_tree_max_topK * max_draft_len`. Defaults to `dynamic_tree_max_topK * max_draft_len` if not set. -When `use_dynamic_tree=True`, the dynamic tree CUDA buffers are pre-allocated based on the `LLM`'s `max_batch_size`; it is propagated internally and must not be passed to `Eagle3DecodingConfig` directly. +When `use_dynamic_tree=True`, the dynamic tree CUDA buffers are pre-allocated based on the `LLM`'s `max_batch_size`, which is propagated internally and must not be passed directly to `Eagle3DecodingConfig`. ```python from tensorrt_llm.llmapi import Eagle3DecodingConfig diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 6ffb227560d0..a43d6e6f530b 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -1199,6 +1199,21 @@ def _(logits: torch.Tensor, skip_temperature: bool = False) -> torch.Tensor: return logits.new_empty(list(logits.shape), dtype=torch.float32) + @torch.library.register_fake("trtllm::build_draft_prob_indices_out_op") + def _(topkScoreIndices: torch.Tensor, draftProbIndices: torch.Tensor, + topK: int, numDraftTokens: int) -> None: + return None + + @torch.library.register_fake("trtllm::verify_dynamic_tree_rejection_out_op") + def _(candidates: torch.Tensor, draftProbs: torch.Tensor, + targetProbs: torch.Tensor, targetSupportIndices: torch.Tensor, + targetSupportLengths: torch.Tensor, draftProbIndices: torch.Tensor, + retrieveNextToken: torch.Tensor, retrieveNextSibling: torch.Tensor, + treeValid: torch.Tensor, acceptIndex: torch.Tensor, + acceptTokenNum: torch.Tensor, acceptToken: torch.Tensor, + numSpecStep: int, seed: torch.Tensor, offset: torch.Tensor) -> None: + return None + @torch.library.register_fake( "trtllm::compute_draft_probs_for_dynamic_tree_rejection_op") def _(draftLogits: torch.Tensor, diff --git a/tensorrt_llm/_torch/speculative/one_model_sampler.py b/tensorrt_llm/_torch/speculative/one_model_sampler.py index ae1dda63e658..7e3b06b383cc 100644 --- a/tensorrt_llm/_torch/speculative/one_model_sampler.py +++ b/tensorrt_llm/_torch/speculative/one_model_sampler.py @@ -116,7 +116,7 @@ def compute_probs_from_logits( """ Compute probabilities from logits with temperature, top-k, and top-p applied. """ - if logits.is_cuda and hasattr(torch.ops.trtllm, "compute_probs_from_logits_op"): + if logits.is_cuda: return torch.ops.trtllm.compute_probs_from_logits_op( logits, temperatures, top_k, top_p, skip_temperature )