[TRTLLM-17389][feat] Serve MiniMax-M3.1-Flash with its bundled DSpark drafter

docs/source/models/supported-models.md

@@ -83,7 +83,7 @@ statuses for the same architecture in the two matrices.
 | `Gemma4ForConditionalGeneration` | Yes                | Yes        | Yes                        | No                    | Yes             | MTP                        | Yes           | No             | Yes                      | Yes                   | Yes             |
 | `Gemma4UnifiedForConditionalGeneration` | Yes                | Yes        | Yes                        | No                    | Yes             | No                         | Yes           | No             | Yes                      | Yes                   | Yes             |
 | `Step3p7ForConditionalGeneration`| Yes               | Yes        | Yes                        | Untested              | Untested        | MTP                        | Yes           | No             | Yes                      | Yes                   | Yes             |
-| `MiniMaxM3SparseForConditionalGeneration` [^12] | Yes               | Yes        | Yes                        | Untested              | Untested        | EAGLE-3 (Linear)           | Yes           | No             | N/A                      | Yes                   | Yes             |
+| `MiniMaxM3SparseForConditionalGeneration` [^12] | Yes               | Yes        | Yes                        | Untested              | Untested        | EAGLE-3 (Linear), DSpark [^21] | Yes           | No             | N/A                      | Yes                   | Yes             |
 
 [^1]: Chunked Prefill for MLA can only be enabled on SM90/SM100/SM103/SM107/SM120.
 [^2]: KV cache reuse for MLA can only be enabled on SM90/SM100/SM103/SM107/SM120/SM121 and in BF16/FP8 KV cache dtype.
@@ -101,6 +101,7 @@ statuses for the same architecture in the two matrices.
 [^18]: NGram and standalone Suffix Automaton (SA) use model-free drafting on the PyTorch backend, so they are not listed in individual entries. This does not imply universal end-to-end support: compatibility depends on each model's multi-token verification and cache-management paths and may be untested or explicitly restricted.
 [^19]: KV cache reuse for hybrid recurrent-attention models requires an explicit recurrent-state snapshot policy, such as `kv_cache_config.mamba_state_config.periodic_snapshot_interval`; the model default disables reuse when no snapshot policy is configured.
 [^20]: Supports text, image, and video inputs, MTP (including with attention data parallelism), and FP8 KV cache. Image and video inputs require `transformers==5.17.0`. Beam search is not supported. See the [GLM-5.3-Flash deployment guide](../deployment-guide/deployment-guide-for-glm-5.3-flash-on-trtllm.md) for setup and feature-specific requirements.
+[^21]: MiniMax-M3.1-Flash ships a standalone DSpark drafter in its `dspark/` directory (architecture `DSparkMiniMaxDraftModel`). Load it with `speculative_config: {decoding_type: DSpark, speculative_model: <checkpoint>/dspark}`. Keep KV cache block reuse enabled (the default): with `kv_cache_config.enable_block_reuse: false`, or on a disaggregated generation server, prompts of about 131k tokens or more fail with `Draft KV cache context resize failed`.
 
 # Encoder-Decoder Feature Support Matrix (PyTorch Backend)
 

tensorrt_llm/_torch/models/_arch_index.py

@@ -45,6 +45,7 @@ def is_builtin_zoo_module(module_name: str) -> bool:
     "Cosmos3EdgeForConditionalGeneration": "modeling_cosmos3_edge",
     "Cosmos3ForConditionalGeneration": "modeling_cosmos3",
     "DeciLMForCausalLM": "modeling_nemotron_nas",
+    "DSparkMiniMaxDraftModel": "modeling_dspark",
     "DeepseekV32ForCausalLM": "modeling_deepseekv3",
     "DeepseekV3ForCausalLM": "modeling_deepseekv3",
     "DeepseekV41ForCausalLM": "modeling_deepseekv41",

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
 
