[None][fix] MiniMax-M3.1 DSpark AL parity: non-causal draft blocks + opt-in golden sampling

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,

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)

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)

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,

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