diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index 32c5c246baf8..1d51a28dfba9 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # +import math from dataclasses import replace from typing import Dict, Optional @@ -28,6 +29,7 @@ get_dflash_trtllm_gen_ops, ) from ..speculative.interface import SpeculativeDecodingMode +from ..utils import maybe_compile from .modeling_utils import FUSED_MODULE_COMPONENTS, get_model_architecture, register_draft_model @@ -63,6 +65,56 @@ def dspark_layer_window_size( return (swa_window - 1, swa_window - 1) +@maybe_compile +def dflash_merge_later_block_keys_( + out: torch.Tensor, + lse_log2: torch.Tensor, + query: torch.Tensor, + block_key: torch.Tensor, + block_value: torch.Tensor, + scale: float, + window_right: int, +) -> None: + """Add each block query's later keys ``j < k <= j + window_right`` to its causal attention + ``out`` [B * blk, Hq, D], in place, given the base-2 log-sum-exp ``lse_log2`` [B * blk, Hq]. + The softmax over the few block keys is unrolled so the merge compiles to one pointwise kernel.""" + B, blk, num_heads, head_dim = query.shape + num_kv_heads = block_key.shape[2] + group = num_heads // num_kv_heads + q = query.float().reshape(B, blk, num_kv_heads, group, 1, head_dim) + keys = block_key.float().permute(0, 2, 1, 3)[:, None, :, None] # [B, 1, Hkv, 1, blk, D] + logits = (q * keys).sum(-1) * scale # [B, query j, Hkv, g, key k] + offsets = torch.arange(blk, device=query.device) + distance = offsets.view(1, -1) - offsets.view(-1, 1) # [j, k] = k - j + later = distance > 0 + if window_right >= 0: + later &= distance <= window_right + later_logits = [ + logits[..., k].masked_fill(~later[:, k].view(1, blk, 1, 1), float("-inf")) + for k in range(blk) + ] + max_logit = later_logits[0] + for logit in later_logits[1:]: + max_logit = torch.maximum(max_logit, logit) + # A query without later keys has only -inf logits; its weight is 0 below. + max_logit = torch.where(torch.isfinite(max_logit), max_logit, 0.0) + exp_logits = [torch.exp(logit - max_logit) for logit in later_logits] + sum_exp = exp_logits[0] + for exp_logit in exp_logits[1:]: + sum_exp = sum_exp + exp_logit + has_later = sum_exp > 0 + sum_exp = torch.where(has_later, sum_exp, 1.0) + lse_causal = lse_log2.float().reshape(B, blk, num_kv_heads, group) * math.log(2.0) + # exp(lse_later) / (exp(lse_causal) + exp(lse_later)) + weight = torch.where(has_later, torch.sigmoid(max_logit + torch.log(sum_exp) - lse_causal), 0.0) + merged = out.float().reshape(B, blk, num_kv_heads, group, head_dim) * (1 - weight)[..., None] + value = block_value.float() + prob_scale = weight / sum_exp + for k in range(blk): + merged = merged + (prob_scale * exp_logits[k])[..., None] * value[:, k, None, :, None] + out.copy_(merged.reshape(B * blk, num_heads, head_dim)) + + def _is_dflash2_architecture(config: PretrainedConfig) -> bool: """Whether a draft checkpoint advertises itself as a DFlash 2 drafter. @@ -1385,30 +1437,6 @@ def _resolve_block_attention(self, layer_idx: int) -> tuple[bool, tuple[int, int window_size = swa_window return causal, window_size - def validate_block_attention_windows(self) -> None: - """Refuse windows the TRTLLM block-decode kernel cannot apply. - - TRTLLM-Gen only implements sliding windows for causal masks: flashinfer - rejects ``causal=False`` with any finite ``window_left`` outright. A - non-causal windowed layer must therefore run on VANILLA or FA4, which - take a two-sided window. - """ - if self.dflash_attention_backend != "TRTLLM": - return - num_layers = getattr(self.config, "num_hidden_layers", None) - if num_layers is None: - num_layers = len(self.model.layers) - for layer_idx in range(num_layers): - causal, (window_left, _) = self._resolve_block_attention(layer_idx) - if causal or window_left < 0: - continue - raise ValueError( - f"DFlash draft layer {layer_idx} attends non-causally within a " - f"{window_left + 1}-token sliding window. The TRTLLM DFlash " - "attention backend does not support non-causal sliding-window " - "attention. Use attention_backend=VANILLA or FA4." - ) - def _prepare_dflash_trtllm_gen_buffers( self, dtype: torch.dtype, @@ -1431,6 +1459,8 @@ def _prepare_dflash_trtllm_gen_buffers( rotary_embedding_dim=0, fp8_context_fmha=False, ) + # Softmax stats flashinfer carves off the front when the decode returns its log-sum-exp. + workspace_bytes += 8 * num_heads * max_batch_size * 256 * -(-block_size // 256) + (1 << 20) device = torch.device(device) is_capturing = torch.cuda.is_current_stream_capturing() if self._dflash_trtllm_gen_device != device: @@ -1873,8 +1903,10 @@ def dflash_forward( q_flat = Q_bshd.reshape(-1, num_heads_per_rank, head_dim) out_flat = out.reshape(-1, num_heads_per_rank, head_dim) window_left = window_size[0] - if causal: - trtllm_gen_ops.batch_decode_with_kv_cache( + # trtllm-gen windows only causal attention; a non-causal window adds the later block keys. + merge_later_keys = not causal and window_left >= 0 and block_size > 1 + if causal or window_left >= 0: + decoded = trtllm_gen_ops.batch_decode_with_kv_cache( query=q_flat, kv_cache=(layer_cache[:, 0], layer_cache[:, 1]), workspace_buffer=self._dflash_trtllm_gen_workspace, @@ -1897,8 +1929,19 @@ def dflash_forward( kv_cache_sf=None, uses_shared_paged_kv_idx=True, bmm1_scale_log2=None, + return_lse=merge_later_keys, multi_ctas_kv_counter_buffer=self._dflash_trtllm_gen_counters, ) + if merge_later_keys: + dflash_merge_later_block_keys_( + out_flat, + decoded[1], + Q_bshd, + k_noise_bshd, + v_noise_bshd, + head_dim**-0.5, + window_size[1], + ) else: cum_seq_lens_q = torch.arange( 0, diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 1e2edea4800d..52bf168d8da2 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -2772,6 +2772,11 @@ def __init__( self.mask_token_id = int(spec_mask_token_id) self._init_dspark_heads(draft_config.pretrained_config) + def _get_attention_mask_args(self, layer_idx): + """DSpark drafters are trained non-causal; a sliding window bounds both sides.""" + _, (window_left, _) = super()._get_attention_mask_args(layer_idx) + return False, (window_left, window_left) + def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): """Take the DSpark head weights, then hand the rest to DFlash.""" weights, _ = self._take_dspark_head_weights(weights) diff --git a/tensorrt_llm/_torch/speculative/block_verification.py b/tensorrt_llm/_torch/speculative/block_verification.py new file mode 100644 index 000000000000..236430726a05 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/block_verification.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Block verification (Sun et al., 2024, arXiv:2403.10444) for one-model speculative decoding. + +For drafts ``x_j ~ q_j`` (``j < K``), target rows ``p_j`` (bonus row ``p_K``) and uniforms ``u_j``: +``w_j = min(1, w_{j-1} p_j(x_j) / q_j(x_j))``, ``r_j = sum_x max(w_j p_{j+1}(x) - q_{j+1}(x), 0)``, +``h_j = r_j / (r_j + 1 - w_j)`` and ``h_{K-1} = w_{K-1}``. ``L`` is the largest ``j + 1`` with +``u_j <= h_j`` (else 0); the emitted token is drawn from ``max(w_{L-1} p_L - q_L, 0)``, or ``p_K`` if ``L = K``. +""" + +from typing import Callable + +import torch + +_MASK32 = 0xFFFFFFFF + + +def _mul_u32(values: torch.Tensor, constant: int) -> torch.Tensor: + """``(values * constant) mod 2**32``, multiplying 16-bit halves so int64 never overflows.""" + low = values * (constant & 0xFFFF) + high = ((values * (constant >> 16)) & 0xFFFF) << 16 + return (low + high) & _MASK32 + + +def _mix_u32(values: torch.Tensor) -> torch.Tensor: + """The lowbias32 finalizer.""" + values = values ^ (values >> 16) + values = _mul_u32(values, 0x7FEB352D) + values = values ^ (values >> 15) + values = _mul_u32(values, 0x846CA68B) + return values ^ (values >> 16) + + +def counter_uniforms(seed: torch.Tensor, offset: torch.Tensor, num_draws: int) -> torch.Tensor: + """``[rows, num_draws]`` float64 uniforms in (0, 1), hashed from each row's ``(seed, offset)``. + + The same key gives the same values on every rank and in every graph replay; a new offset + gives fresh ones. + """ + seed = seed.to(torch.int64).unsqueeze(-1) + offset = offset.to(torch.int64).unsqueeze(-1) + draws = torch.arange(num_draws, dtype=torch.int64, device=seed.device) + key = _mix_u32((seed & _MASK32) ^ 0x5BD1E995) + key = _mix_u32(key ^ ((seed >> 32) & _MASK32)) + key = _mix_u32(key ^ (offset & _MASK32)) + key = _mix_u32(key ^ ((offset >> 32) & _MASK32)) + bits = _mix_u32(_mix_u32(key ^ draws)) + return (bits.to(torch.float64) + 0.5) * 2.0**-32 + + +def block_verify( + draft_tokens: torch.Tensor, + draft_probs: torch.Tensor, + target_probs: torch.Tensor, + uniforms: torch.Tensor, + sample_emitted: Callable[[torch.Tensor], torch.Tensor], +) -> tuple[torch.Tensor, torch.Tensor]: + """Verify ``[B, K]`` drafts; ``sample_emitted`` draws from ``[B, V]`` nonnegative weights. + + Returns the token-level kernel's layout: ``[B, K + 1]`` int32 tokens (accepted drafts, the emitted + token, then -1) and ``[B]`` int32 counts including the emitted token. No host sync (graph-safe). + """ + num_drafts = draft_tokens.shape[1] + vocab_size = target_probs.shape[-1] + tokens = draft_tokens.long().clamp(0, vocab_size - 1).unsqueeze(-1) + p_x = target_probs[:, :num_drafts].gather(-1, tokens).squeeze(-1).float() + q_x = draft_probs.gather(-1, tokens).squeeze(-1).float() + # A draft its proposal gives no mass (or a non-finite row) is never accepted. + log_ratio = torch.where(q_x > 0, p_x.log() - q_x.log(), torch.full_like(q_x, float("-inf"))) + log_ratio = torch.nan_to_num(log_ratio, nan=float("-inf")) + log_w = torch.empty_like(log_ratio) + running = torch.zeros_like(log_ratio[:, 0]) + for j in range(num_drafts): + running = torch.clamp(running + log_ratio[:, j], max=0.0) + log_w[:, j] = running + w = log_w.exp() + h = w.clone() + if num_drafts > 1: + residual = w[:, :-1, None] * target_probs[:, 1:num_drafts].float() + residual.sub_(draft_probs[:, 1:].float()).clamp_min_(0.0) + mass = residual.sum(-1) + denom = mass + 1.0 - w[:, :-1] + h[:, :-1] = torch.where(denom > 0, mass / denom, torch.ones_like(mass)) + + positions = torch.arange(1, num_drafts + 1, dtype=torch.int64, device=draft_tokens.device) + accepted_len = (positions * (uniforms.to(torch.float64) <= h.to(torch.float64))).amax(-1) + row = accepted_len.view(-1, 1, 1) + p_row = target_probs.gather(1, row.expand(-1, 1, vocab_size)).squeeze(1).float() + q_row = draft_probs.gather(1, row.clamp(max=num_drafts - 1).expand(-1, 1, vocab_size)) + w_prev = torch.cat([torch.ones_like(w[:, :1]), w], dim=1).gather(1, accepted_len.view(-1, 1)) + is_bonus = (accepted_len == num_drafts).unsqueeze(-1) + residual = torch.where( + is_bonus, p_row, (w_prev * p_row - q_row.squeeze(1).float()).clamp_min(0.0) + ) + # Fall back to the target row if rounding leaves the residual without mass. + residual = torch.where(residual.sum(-1, keepdim=True) > 0, residual, p_row) + emitted = sample_emitted(residual).long().clamp(0, vocab_size - 1) + + slots = torch.arange(num_drafts + 1, dtype=torch.int64, device=draft_tokens.device) + drafts = torch.cat([draft_tokens.long(), torch.zeros_like(draft_tokens[:, :1]).long()], dim=1) + accepted_tokens = torch.where( + slots < accepted_len[:, None], drafts, torch.full_like(drafts, -1) + ) + accepted_tokens = torch.where(slots == accepted_len[:, None], emitted[:, None], accepted_tokens) + return accepted_tokens.to(torch.int32), (accepted_len + 1).to(torch.int32) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 410a12f91512..46eb3157dbeb 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -1423,8 +1423,7 @@ def _lazy_init_ctx_buffers( ): # Gated on the drafter that actually runs the trtllm-gen ops: # one paging its own block decode the same way is not shaped - # like them. No non-causal sliding window there either. - draft_model.validate_block_attention_windows() + # like them. has_context_attention = any( not draft_model._get_attention_mask_args(layer_idx)[0] for layer_idx in range(L) ) @@ -2479,11 +2478,8 @@ def _forward_impl( ) vocab_size = gen_logits.shape[-1] - gen_draft_tokens = self.sample_draft_tokens( - gen_logits, - spec_metadata, - batch_size, - num_contexts=num_contexts, + gen_draft_tokens, vocab_size = self._draft_block_tokens( + draft_model, gen_logits, inputs, spec_metadata, batch_size, num_contexts ) # Opt-in confidence-calibration recording (env-gated). The @@ -2598,6 +2594,15 @@ def _refine_block_logits( """ return gen_logits + def _draft_block_tokens( + self, draft_model, gen_logits, inputs, spec_metadata, batch_size, num_contexts + ): + """The block's draft tokens and the vocab width they index.""" + tokens = self.sample_draft_tokens( + gen_logits, spec_metadata, batch_size, num_contexts=num_contexts + ) + return tokens, gen_logits.shape[-1] + def _apply_dflash2_selector( self, draft_model, diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 322ef523a670..78445edd64fc 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -1251,6 +1251,10 @@ def _draft_slot_ids( def _try_fused_draft_head_and_refine( self, draft_model, hidden, attn_metadata, inputs, spec_metadata, num_gens, k ): + if self._samples_block_ancestrally(draft_model, spec_metadata): + return draft_model.logits_processor( + hidden, draft_model.lm_head, attn_metadata, True + ).reshape(num_gens, k, -1) # Only the M3.1 SM107 TP2 target's explicit A1 processor opts into this adapter. if not type(self).__module__.endswith("minimax_m31_nvfp4__sm_107__tp2.dspark"): return None @@ -1262,6 +1266,70 @@ def _try_fused_draft_head_and_refine( self, draft_model, hidden, attn_metadata, inputs, spec_metadata, num_gens, k ) + def _samples_block_ancestrally(self, draft_model, spec_metadata) -> bool: + """Sampled drafts follow the Markov head through the tokens actually drawn.""" + return bool( + getattr(draft_model, "has_markov_head", False) + and spec_metadata.wants_advanced_draft_sampling + ) + + def _draft_block_tokens( + self, draft_model, gen_logits, inputs, spec_metadata, batch_size, num_contexts + ): + if self._samples_block_ancestrally(draft_model, spec_metadata): + return self._sample_block_ancestrally( + draft_model, gen_logits, inputs, spec_metadata, num_contexts, batch_size + ) + return super()._draft_block_tokens( + draft_model, gen_logits, inputs, spec_metadata, batch_size, num_contexts + ) + + def _sample_block_ancestrally( + self, draft_model, gen_logits, inputs, spec_metadata, num_contexts, batch_size + ): + """Draw position k from ``softmax((base_k + bias(d_{k-1})) / T)``, ``d_{k-1}`` the token drawn + before it (the anchor for k = 0), with the temperature only, and store that distribution as + the proposal. Returns the int32 tokens and the stored vocab width.""" + from ..models.modeling_speculative import dspark_markov_step_bias + + num_gens, num_drafts, _ = gen_logits.shape + # The chain indexes the full Markov table; verification reads full-vocab q. + gen_logits = self.maybe_gather_sharded_draft_logits(gen_logits, spec_metadata) + vocab = gen_logits.shape[-1] + # Zero the CUDA-graph padding rows (uninitialized state) so every draw is legal. + gen_logits = self._zero_padding_rows( + gen_logits.reshape(num_gens * num_drafts, vocab), + spec_metadata, + num_contexts, + batch_size, + num_drafts, + ).reshape(num_gens, num_drafts, vocab) + params = self._temperature_only_params( + spec_metadata.request_temperatures[num_contexts:batch_size] + ) + draft_probs = spec_metadata.draft_probs + slots = spec_metadata.batch_slot_ids[num_contexts:batch_size] + tokens = inputs["first_prev_tokens"] + drafts = [] + for step in range(num_drafts): + bias = dspark_markov_step_bias(tokens, draft_model.markov_w1, draft_model.markov_w2) + seed, offset = self._rng_state_per_request( + spec_metadata, num_contexts, batch_size, slot=spec_metadata.rng_draft_slot(step) + ) + tokens, probs = self._sample_from_logits_with_probs( + spec_metadata.advanced_sampling_mode, + gen_logits[:, step] + bias.to(gen_logits.dtype), + *params, + seed=seed, + offset=offset, + ) + if draft_probs is not None: + draft_probs[slots, step, :vocab] = probs + drafts.append(tokens) + if draft_probs is not None: + spec_metadata.draft_probs_last_dim = vocab + return torch.stack(drafts, dim=1).to(torch.int32), vocab + def _refine_block_logits( self, draft_model, @@ -1295,8 +1363,9 @@ def _apply_dspark_markov_bias( logits, where prev_0 is the anchor (last accepted) token and prev_{i>0} is the greedy token from step i-1's biased logits. Greedy per-position argmax of the returned logits therefore reproduces the reference - sampled chain; the rejection-sampling path samples from the same - biased distributions (proposal conditioned on the greedy chain). + chain. Sampled drafts do not come through here: they are drawn by + :meth:`_sample_block_ancestrally`, which conditions each position on + the token actually drawn before it. Handles a TP vocab-sharded draft lm_head by slicing markov_w2's rows to this rank's contiguous shard and chaining through the TP-aware diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index c188d5d231cf..4f3344b52aa9 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -49,6 +49,7 @@ from ..pyexecutor.sampler.ops import custom as fused_sampling from ..pyexecutor.sampler.ops import flashinfer as flashinfer_sampling from ..pyexecutor.sampler.ops.vanilla import greedy_search_sampling_batch +from .block_verification import block_verify, counter_uniforms # Logit given to token 0 of a CUDA-graph padding row (see # SpecWorkerBase._zero_padding_rows). Finite in @@ -733,6 +734,9 @@ def is_ragged_verify(self) -> bool: # Drop the sampler tail's no-op bookkeeping launches, bitwise the same # outputs (RuntimeHooks.spec_tail_bookkeep; see tail_bookkeep_draft_len). tail_bookkeep: bool = False + # Verification rule of the rejection path, "standard" or "block" (deploy-time; + # from DecodingBaseConfig.rejection_sample_method). + rejection_sample_method: str = "standard" # Advanced-sampling specialization (deploy-time; from DecodingBaseConfig.advanced_sampling_mode). advanced_sampling_mode: AdvancedSamplingMode = AdvancedSamplingMode.FULL # Whether the occurrence penalties are enabled (deploy-time; from @@ -2695,6 +2699,17 @@ def _uniform_padding_probs(probs: torch.Tensor, probs.masked_fill_(is_padding.unsqueeze(-1), 1.0 / probs.shape[-1]) return probs + @staticmethod + def _temperature_only_params( + temperatures: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Sampler arguments for ``temperatures`` without filters (device-side fills, graph-safe).""" + return (temperatures, + torch.full_like(temperatures, + torch.iinfo(torch.int32).max, + dtype=torch.int32), + torch.ones_like(temperatures), torch.zeros_like(temperatures)) + def advanced_sample_draft(self, logits: torch.Tensor, spec_metadata: "SpecMetadata", @@ -3052,19 +3067,38 @@ def _sample_and_accept_draft_tokens_rejection( batch_size, slot=spec_metadata.rng_accept_slot) + use_block = spec_metadata.rejection_sample_method == "block" slot_out = None if slot_rows is not None: - slot_out = _CHAIN_SPEC_SLOT_OVERRIDE(slot_rows[0], slot_rows[1], - runtime_draft_len, - full_draft_tokens, - target_probs, seed, offset) + if not use_block: + slot_out = _CHAIN_SPEC_SLOT_OVERRIDE( + slot_rows[0], slot_rows[1], runtime_draft_len, + full_draft_tokens, target_probs, seed, offset) if slot_out is None: - # Declined: the gathered copy, as without the override. + # Declined, or block verification: the gathered copy, as + # without the override. full_draft_probs = slot_rows[0][ slot_rows[1], :runtime_draft_len, :vocab_size] if slot_out is not None: gen_accepted, gen_num_accepted = slot_out + elif use_block and not is_ragged: + unit = torch.ones_like(seed, dtype=torch.float32) + gen_accepted, gen_num_accepted = block_verify( + full_draft_tokens, full_draft_probs, target_probs, + counter_uniforms(seed, offset, runtime_draft_len), + lambda weights: self._sample_from_logits( + spec_metadata.advanced_sampling_mode, + weights.log(), + *self._temperature_only_params(unit), + seed=seed, + offset=offset)) else: + if use_block: + logger.warning_once( + "rejection_sample_method='block' does not cover ragged " + "(confidence-scheduled) verification; these steps use " + "token-level rejection sampling.", + key="block_verification_ragged_fallback") gen_accepted, gen_num_accepted = rejection_sampling_one_model( draft_probs=full_draft_probs, draft_token_ids=full_draft_tokens, diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 2bb15ed1d727..55ab8a890275 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -359,6 +359,7 @@ def get_spec_metadata(spec_config, # dispatcher routes to a backend that takes no min_p argument -- so the filter is # dropped with nothing raised. One assignment here keeps them in step. metadata.advanced_sampling_mode = spec_config.advanced_sampling_mode + metadata.rejection_sample_method = spec_config.rejection_sample_method if num_seq_slots is not None: metadata.num_seq_slots = num_seq_slots return metadata diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 41f3c3976f73..66366bf53f8e 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2367,6 +2367,16 @@ class DecodingBaseConfig(StrictBaseModel): "(default) to use exact-match verification on non-greedy batches. " "The non-dynamic-tree one-model path requires FlashInfer.") + rejection_sample_method: Literal["standard", "block"] = Field( + default="standard", + status="prototype", + description= + "Verification rule of rejection sampling (use_rejection_sampling=True): " + "'standard' is token-level rejection sampling; 'block' is block " + "verification (Sun et al., 2024, arXiv:2403.10444), which is lossless " + "and accepts at least as many tokens. Linear-chain one-model paths " + "only; ragged (confidence-scheduled) steps use 'standard'.") + allow_advanced_sampling: bool = Field( default=False, status="deprecated", @@ -7241,6 +7251,14 @@ def validate_speculative_config(self): # DraftTarget one-model) share the SpecMetadata slot-indexed # dispatch. + if self.speculative_config.rejection_sample_method == "block" and ( + not self.speculative_config.use_rejection_sampling + or getattr(self.speculative_config, "use_dynamic_tree", + False)): + raise ValueError( + "rejection_sample_method='block' requires " + "use_rejection_sampling=True and no dynamic tree.") + if isinstance(self.speculative_config, PARDDecodingConfig): assert self.speculative_config.max_draft_len > 0, "PARD max_draft_len must be > 0" diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 4b6941b8f505..57f2f51678dd 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1784,6 +1784,15 @@ "kind": "value", "path": "speculative_config.num_nextn_predict_layers" }, + { + "allowed_values": [ + "standard", + "block" + ], + "capture_policy": "literal", + "kind": "categorical", + "path": "speculative_config.rejection_sample_method" + }, { "capture_policy": "float", "kind": "value", diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash2_semantics.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash2_semantics.py index 0867979604ae..88a61d81392f 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash2_semantics.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash2_semantics.py @@ -621,47 +621,6 @@ def test_is_causal_is_ignored_for_a_plain_dflash_drafter(): assert wrapper._get_attention_mask_args(1) == (False, (-1, -1)) -# --------------------------------------------------------------------------- -# TRTLLM block-decode backend: no non-causal sliding window -# --------------------------------------------------------------------------- - - -def test_released_drafter_window_runs_on_two_sided_backends(): - """VANILLA and FA4 hand both window bounds to their kernels.""" - for backend in ("VANILLA", "FA4"): - _mask_wrapper(_all_sliding_config(2048), backend=backend).validate_block_attention_windows() - - -def test_trtllm_backend_rejects_the_released_drafter_window(): - """flashinfer's trtllm-gen context kernel refuses causal=False with any - finite window_left, so the released non-causal 2048 window cannot run.""" - wrapper = _mask_wrapper(_all_sliding_config(2048), backend="TRTLLM") - with pytest.raises(ValueError, match="non-causally within a 2048-token sliding window"): - wrapper.validate_block_attention_windows() - - -def test_trtllm_backend_accepts_causal_or_unwindowed_layers(): - # A causal sliding window is supported by the kernel. - _mask_wrapper( - _all_sliding_config(4, is_causal=True), backend="TRTLLM" - ).validate_block_attention_windows() - # No window at all, either. - _mask_wrapper( - SimpleNamespace(num_hidden_layers=2, layer_types=["full_attention"] * 2), - backend="TRTLLM", - ).validate_block_attention_windows() - - -def test_trtllm_window_check_reads_the_use_swa_override(): - """A use_swa window replaces the checkpoint-inferred one in the block - decode, so the check must resolve the same window the forward uses.""" - config = SimpleNamespace(num_hidden_layers=2, layer_types=["full_attention"] * 2) - wrapper = _mask_wrapper(config, backend="TRTLLM", layer_windows=[(3, 3), (3, 3)]) - assert wrapper._resolve_block_attention(0) == (False, (3, 3)) - with pytest.raises(ValueError, match="draft layer 0"): - wrapper.validate_block_attention_windows() - - # --------------------------------------------------------------------------- # Tiny end-to-end drafter: config parsing, weight loading, block-decode # parity vs an fp32 eager oracle (needs CUDA + flash_attn). diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py new file mode 100644 index 000000000000..91e7a90da967 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py @@ -0,0 +1,150 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Non-causal sliding-window DSpark block attention on the TRTLLM backend. + +The drafter serves it as the trtllm-gen causal windowed decode, with its log-sum-exp, plus +``dflash_merge_later_block_keys_``; it must equal dense non-causal attention. +""" + +import math +from unittest.mock import patch + +import pytest +import torch + +from tensorrt_llm._torch.models.modeling_dflash import ( + DFlashForCausalLM, + dflash_merge_later_block_keys_, +) + +HEAD_DIM, PAGE_SIZE = 128, 64 + + +def _trtllm_gen_ops(): + if not torch.cuda.is_available(): + pytest.skip("needs CUDA") + from tensorrt_llm._torch.speculative.dflash_attention import ( + dflash_trtllm_gen_unavailability_reason, + get_dflash_trtllm_gen_ops, + ) + + reason = dflash_trtllm_gen_unavailability_reason() + if reason is not None: + pytest.skip(f"DFlash TRTLLM backend unavailable: {reason}") + return get_dflash_trtllm_gen_ops() + + +def _check_block_attention(ctx_lens, block_size, num_heads, num_kv_heads, window): + """Runs the drafter's lowering over a shuffled page pool and compares it to dense fp32 attention.""" + ops = _trtllm_gen_ops() + dev = torch.device("cuda") + B = len(ctx_lens) + seq_lens = [c + block_size for c in ctx_lens] + num_pages = [math.ceil(s / PAGE_SIZE) for s in seq_lens] + max_pages = max(num_pages) + pages = torch.randperm(sum(num_pages), dtype=torch.int32).split(num_pages) + block_tables = torch.zeros(B, max_pages, dtype=torch.int32, device=dev) + cache = torch.zeros( + sum(num_pages), 2, num_kv_heads, PAGE_SIZE, HEAD_DIM, dtype=torch.bfloat16, device=dev + ) + keys = [torch.randn(s, num_kv_heads, HEAD_DIM, device=dev).bfloat16() for s in seq_lens] + values = [torch.randn(s, num_kv_heads, HEAD_DIM, device=dev).bfloat16() for s in seq_lens] + for b in range(B): + block_tables[b, : num_pages[b]] = pages[b] + for pos in range(0, ctx_lens[b], PAGE_SIZE): + n = min(PAGE_SIZE, ctx_lens[b] - pos) + page = int(pages[b][pos // PAGE_SIZE]) + cache[page, 0, :, :n] = keys[b][pos : pos + n].transpose(0, 1) + cache[page, 1, :, :n] = values[b][pos : pos + n].transpose(0, 1) + query = torch.randn(B, block_size, num_heads, HEAD_DIM, device=dev).bfloat16() + block_key = torch.stack([k[-block_size:] for k in keys]) + block_value = torch.stack([v[-block_size:] for v in values]) + seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=dev) + ops.append_paged_kv_cache( + append_key=block_key.flatten(0, 1), + append_value=block_value.flatten(0, 1), + batch_indices=torch.arange(B, dtype=torch.int32, device=dev).repeat_interleave(block_size), + positions=(seq_lens_t.view(-1, 1) - block_size + torch.arange(block_size, device=dev)) + .flatten() + .int(), + paged_kv_cache=cache, + kv_indices=block_tables.flatten(), + kv_indptr=torch.arange(0, (B + 1) * max_pages, max_pages, dtype=torch.int32, device=dev), + kv_last_page_len=(seq_lens_t - 1) % PAGE_SIZE + 1, + kv_layout="HND", + ) + drafter = DFlashForCausalLM.__new__(DFlashForCausalLM) + torch.nn.Module.__init__(drafter) + drafter._dflash_trtllm_gen_ops = ops + drafter._dflash_trtllm_gen_device = drafter._dflash_trtllm_gen_sm_count = None + drafter._dflash_trtllm_gen_workspace = drafter._dflash_trtllm_gen_counters = None + with patch.object(DFlashForCausalLM, "_prepare_dflash_index_buffers"): + drafter._prepare_dflash_trtllm_gen_buffers( + torch.bfloat16, dev, B, block_size, num_heads, num_kv_heads, HEAD_DIM + ) + + w = window - 1 + out = torch.empty(B * block_size, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=dev) + _, lse = ops.batch_decode_with_kv_cache( + query=query.flatten(0, 1), + kv_cache=(cache[:, 0], cache[:, 1]), + workspace_buffer=drafter._dflash_trtllm_gen_workspace, + block_tables=block_tables, + seq_lens=seq_lens_t, + max_seq_len=max_pages * PAGE_SIZE, + bmm1_scale=HEAD_DIM**-0.5, + bmm2_scale=1.0, + window_left=w, + out=out, + sinks=None, + enable_pdl=False, + kv_layout="HND", + backend="trtllm-gen", + q_len_per_req=block_size, + max_q_len=None, + cum_seq_lens_q=None, + kv_cache_sf=None, + uses_shared_paged_kv_idx=True, + bmm1_scale_log2=None, + return_lse=True, + multi_ctas_kv_counter_buffer=drafter._dflash_trtllm_gen_counters, + ) + dflash_merge_later_block_keys_(out, lse, query, block_key, block_value, HEAD_DIM**-0.5, w) + + group = num_heads // num_kv_heads + for b, c in enumerate(ctx_lens): + pos_q = c + torch.arange(block_size, device=dev).view(-1, 1) + pos_k = torch.arange(c + block_size, device=dev).view(1, -1) + logits = torch.einsum( + "jhd,lhd->hjl", query[b].float(), keys[b].float().repeat_interleave(group, 1) + ) + logits = logits.masked_fill((pos_k - pos_q).abs() > w, float("-inf")) * HEAD_DIM**-0.5 + expected = torch.einsum( + "hjl,lhd->jhd", logits.softmax(-1), values[b].float().repeat_interleave(group, 1) + ) + got = out[b * block_size : (b + 1) * block_size].float() + torch.testing.assert_close(got, expected, atol=2e-2, rtol=2e-2) + + +@pytest.mark.parametrize("block_size", [3, 7]) +@pytest.mark.parametrize("heads", [(16, 1), (64, 4)]) +def test_non_causal_window_matches_dense_attention(block_size, heads): + torch.manual_seed(block_size) + window = 64 + ctx_lens = [ + 1, + window - block_size - 1, + window - block_size, + window - 1, + window, + window + 1, + 3 * window + 17, + ] + _check_block_attention(ctx_lens, block_size, *heads, window) + + +def test_workspace_holds_the_lse_stats(): + """At 64 heads, 512 requests and 3 block tokens, the decode's softmax stats alone outgrow the + plain trtllm-gen workspace, which flashinfer would refuse.""" + torch.manual_seed(0) + _check_block_attention([1 + (97 * b) % 300 for b in range(512)], 3, 64, 4, 4096) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sampling.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sampling.py new file mode 100644 index 000000000000..870e76cf3c07 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sampling.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DSpark ancestral drafting and block verification under rejection sampling.""" + +import math +import types +from unittest.mock import patch + +import torch + +from tensorrt_llm._torch.models.modeling_speculative import dspark_markov_step_bias +from tensorrt_llm._torch.speculative.block_verification import block_verify +from tensorrt_llm._torch.speculative.dspark import DSparkWorker +from tensorrt_llm._torch.speculative.interface import SpecWorkerBase + +VOCAB, K = 3, 3 + + +def _prefix_model(generator, sharpness): + """A next-token distribution for every prefix of length 0..K.""" + num_prefixes = sum(VOCAB**m for m in range(K + 1)) + logits = sharpness * torch.randn(num_prefixes, VOCAB, generator=generator, dtype=torch.float64) + return torch.softmax(logits, -1) + + +def _rows(model, prefixes): + ids = torch.zeros(prefixes.shape[0], dtype=torch.int64) + for m in range(prefixes.shape[1]): + ids = ids * VOCAB + prefixes[:, m] + return model[ids + sum(VOCAB**m for m in range(prefixes.shape[1]))] + + +def test_block_verification_is_lossless(): + """Blocks verified, then completed by sampling the target, are distributed as the target's own samples.""" + generator = torch.Generator().manual_seed(0) + target, draft = _prefix_model(generator, 1.5), _prefix_model(generator, 1.0) + num = 200_000 + drafts = torch.empty(num, 0, dtype=torch.int64) + draft_rows = [] + for _ in range(K): + draft_rows.append(_rows(draft, drafts)) + drafts = torch.cat([drafts, torch.multinomial(draft_rows[-1], 1, generator=generator)], 1) + target_rows = [_rows(target, drafts[:, :j]) for j in range(K + 1)] + tokens, num_accepted = block_verify( + drafts, + torch.stack(draft_rows, 1).float(), + torch.stack(target_rows, 1).float(), + torch.rand(num, K, generator=generator, dtype=torch.float64), + lambda weights: torch.multinomial(weights, 1, generator=generator).squeeze(-1), + ) + + sequences = torch.zeros(num, K + 1, dtype=torch.int64) + for m in range(K + 1): + continued = torch.multinomial(_rows(target, sequences[:, :m]), 1, generator=generator) + sequences[:, m] = torch.where( + m < num_accepted.long(), tokens[:, m].long(), continued.squeeze(-1) + ) + grid = torch.cartesian_prod(*[torch.arange(VOCAB)] * (K + 1)) + expected = torch.full((grid.shape[0],), float(num), dtype=torch.float64) + for m in range(K + 1): + expected *= _rows(target, grid[:, :m]).gather(-1, grid[:, m : m + 1]).squeeze(-1) + ids = (sequences * VOCAB ** torch.arange(K, -1, -1)).sum(-1) + observed = torch.bincount(ids, minlength=grid.shape[0]).double() + big = expected >= 5 # pool the rare sequences + observed = torch.cat([observed[big], observed[~big].sum().view(1)]) + expected = torch.cat([expected[big], expected[~big].sum().view(1)]) + stat = ((observed - expected) ** 2 / expected).sum().item() + dof = observed.numel() - 1 + critical = dof * (1 - 2 / (9 * dof) + 3.719 * math.sqrt(2 / (9 * dof))) ** 3 # p = 1e-4 + assert stat < critical, f"chi-square {stat:.1f} >= {critical:.1f}" + + +def test_ancestral_drafts_condition_each_position_on_the_drawn_token(): + """Position k is drawn, with the temperature only, from base_k plus the Markov bias of the token drawn before it; + the distribution it was drawn from is stored as its proposal.""" + generator = torch.Generator().manual_seed(1) + num_gens, vocab, rank = 4, 32, 8 + drafter = types.SimpleNamespace( + has_markov_head=True, + markov_w1=torch.randn(vocab, rank, generator=generator), + markov_w2=torch.randn(vocab, rank, generator=generator), + ) + base = torch.randn(num_gens, K, vocab, generator=generator) + anchors = torch.randint(vocab, (num_gens,), generator=generator) + meta = types.SimpleNamespace( + request_temperatures=torch.full((num_gens,), 0.7), + request_seeds=torch.arange(num_gens, dtype=torch.int64), + request_offsets=torch.zeros(num_gens, dtype=torch.int64), + rng_draft_slot=lambda step: K + 1 + step, + batch_slot_ids=torch.arange(num_gens).flip(0), + draft_probs=torch.zeros(num_gens, K, vocab), + draft_probs_last_dim=0, + dummy_slot_row=0, + advanced_sampling_mode=None, + is_all_greedy_sample=False, + wants_advanced_draft_sampling=True, + draft_vocab_size=vocab, + vocab_size=vocab, + ) + calls = [] + + def sampler(mode, logits, temperatures, top_ks, top_ps, min_ps, *, seed, offset): + probs = torch.softmax(logits / temperatures[:, None], -1) + calls.append((logits, probs, top_ks, top_ps, min_ps)) + return torch.multinomial(probs, 1, generator=generator).squeeze(-1).int(), probs + + worker = DSparkWorker.__new__(DSparkWorker) + torch.nn.Module.__init__(worker) + worker.mapping = None + assert worker._samples_block_ancestrally(drafter, meta) + with patch.object(SpecWorkerBase, "_sample_from_logits_with_probs", staticmethod(sampler)): + tokens, stored_vocab = worker._sample_block_ancestrally( + drafter, base, {"first_prev_tokens": anchors}, meta, 0, num_gens + ) + + assert tokens.dtype == torch.int32 and stored_vocab == vocab == meta.draft_probs_last_dim + prev = anchors + for step, (logits, probs, top_ks, top_ps, min_ps) in enumerate(calls): + bias = dspark_markov_step_bias(prev, drafter.markov_w1, drafter.markov_w2) + torch.testing.assert_close(logits, base[:, step] + bias) + assert ( + (top_ks == torch.iinfo(torch.int32).max).all() + and (top_ps == 1).all() + and (min_ps == 0).all() + ) + torch.testing.assert_close(meta.draft_probs[meta.batch_slot_ids, step], probs) + prev = tokens[:, step] + assert len(calls) == K diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py index 0249d62b4efa..ba8383283466 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py @@ -362,6 +362,7 @@ def test_rejection_preserves_sampling_params_and_caps_acceptance(monkeypatch, ra widths = [2, 4] if ragged else [4, 4] metadata = _penalty_mapping_meta([4, 2, 9], widths if ragged else None) metadata.rng_accept_slot = 2 * worker.max_draft_len + 1 + metadata.rejection_sample_method = "standard" metadata.dummy_slot_row = 9 rows = 1 + sum(widths) metadata.temperatures = torch.arange(rows, dtype=torch.float32) + 0.5 @@ -507,6 +508,7 @@ def rejection(**kwargs): lens = torch.tensor(local_windows * trials, dtype=torch.int32) metadata = _penalty_mapping_meta([5, *slots.tolist()], lens.tolist() if ragged else None) metadata.rng_accept_slot = 2 * worker.max_draft_len + 1 + metadata.rejection_sample_method = "standard" metadata.dummy_slot_row = 9 metadata.runtime_draft_len = width metadata.enable_penalty = False