@@ -22,6 +23,7 @@
 from ..distributed import AllReduce
 from ..pyexecutor.config_utils import _is_sliding_attention_layer, get_layer_attention_window
 from ..speculative.dflash_attention import (
+    DFLASH_CONTEXT_LOWERING,
     dflash_trtllm_gen_unavailability_reason,
     get_dflash_fa4_fwd,
     get_dflash_flash_attention,
@@ -32,6 +34,11 @@
 from ..speculative.spec_tail_kernels import DraftBlockPlacement
 from .modeling_utils import FUSED_MODULE_COMPONENTS, get_model_architecture, register_draft_model
 
+# trtllm-gen reports the log-sum-exp in base 2; the merge and the MLA fixup
+# (modeling_dspark) multiply by this to get the natural-log LSE flashinfer's
+# merge expects. One constant, so a flashinfer convention change is fixed once.
+TRTLLM_GEN_LSE_LN2 = math.log(2.0)
+
 
 def dspark_layer_window_size(
     use_swa: bool, swa_window: int, layer_types, layer_idx: int
@@ -65,6 +72,129 @@ def dspark_layer_window_size(
     return (swa_window - 1, swa_window - 1)
 
 
+def dflash_merge_future_block_attention(
+    out_causal: torch.Tensor,
+    lse_causal: torch.Tensor,
+    query: torch.Tensor,
+    block_key: torch.Tensor,
+    block_value: torch.Tensor,
+    scale: float,
+    window_right: int,
+) -> torch.Tensor:
+    """Extend a causal block attention to a non-causal one, exactly.
+
+    A non-causal block query ``j`` attends the keys a causal one does (context
+    within the left window, block keys ``<= j``) plus the block keys after it,
+    ``j < k <= j + window_right``. Given the causal output and its natural-log
+    log-sum-exp, this adds that second key set's softmax mass and renormalizes:
+    the result is the non-causal attention up to rounding. It is how the TRTLLM
+    backend serves a non-causal sliding-window layer, which the trtllm-gen
+    context kernel rejects while its causal decode kernel windows natively.
+
+    Args:
+        out_causal: [B * blk, Hq, D] causal attention output.
+        lse_causal: [B * blk, Hq] float32 log-sum-exp of the scaled causal logits.
+        query: [B, blk, Hq, D] block queries, as the kernel saw them.
+        block_key / block_value: [B, blk, Hkv, D] block K/V, as appended to the cache.
+        scale: softmax scale the kernel applied to ``q @ k``.
+        window_right: inclusive right window in tokens; -1 means unbounded.
+    Returns:
+        [B * blk, Hq, D] in ``out_causal.dtype``.
+    """
+    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, head_dim)
+    logits = torch.einsum("bjhgd,bkhd->bhgjk", q, block_key.float()) * scale
+    offsets = torch.arange(blk, device=query.device)
+    distance = offsets.view(1, -1) - offsets.view(-1, 1)  # [j, k] = k - j
+    future = distance > 0
+    if window_right >= 0:
+        future &= distance <= window_right
+    logits = logits.masked_fill(~future, float("-inf"))
+    lse_future = torch.logsumexp(logits, dim=-1)  # [B, Hkv, g, blk], -inf for no future keys
+    lse_future_finite = torch.where(
+        torch.isfinite(lse_future), lse_future, torch.zeros_like(lse_future)
+    )
+    probs = torch.exp(logits - lse_future_finite.unsqueeze(-1))
+    out_future = torch.einsum("bhgjk,bkhd->bjhgd", probs, block_value.float())
+    out_future = out_future.reshape(B * blk, num_heads, head_dim)
+    lse_future = lse_future.permute(0, 3, 1, 2).reshape(B * blk, num_heads)
+
+    lse_total = torch.logaddexp(lse_causal, lse_future)
+    weight_causal = torch.exp(lse_causal - lse_total).unsqueeze(-1)
+    weight_future = torch.exp(lse_future - lse_total).unsqueeze(-1)
+    merged = weight_causal * out_causal.float() + weight_future * out_future
+    return merged.to(out_causal.dtype)
+
+
+def dflash_trtllm_windowed_block_attention(
+    trtllm_gen_ops,
+    query: torch.Tensor,
+    block_key: torch.Tensor,
+    block_value: torch.Tensor,
+    layer_cache: torch.Tensor,
+    block_tables: torch.Tensor,
+    seq_lens: torch.Tensor,
+    max_seq_len: int,
+    causal: bool,
+    window_size: tuple[int, int],
+    workspace: torch.Tensor,
+    counters: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    """Block attention through the trtllm-gen decode kernel, into ``out``.
+
+    Serves a causal layer (windowed or not) and a non-causal windowed one; a
+    non-causal unwindowed layer takes the context kernel instead. trtllm-gen
+    windows only causal attention, so the non-causal windowed case runs the
+    causal decode with its log-sum-exp and then adds the block keys after each
+    query (:func:`dflash_merge_future_block_attention`).
+
+    Args:
+        query: [B, blk, Hq, D] block queries.
+        block_key / block_value: [B, blk, Hkv, D], already appended to ``layer_cache``.
+        layer_cache: [pages, 2, Hkv, page_size, D] HND pages of this layer.
+        seq_lens: [B] context plus block tokens per request.
+        window_size: (left, right) inclusive token bounds, -1 for none.
+        out: [B * blk, Hq, D] output.
+    """
+    B, block_size, num_heads, head_dim = query.shape
+    window_left, window_right = window_size
+    merge_future = not causal and window_left >= 0
+    decoded = trtllm_gen_ops.batch_decode_with_kv_cache(
+        query=query.reshape(-1, num_heads, head_dim),
+        kv_cache=(layer_cache[:, 0], layer_cache[:, 1]),
+        workspace_buffer=workspace,
+        block_tables=block_tables,
+        seq_lens=seq_lens,
+        max_seq_len=max_seq_len,
+        bmm1_scale=head_dim**-0.5,
+        bmm2_scale=1.0,
+        window_left=window_left,
+        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=merge_future,
+        multi_ctas_kv_counter_buffer=counters,
+    )
+    if merge_future:
+        lse = decoded[1] * TRTLLM_GEN_LSE_LN2
+        out.copy_(
+            dflash_merge_future_block_attention(
+                out, lse, query, block_key, block_value, head_dim**-0.5, window_right
+            )
+        )
+
+
 def declares_lilicorr(config: PretrainedConfig) -> bool:
     """Whether the checkpoint requires the LiLiCorr candidate scorer."""
     return is_lilicorr_checkpoint(
@@ -1364,29 +1494,40 @@ 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.
+    def validate_block_attention_windows(self) -> list[str]:
+        """Report how the TRTLLM block decode serves each draft layer's mask.
 
         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.
+        rejects ``causal=False`` with any finite ``window_left``. A non-causal
+        windowed layer runs the causal windowed decode plus the block keys
+        after each query (``dflash_merge_future_block_attention``), which is
+        the same attention up to rounding.
+
+        Returns:
+            The lowering per draft layer; empty for the other backends, which
+            take both window bounds directly.
         """
         if self.dflash_attention_backend != "TRTLLM":
-            return
+            return []
         num_layers = getattr(self.config, "num_hidden_layers", None)
         if num_layers is None:
             num_layers = len(self.model.layers)
+        lowerings = []
         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."
+            causal, window_size = self._resolve_block_attention(layer_idx)
+            if causal:
+                lowering = "causal decode"
+            elif window_size[0] >= 0:
+                lowering = "causal decode + future-block merge"
+            else:
+                lowering = DFLASH_CONTEXT_LOWERING
+            logger.info_once(
+                f"DFlash draft layer {layer_idx}: causal={causal}, "
+                f"window={window_size}, TRTLLM lowering: {lowering}.",
+                key=f"dflash_block_attention_{id(self)}_{layer_idx}",
             )
+            lowerings.append(lowering)
+        return lowerings
 
     def _prepare_dflash_trtllm_gen_buffers(
         self,
@@ -1760,29 +1901,21 @@ 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(
-                        query=q_flat,
-                        kv_cache=(layer_cache[:, 0], layer_cache[:, 1]),
-                        workspace_buffer=self._dflash_trtllm_gen_workspace,
-                        block_tables=block_tables,
-                        seq_lens=seq_lens_after,
-                        max_seq_len=pages_per_slot * page_size,
-                        bmm1_scale=head_dim**-0.5,
-                        bmm2_scale=1.0,
-                        window_left=window_left,
-                        out=out_flat,
-                        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,
-                        multi_ctas_kv_counter_buffer=self._dflash_trtllm_gen_counters,
+                if causal or window_left >= 0:
+                    dflash_trtllm_windowed_block_attention(
+                        trtllm_gen_ops,
+                        Q_bshd,
+                        k_noise_bshd,
+                        v_noise_bshd,
+                        layer_cache,
+                        block_tables,
+                        seq_lens_after,
+                        pages_per_slot * page_size,
+                        causal,
+                        window_size,
+                        self._dflash_trtllm_gen_workspace,
+                        self._dflash_trtllm_gen_counters,
+                        out_flat,
                     )
                 else:
                     cum_seq_lens_q = torch.arange(