From 54fd8b8e8a8a6a8717a2378f0b5974e4e370e500 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Fri, 23 Oct 2026 13:12:45 -0700 Subject: [PATCH 1/5] [None][feat] Serve MiniMax-M3.1-Flash with its bundled DSpark drafter The drafter's VL-shaped config.json is flattened by architecture, with its sliding-window layers mapped to the DFlash use_swa spelling. Its checkpoint tensors are renamed out of their namespace, MXFP8-dequantized and Gemma norm-folded at load, and run on the GQA DSpark block decode. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/models/_arch_index.py | 1 + tensorrt_llm/_torch/models/modeling_dspark.py | 259 +++++++++++++++++- .../_torch/pyexecutor/config_utils.py | 76 +++++ .../hw_agnostic/test_minimax_dspark_draft.py | 242 ++++++++++++++++ 4 files changed, 571 insertions(+), 7 deletions(-) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_minimax_dspark_draft.py diff --git a/tensorrt_llm/_torch/models/_arch_index.py b/tensorrt_llm/_torch/models/_arch_index.py index 2b4d200a35ad..42cea89baf97 100644 --- a/tensorrt_llm/_torch/models/_arch_index.py +++ b/tensorrt_llm/_torch/models/_arch_index.py @@ -43,6 +43,7 @@ def is_builtin_zoo_module(module_name: str) -> bool: "Cohere2ForCausalLM": "modeling_cohere2", "Cosmos3ForConditionalGeneration": "modeling_cosmos3", "DeciLMForCausalLM": "modeling_nemotron_nas", + "DSparkMiniMaxDraftModel": "modeling_dspark", "DeepseekV32ForCausalLM": "modeling_deepseekv3", "DeepseekV3ForCausalLM": "modeling_deepseekv3", "DeepseekV41ForCausalLM": "modeling_deepseekv41", diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index ef6f20365727..619519d1f73c 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -83,20 +83,22 @@ class or the edge would become a cycle. import os import re from functools import lru_cache -from typing import Dict, List, NamedTuple, Optional +from typing import Any, Dict, List, NamedTuple, Optional import torch import torch.nn.functional as F from torch import nn from transformers import PretrainedConfig -from tensorrt_llm.functional import RopeEmbeddingUtils +from tensorrt_llm.functional import PositionEmbeddingType, RopeEmbeddingUtils from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping from tensorrt_llm.quantization.mode import QuantAlgo from ..._utils import get_sm_version, is_sm_100f from ...models.modeling_utils import QuantConfig +from ..attention.attention import Attention +from ..attention.backends.interface import PositionalEmbeddingParams, RopeParams from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from ..distributed import AllReduceParams from ..model_config import ModelConfig @@ -106,7 +108,7 @@ class or the edge would become a cycle. from ..modules.linear import Linear, TensorParallelMode from ..modules.mhc.hyper_connection import HCHead from ..modules.rms_norm import RMSNorm -from ..pyexecutor.config_utils import is_mla +from ..pyexecutor.config_utils import MINIMAX_DSPARK_DRAFT_ARCH, is_mla from ..speculative.interface import SpeculativeDecodingMode from ..utils import AuxStreamType from .modeling_deepseekv4 import ( @@ -125,7 +127,12 @@ class or the edge would become a cycle. build_markov_head, dspark_markov_chain_logits, ) -from .modeling_utils import DecoderModel, register_auto_model, register_draft_model +from .modeling_utils import ( + DecoderModel, + DecoderModelForCausalLM, + register_auto_model, + register_draft_model, +) if IS_CUTLASS_DSL_AVAILABLE: from ..custom_ops.dspark_attention_custom_op import ( @@ -3505,6 +3512,230 @@ def load_weights_from_target_model(self, target_model: torch.nn.Module) -> None: self.lm_head = target_model.lm_head +# ---------------------------------------------------------------------------- +# MiniMax-M3 standalone DSpark drafter (architecture ``DSparkMiniMaxDraftModel``). +# +# A GQA backbone (64 heads / 4 KV heads / head_dim 128 at hidden 6144), per-head +# Gemma Q/K norm, NeoX RoPE over the first ``rotary_dim`` channels of each head, +# a dense swiglu-oai MLP and Gemma boundary norms: the dense-layer block of the +# MiniMax-M3 target. Weights ship MXFP8 (e4m3 + UE8M0 1x32 block scales) under +# ``language_model.model.dspark.layers..decoder_layer.*``; the heads are +# ``language_model.model.dspark.{fc,hidden_norm,final_norm,markov_head.*, +# confidence_head.proj.*}``. +# ---------------------------------------------------------------------------- + +MINIMAX_DSPARK_WEIGHT_PREFIX = "language_model.model.dspark." +_MINIMAX_DSPARK_DECODER_LAYER_INFIX = ".decoder_layer." +_MINIMAX_DSPARK_RENAMES = {"final_norm.weight": "norm.weight"} +# Every RMSNorm in this drafter is a Gemma norm, (1 + weight) * x. The backbone +# and the wrapper-owned ``hidden_norm`` are plain norms, so the stored weights +# are pre-incremented at load (the same fold the MiniMax-M3 target applies to +# its layer-boundary norms). +_MINIMAX_DSPARK_GEMMA_NORM_SUFFIXES = ( + "input_layernorm.weight", + "post_attention_layernorm.weight", + "q_norm.weight", + "k_norm.weight", + "hidden_norm.weight", +) +_MINIMAX_DSPARK_GEMMA_NORM_KEYS = ("norm.weight",) +_MXFP8_BLOCK_SIZE = 32 +_MXFP8_SCALE_SUFFIX = ".weight_scale_inv" + + +def remap_minimax_dspark_weight_name(name: str) -> str: + """Checkpoint tensor name -> the name the DFlash/DSpark loaders expect. + + Strips the ``language_model.model.dspark.`` namespace and the + ``.decoder_layer.`` level, and renames ``final_norm`` to ``norm``. Names + outside the namespace are returned unchanged. + """ + if not name.startswith(MINIMAX_DSPARK_WEIGHT_PREFIX): + return name + name = name[len(MINIMAX_DSPARK_WEIGHT_PREFIX) :] + name = name.replace(_MINIMAX_DSPARK_DECODER_LAYER_INFIX, ".") + return _MINIMAX_DSPARK_RENAMES.get(name, name) + + +def dequantize_mxfp8_weight( + weight: torch.Tensor, scale_inv: torch.Tensor, dtype: torch.dtype +) -> torch.Tensor: + """``weight`` (e4m3, ``[out, in]``) times ``2 ** (scale_inv - 127)`` per + ``[1, 32]`` block (``scale_inv``: UE8M0 as uint8, ``[out, in / 32]``), in + ``dtype``.""" + w = weight.to(torch.float32) + exponent = scale_inv.to(torch.int32) - 127 + scale = torch.exp2(exponent.to(torch.float32)) + scale = scale.repeat_interleave(_MXFP8_BLOCK_SIZE, dim=1)[:, : w.shape[1]] + return (w * scale).to(dtype) + + +def prepare_minimax_dspark_weights(weights: Dict, dtype: torch.dtype) -> Dict: + """Rename, dequantize and norm-fold a MiniMax DSpark drafter checkpoint. + + Returns a new dict keyed by the DFlash/DSpark loader names, with every + MXFP8 ``weight`` / ``weight_scale_inv`` pair replaced by one ``dtype`` + tensor and every Gemma norm weight incremented by one. + """ + renamed = {remap_minimax_dspark_weight_name(k): v for k, v in weights.items()} + out = {} + for name, value in renamed.items(): + if name.endswith(_MXFP8_SCALE_SUFFIX): + continue + tensor = value[:] if hasattr(value, "__getitem__") else value + scale_name = name[: -len(".weight")] + _MXFP8_SCALE_SUFFIX + if name.endswith(".weight") and scale_name in renamed: + scale = renamed[scale_name] + scale = scale[:] if hasattr(scale, "__getitem__") else scale + tensor = dequantize_mxfp8_weight(tensor, scale, dtype) + elif name.endswith(_MINIMAX_DSPARK_GEMMA_NORM_SUFFIXES) or name in ( + _MINIMAX_DSPARK_GEMMA_NORM_KEYS + ): + tensor = tensor + 1.0 + out[name] = tensor + return out + + +class MiniMaxDSparkAttention(Attention): + """GQA attention of the MiniMax DSpark drafter: fused ``qkv_proj`` + + ``o_proj`` with per-head ``q_norm`` / ``k_norm`` and partial NeoX RoPE. + + Only the projections, norms and RoPE parameters are read; the block decode + is ``DFlashForCausalLM.dflash_forward``. ``self.pretrained_config`` is what + that decode reads ``partial_rotary_factor`` from. + """ + + def __init__(self, model_config: ModelConfig, layer_idx: int) -> None: + config = model_config.pretrained_config + self.pretrained_config = config + head_dim = int( + getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + ) + super().__init__( + hidden_size=config.hidden_size, + num_attention_heads=config.num_attention_heads, + num_key_value_heads=config.num_key_value_heads, + max_position_embeddings=config.max_position_embeddings, + bias=False, + pos_embd_params=PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, + rope=RopeParams.from_config(config), + ), + rope_fusion=False, + layer_idx=layer_idx, + dtype=config.torch_dtype, + dense_bias=False, + config=model_config, + head_dim=head_dim, + ) + self.q_norm = RMSNorm( + hidden_size=head_dim, eps=config.rms_norm_eps, dtype=config.torch_dtype + ) + self.k_norm = RMSNorm( + hidden_size=head_dim, eps=config.rms_norm_eps, dtype=config.torch_dtype + ) + + +class MiniMaxDSparkDecoderLayer(DecoderLayer): + """One drafter layer: GQA projections + dense swiglu-oai MLP + the two norms.""" + + def __init__(self, model_config: ModelConfig, layer_idx: int) -> None: + super().__init__() + config = model_config.pretrained_config + self.layer_idx = layer_idx + self.self_attn = MiniMaxDSparkAttention(model_config, layer_idx) + intermediate_size = int( + getattr(config, "dspark_ffn_hidden_size", None) + or getattr(config, "dense_intermediate_size", config.intermediate_size) + ) + self.mlp = GatedMLP( + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + bias=False, + activation=torch.nn.functional.silu, + swiglu_alpha=float(getattr(config, "swiglu_alpha", 1.702)), + swiglu_beta=1.0, + swiglu_limit=float(getattr(config, "swiglu_limit", 7.0)), + dtype=config.torch_dtype, + config=model_config, + overridden_tp_size=1 if model_config.mapping.enable_attention_dp else None, + layer_idx=layer_idx, + ) + self.input_layernorm = RMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype + ) + self.post_attention_layernorm = RMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype + ) + + def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor: + raise NotImplementedError( + "MiniMaxDSparkDecoderLayer is driven by DFlashForCausalLM.dflash_forward, " + "not by the generic decoder loop." + ) + + +class MiniMaxDSparkModel(DecoderModel): + """Weight container for the MiniMax DSpark block decode.""" + + def __init__(self, model_config: ModelConfig) -> None: + super().__init__(model_config) + config = model_config.pretrained_config + # Replaced by the target's embedding in load_weights_from_target_model; + # the checkpoint ships none. + self.embed_tokens = Embedding( + config.vocab_size, + config.hidden_size, + dtype=config.torch_dtype, + mapping=model_config.mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=True, + ) + self.layers = nn.ModuleList( + [ + MiniMaxDSparkDecoderLayer(model_config, layer_idx) + for layer_idx in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype + ) + + def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor: + raise NotImplementedError( + "MiniMaxDSparkModel is a weight container for the DSpark block decode; " + "see DFlashForCausalLM.dflash_forward." + ) + + +@register_auto_model("DSparkMiniMaxDraftModel") +class MiniMaxDSparkBackbone(DecoderModelForCausalLM[MiniMaxDSparkModel, PretrainedConfig]): + """Registry entry for the MiniMax DSpark drafter checkpoint + (``config.architectures == ["DSparkMiniMaxDraftModel"]``).""" + + def __init__(self, model_config: ModelConfig) -> None: + config = model_config.pretrained_config + super().__init__( + MiniMaxDSparkModel(model_config), + config=model_config, + hidden_size=config.hidden_size, + vocab_size=config.vocab_size, + ) + + +class MiniMaxDSparkForCausalLM(GQADSparkForCausalLM): + """DSpark drafter on the MiniMax-M3 GQA backbone, from a standalone checkpoint. + + The GQA block decode and head set are inherited; this class only translates + the checkpoint (namespace, MXFP8 dequantization, Gemma norm fold) before + the DFlash loader sees it. See :func:`prepare_minimax_dspark_weights`. + """ + + def load_weights(self, weights: Dict, weight_mapper: Any = None, **kwargs: Any) -> Any: + weights = prepare_minimax_dspark_weights(weights, self.config.torch_dtype) + return super().load_weights(weights, weight_mapper=weight_mapper, **kwargs) + + def draft_is_embedded_in_target(model_config) -> bool: """True when the DSpark draft weights live inside the target checkpoint. @@ -3562,14 +3793,21 @@ def _build_dspark_draft(model_config, draft_config, lm_head, model): # what differs between the two wrappers, and everything below it # (``DFlashForCausalLM.__init__``) already resolves the backbone through the # model registry, so keying on model_type a second time would force a new - # entry for every GQA family that already works. - if is_mla(draft_config.pretrained_config): + # entry for every GQA family that already works. The MiniMax drafter is the + # one architecture-keyed exception: its checkpoint needs a translation step + # (namespace, MXFP8, Gemma norms) that the generic GQA loader has no hook for. + draft_arches = getattr(draft_config.pretrained_config, "architectures", None) or [] + if draft_arches and draft_arches[0] == MINIMAX_DSPARK_DRAFT_ARCH: + drafter_cls = MiniMaxDSparkForCausalLM + elif is_mla(draft_config.pretrained_config): return MLADSparkForCausalLM( draft_config, dflash_attention_backend=model_config.spec_config.attention_backend, ) + else: + drafter_cls = GQADSparkForCausalLM # The drafter's own ModelConfig has no spec_config, so the worker's mask id is handed over here. - return GQADSparkForCausalLM( + return drafter_cls( draft_config, dflash_attention_backend=model_config.spec_config.attention_backend, spec_mask_token_id=model_config.spec_config.mask_token_id, @@ -3584,6 +3822,13 @@ def _build_dspark_draft(model_config, draft_config, lm_head, model): # Standalone flavour. "GQADSparkForCausalLM", "MLADSparkForCausalLM", + # MiniMax-M3 GQA drafter and its checkpoint translation. + "MiniMaxDSparkForCausalLM", + "MiniMaxDSparkBackbone", + "MiniMaxDSparkModel", + "prepare_minimax_dspark_weights", + "remap_minimax_dspark_weight_name", + "dequantize_mxfp8_weight", # MLA drafter backbone (a weight container for the block decode). "K3DsparkForCausalLM", "K3DsparkModel", diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index 24b094deb4d4..097e73aae71f 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -972,6 +972,78 @@ def match_nemotron_h_layer_types(model_config, layer_types): return [_NEMOTRON_H_LEGACY_LAYER_TYPES.get(t, t) for t in layer_types] +MINIMAX_DSPARK_DRAFT_ARCH = "DSparkMiniMaxDraftModel" + +# Keys of the standalone MiniMax DSpark drafter config.json that describe the +# target checkpoint it was exported next to, not the drafter itself. The +# drafter's MXFP8 linears are dequantized to the compute dtype at load, so the +# quantization block is dropped and the backbone is built unquantized. +_MINIMAX_DSPARK_DROPPED_KEYS = frozenset({ + "text_config", + "vision_config", + "quantization_config", + "architectures", + "model_type", + "_name_or_path", +}) + + +def build_minimax_dspark_draft_config( + config_dict: dict) -> transformers.PretrainedConfig: + """Flat drafter config for a standalone MiniMax-M3 DSpark checkpoint. + + The drafter ships a MiniMax VL-shaped config.json: ``model_type`` + ``minimax_m3_vl`` with the backbone geometry nested under ``text_config``, + no ``auto_map`` and no modeling code, so ``AutoConfig.from_pretrained`` + cannot resolve it. The drafter's own fields (``dspark_*``, + ``dflash_config``, ``hidden_size``, ``vocab_size``, ...) are duplicated + at the top level and inside ``text_config``; the text side wins where both + are present. The result is a plain ``PretrainedConfig`` with every field as + an attribute, which is what ``DFlashForCausalLM`` and the registered + ``DSparkMiniMaxDraftModel`` backbone read. + + ``dflash_config`` gets ``use_swa`` / ``swa_window_size`` when the drafter + declares ``sliding_attention`` layers with a ``sliding_window`` but not the + DFlash spelling: the block decode keys its per-layer window on the former, + and the HF reference these drafters are distilled with applies the window + on exactly those layers. + """ + flat = { + k: v + for k, v in config_dict.items() if k not in _MINIMAX_DSPARK_DROPPED_KEYS + } + text = config_dict.get("text_config") or {} + flat.update({ + k: v + for k, v in text.items() if k not in _MINIMAX_DSPARK_DROPPED_KEYS + }) + dflash = dict(flat.get("dflash_config") or {}) + layer_types = dflash.get("layer_types") or flat.get("layer_types") or [] + sliding_window = dflash.get("sliding_window") or flat.get("sliding_window") + if ("use_swa" not in dflash and sliding_window + and "sliding_attention" in layer_types): + dflash["use_swa"] = True + dflash["swa_window_size"] = int(sliding_window) + dflash.setdefault("layer_types", list(layer_types)) + flat["dflash_config"] = dflash + flat["architectures"] = [MINIMAX_DSPARK_DRAFT_ARCH] + # Not ``minimax_m3_vl``: that model_type selects the VL target's + # text_config unwrapping and input-processor registration, neither of + # which applies to the flattened drafter. + flat["model_type"] = "minimax_dspark_draft" + dtype = flat.pop("dtype", None) or flat.pop("torch_dtype", + None) or "bfloat16" + flat.pop("torch_dtype", None) + if isinstance(dtype, str): + dtype = str_dtype_to_torch(dtype) + config = transformers.PretrainedConfig() + for key, value in flat.items(): + setattr(config, key, value) + config.dtype = dtype + config.torch_dtype = dtype + return config + + def load_pretrained_config(model_name_or_path: str, trust_remote_code: bool = False, checkpoint_format: Optional[str] = None, @@ -1113,6 +1185,10 @@ def load_pretrained_config(model_name_or_path: str, vision_config = getattr(model_config, "vision_config", None) if getattr(vision_config, "num_hidden_layers", 0) > 0: model_config.architectures = ["DeepseekV41ForConditionalGeneration"] + elif architectures and architectures[0] == MINIMAX_DSPARK_DRAFT_ARCH: + # Standalone MiniMax-M3 DSpark drafter: VL-shaped config.json with + # no auto_map, flattened here (see build_minimax_dspark_draft_config). + model_config = build_minimax_dspark_draft_config(config_dict) elif model_type in _CONFIG_REGISTRY: config_class = _CONFIG_REGISTRY[model_type] model_config = config_class.from_pretrained(model_name_or_path, diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_minimax_dspark_draft.py b/tests/unittest/_torch/speculative/hw_agnostic/test_minimax_dspark_draft.py new file mode 100644 index 000000000000..3e562a454368 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_minimax_dspark_draft.py @@ -0,0 +1,242 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MiniMax-M3 standalone DSpark drafter: config flattening and checkpoint +translation (CPU only; no drafter module is built here).""" + +import json +import pathlib + +import pytest +import torch + +from tensorrt_llm._torch.models.modeling_dspark import ( + GQADSparkForCausalLM, + MiniMaxDSparkBackbone, + MiniMaxDSparkForCausalLM, + dequantize_mxfp8_weight, + prepare_minimax_dspark_weights, + remap_minimax_dspark_weight_name, +) +from tensorrt_llm._torch.models.modeling_utils import get_registered_model_class +from tensorrt_llm._torch.pyexecutor.config_utils import ( + MINIMAX_DSPARK_DRAFT_ARCH, + build_minimax_dspark_draft_config, + load_pretrained_config, +) + +# The shipped drafter config.json, reduced to the fields the loaders read. The +# backbone geometry sits under text_config; the dspark_* / dflash_config knobs +# are duplicated at both levels, as in the checkpoint. +_TEXT = { + "architectures": ["DSparkMiniMaxDraftModel"], + "dtype": "bfloat16", + "hidden_size": 6144, + "intermediate_size": 3072, + "num_hidden_layers": 5, + "num_attention_heads": 64, + "num_key_value_heads": 4, + "head_dim": 128, + "vocab_size": 200064, + "max_position_embeddings": 1048576, + "rms_norm_eps": 1e-06, + "use_gemma_norm": True, + "rope_theta": 10000000, + "rotary_dim": 64, + "partial_rotary_factor": 0.5, + "hidden_act": "swigluoai", + "use_qk_norm": True, + "dense_intermediate_size": 12288, + "qk_norm_type": "per_head", + "swiglu_alpha": 1.702, + "swiglu_limit": 7.0, + "sparse_attention_config": None, + "dspark_block_size": 7, + "dspark_markov_rank": 256, + "dspark_markov_head_type": "vanilla", + "dspark_noise_token_id": 200058, + "dspark_use_dense_ffn": True, + "dspark_ffn_hidden_size": 12288, + "dspark_target_layer_ids": [3, 17, 31, 45, 59], + "dspark_config": { + "markov_rank": 256, + "markov_head_type": "vanilla", + "mask_token_id": 200058, + "use_dense_ffn": True, + "ffn_hidden_size": 12288, + }, + "dflash_config": { + "num_hidden_layers": 5, + "num_target_layers": 60, + "block_size": 7, + "mask_token_id": 200058, + "target_layer_ids": [3, 17, 31, 45, 59], + "layer_types": ["sliding_attention"] * 5, + "sliding_window": 4096, + }, + "enable_confidence_head": True, + "layer_types": ["sliding_attention"] * 5, + "sliding_window": 4096, + "model_type": "", +} + +_CONFIG = { + "transformers_version": "5.12.1", + "architectures": ["DSparkMiniMaxDraftModel"], + "dtype": "bfloat16", + "model_type": "minimax_m3_vl", + "text_config": _TEXT, + "quantization_config": { + "quant_method": "mxfp8", + "weight_block_size": [1, 32], + "ignored_layers": ["lm_head"], + }, + "tie_word_embeddings": False, + "dspark_block_size": 7, + "dspark_target_layer_ids": [3, 17, 31, 45, 59], + "dflash_config": dict(_TEXT["dflash_config"]), + "hidden_size": 6144, + "vocab_size": 200064, + "num_hidden_layers": 5, +} + + +def test_flattened_config_exposes_backbone_and_drafter_fields() -> None: + cfg = build_minimax_dspark_draft_config(_CONFIG) + assert cfg.architectures == [MINIMAX_DSPARK_DRAFT_ARCH] + assert cfg.model_type != "minimax_m3_vl" + assert cfg.num_attention_heads == 64 + assert cfg.num_key_value_heads == 4 + assert cfg.head_dim == 128 + assert cfg.rotary_dim == 64 and cfg.partial_rotary_factor == 0.5 + assert cfg.dspark_target_layer_ids == [3, 17, 31, 45, 59] + assert cfg.dspark_markov_rank == 256 + assert cfg.dflash_config["mask_token_id"] == 200058 + assert cfg.dflash_config["block_size"] == 7 + # The HF sliding window on 'sliding_attention' layers is translated to the + # DFlash block-decode spelling. + assert cfg.dflash_config["use_swa"] is True + assert cfg.dflash_config["swa_window_size"] == 4096 + # Dequantized at load: the backbone is built unquantized. + assert getattr(cfg, "quantization_config", None) is None + assert not hasattr(cfg, "text_config") + assert cfg.torch_dtype == torch.bfloat16 + + +def test_flattened_config_keeps_existing_dflash_window_spelling() -> None: + config = json.loads(json.dumps(_CONFIG)) + config["text_config"]["dflash_config"]["use_swa"] = False + cfg = build_minimax_dspark_draft_config(config) + assert cfg.dflash_config["use_swa"] is False + assert "swa_window_size" not in cfg.dflash_config + + +def test_load_pretrained_config_routes_by_architecture(tmp_path: pathlib.Path) -> None: + (tmp_path / "config.json").write_text(json.dumps(_CONFIG)) + cfg = load_pretrained_config(str(tmp_path), trust_remote_code=False) + assert cfg.architectures == [MINIMAX_DSPARK_DRAFT_ARCH] + assert cfg.hidden_size == 6144 and cfg.num_hidden_layers == 5 + + +def test_backbone_is_registered_under_the_checkpoint_architecture() -> None: + assert get_registered_model_class(MINIMAX_DSPARK_DRAFT_ARCH) is MiniMaxDSparkBackbone + assert issubclass(MiniMaxDSparkForCausalLM, GQADSparkForCausalLM) + + +@pytest.mark.parametrize( + "name, expected", + [ + ( + "language_model.model.dspark.layers.3.decoder_layer.self_attn.q_proj.weight", + "layers.3.self_attn.q_proj.weight", + ), + ( + "language_model.model.dspark.layers.0.decoder_layer.mlp.down_proj.weight_scale_inv", + "layers.0.mlp.down_proj.weight_scale_inv", + ), + ("language_model.model.dspark.fc.weight", "fc.weight"), + ("language_model.model.dspark.hidden_norm.weight", "hidden_norm.weight"), + ("language_model.model.dspark.final_norm.weight", "norm.weight"), + ( + "language_model.model.dspark.markov_head.markov_w1.weight", + "markov_head.markov_w1.weight", + ), + ( + "language_model.model.dspark.confidence_head.proj.bias", + "confidence_head.proj.bias", + ), + ("lm_head.weight", "lm_head.weight"), + ], +) +def test_weight_name_remap(name: str, expected: str) -> None: + assert remap_minimax_dspark_weight_name(name) == expected + + +def test_mxfp8_dequantization_matches_per_block_reference() -> None: + torch.manual_seed(0) + out_features, in_features = 4, 96 # three 32-wide K blocks + # Build finite E4M3 codes directly: dividing bf16-range samples by block scales as small + # as 2**-13 would overflow E4M3 (max 448) into NaN and make the comparison meaningless. + quantized = (torch.randn(out_features, in_features) * 3).to(torch.float8_e4m3fn) + assert torch.isfinite(quantized.to(torch.float32)).all() + # UE8M0 exponents around 127 (2**0). + exponents = torch.randint(114, 131, (out_features, in_features // 32), dtype=torch.uint8) + scale = torch.exp2(exponents.to(torch.float32) - 127).repeat_interleave(32, dim=1) + expected = (quantized.to(torch.float32) * scale).to(torch.bfloat16) + + got = dequantize_mxfp8_weight(quantized, exponents, torch.bfloat16) + assert got.dtype == torch.bfloat16 + torch.testing.assert_close(got, expected, rtol=0, atol=0) + + +def test_prepare_weights_dequantizes_and_folds_gemma_norms() -> None: + prefix = "language_model.model.dspark." + quantized = torch.ones(2, 64, dtype=torch.float8_e4m3fn) + scale = torch.full((2, 2), 128, dtype=torch.uint8) # 2**1 per block + norm = torch.full((8,), 0.25, dtype=torch.bfloat16) + weights = { + prefix + "layers.0.decoder_layer.self_attn.k_proj.weight": quantized, + prefix + "layers.0.decoder_layer.self_attn.k_proj.weight_scale_inv": scale, + prefix + "layers.0.decoder_layer.self_attn.k_norm.weight": norm.clone(), + prefix + "layers.0.decoder_layer.input_layernorm.weight": norm.clone(), + prefix + "hidden_norm.weight": norm.clone(), + prefix + "final_norm.weight": norm.clone(), + prefix + "fc.weight": torch.zeros(8, 16, dtype=torch.bfloat16), + prefix + "markov_head.markov_w1.weight": norm.clone(), + } + out = prepare_minimax_dspark_weights(weights, torch.bfloat16) + + assert set(out) == { + "layers.0.self_attn.k_proj.weight", + "layers.0.self_attn.k_norm.weight", + "layers.0.input_layernorm.weight", + "hidden_norm.weight", + "norm.weight", + "fc.weight", + "markov_head.markov_w1.weight", + } + assert out["layers.0.self_attn.k_proj.weight"].dtype == torch.bfloat16 + torch.testing.assert_close( + out["layers.0.self_attn.k_proj.weight"], + torch.full((2, 64), 2.0, dtype=torch.bfloat16), + ) + for key in ( + "layers.0.self_attn.k_norm.weight", + "layers.0.input_layernorm.weight", + "hidden_norm.weight", + "norm.weight", + ): + torch.testing.assert_close(out[key], torch.full((8,), 1.25, dtype=torch.bfloat16)) + # Head tensors are not norms and are passed through untouched. + torch.testing.assert_close(out["markov_head.markov_w1.weight"], norm) From 7d172e66c9c6ac21307c69bab4c2e358dc73230d Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:50:37 -0700 Subject: [PATCH 2/5] [None][fix] Run windowed DSpark draft blocks non-causally on TRTLLM DSpark drafters are trained with a non-causal block and a two-sided window, but DFlash resolved every windowed draft layer as causal. GQA DSpark now resolves non-causal (W-1, W-1); the TRTLLM backend serves it as the causal windowed decode plus a log-sum-exp merge of the later block keys. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_dflash.py | 202 +++++++-- tensorrt_llm/_torch/models/modeling_dspark.py | 12 + .../hw_agnostic/test_dflash2_semantics.py | 47 ++- .../test_dspark_noncausal_window.py | 388 ++++++++++++++++++ 4 files changed, 594 insertions(+), 55 deletions(-) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index 14af68286a0c..a73c307077e3 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 @@ -63,6 +64,130 @@ 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: + # trtllm-gen reports the log-sum-exp in base 2. + lse = decoded[1] * math.log(2.0) + 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( @@ -1305,29 +1430,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 = "non-causal context" + 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, @@ -1697,29 +1833,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( diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 619519d1f73c..fa6d0060720d 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -2772,6 +2772,18 @@ 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 block attention is non-causal on every layer. + + The DFlash default reads a windowed sliding layer as causal. DSpark + drafters are trained non-causal (the heads refuse ``causal=true``, see + :class:`_DSparkHeadMixin`), and a windowed layer bounds both sides by + the window, as the HF flash path they are trained under does (see + ``dspark_layer_window_size``). + """ + _, (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/tests/unittest/_torch/speculative/hw_agnostic/test_dflash2_semantics.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash2_semantics.py index e2cfaa264682..2d7d1ee6f515 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash2_semantics.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash2_semantics.py @@ -627,34 +627,46 @@ def test_is_causal_is_ignored_for_a_plain_dflash_drafter(): # --------------------------------------------------------------------------- -# TRTLLM block-decode backend: no non-causal sliding window +# TRTLLM block-decode backend: a non-causal sliding window is lowered to the +# causal windowed decode plus the future-block merge # --------------------------------------------------------------------------- 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() + assert ( + _mask_wrapper( + _all_sliding_config(2048), backend=backend + ).validate_block_attention_windows() + == [] + ) -def test_trtllm_backend_rejects_the_released_drafter_window(): +def test_trtllm_backend_lowers_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.""" + finite window_left, so the released non-causal 2048 window runs as the + causal windowed decode plus the future-block merge.""" 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() + assert wrapper.validate_block_attention_windows() == ["causal decode + future-block merge"] * 2 -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_backend_keeps_causal_and_unwindowed_layers_on_their_kernels(): + # A causal sliding window is supported by the decode kernel directly. + assert ( + _mask_wrapper( + _all_sliding_config(4, is_causal=True), backend="TRTLLM" + ).validate_block_attention_windows() + == ["causal decode"] * 2 + ) + # No window at all runs the non-causal context kernel. + assert ( + _mask_wrapper( + SimpleNamespace(num_hidden_layers=2, layer_types=["full_attention"] * 2), + backend="TRTLLM", + ).validate_block_attention_windows() + == ["non-causal context"] * 2 + ) def test_trtllm_window_check_reads_the_use_swa_override(): @@ -663,8 +675,7 @@ def test_trtllm_window_check_reads_the_use_swa_override(): 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() + assert wrapper.validate_block_attention_windows() == ["causal decode + future-block merge"] * 2 # --------------------------------------------------------------------------- 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..668266bca8b2 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py @@ -0,0 +1,388 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Non-causal sliding-window block attention for DSpark drafters. + +A DSpark block query ``j`` at position ``C + j`` (``C`` context tokens) attends +every key within ``W - 1`` positions on both sides: context keys from +``C + j - W + 1`` and every block key. The TRTLLM backend serves that mask as +the causal windowed decode plus the block keys after each query, merged by +log-sum-exp. These tests check the mask resolution (CPU), the merge against a +dense oracle (CPU), and the full trtllm-gen path against a dense fp32 oracle +over a paged cache (SM100/SM103 only). +""" + +import math +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.models.modeling_dflash import ( + DFlashForCausalLM, + dflash_merge_future_block_attention, + dflash_trtllm_windowed_block_attention, +) +from tensorrt_llm._torch.models.modeling_dspark import GQADSparkForCausalLM + +# --------------------------------------------------------------------------- +# Dense oracle +# --------------------------------------------------------------------------- + + +def _oracle(query, keys, values, ctx_lens, window_left, window_right, causal): + """Dense fp32 block attention with the trained mask. + + query: [B, blk, Hq, D]; keys/values: per request [C_b + blk, Hkv, D]. + Returns ([B * blk, Hq, D] output, [B * blk, Hq] natural-log lse). + """ + B, blk, num_heads, head_dim = query.shape + outs, lses = [], [] + for b in range(B): + k = keys[b].float() + v = values[b].float() + num_kv_heads = k.shape[1] + group = num_heads // num_kv_heads + q = query[b].float() # [blk, Hq, D] + k_rep = k.repeat_interleave(group, dim=1) # [L, Hq, D] + v_rep = v.repeat_interleave(group, dim=1) + logits = torch.einsum("jhd,lhd->hjl", q, k_rep) * head_dim**-0.5 + pos_q = ctx_lens[b] + torch.arange(blk, device=q.device).view(-1, 1) + pos_k = torch.arange(k.shape[0], device=q.device).view(1, -1) + allowed = torch.ones_like(pos_q - pos_k, dtype=torch.bool) + if window_left >= 0: + allowed &= pos_k >= pos_q - window_left + if causal: + allowed &= pos_k <= pos_q + elif window_right >= 0: + allowed &= pos_k <= pos_q + window_right + logits = logits.masked_fill(~allowed, float("-inf")) + lse = torch.logsumexp(logits, dim=-1) # [Hq, blk] + out = torch.einsum("hjl,lhd->jhd", torch.softmax(logits, dim=-1), v_rep) + outs.append(out) + lses.append(lse.transpose(0, 1)) + return torch.cat(outs), torch.cat(lses) + + +# --------------------------------------------------------------------------- +# Merge (CPU): causal oracle + merge == non-causal oracle +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("block_size", [1, 2, 4, 7]) +@pytest.mark.parametrize("window", [3, 8, 64]) +@pytest.mark.parametrize("heads", [(4, 1), (8, 2)]) +def test_merge_matches_the_non_causal_oracle(block_size, window, heads): + torch.manual_seed(0) + num_heads, num_kv_heads = heads + head_dim = 16 + ctx_lens = [0, 1, window - block_size, window - 1, window, window + 5, 3 * window + 2] + ctx_lens = [c for c in ctx_lens if c >= 0] + B = len(ctx_lens) + query = torch.randn(B, block_size, num_heads, head_dim) + keys = [torch.randn(c + block_size, num_kv_heads, head_dim) for c in ctx_lens] + values = [torch.randn(c + block_size, num_kv_heads, head_dim) for c in ctx_lens] + block_key = torch.stack([k[-block_size:] for k in keys]) + block_value = torch.stack([v[-block_size:] for v in values]) + w = window - 1 + + out_causal, lse_causal = _oracle(query, keys, values, ctx_lens, w, -1, causal=True) + expected, _ = _oracle(query, keys, values, ctx_lens, w, w, causal=False) + merged = dflash_merge_future_block_attention( + out_causal, lse_causal, query, block_key, block_value, head_dim**-0.5, w + ) + torch.testing.assert_close(merged, expected, atol=1e-5, rtol=1e-5) + if block_size > 1: + # The causal output alone is a different attention: the merge matters. + assert not torch.allclose(out_causal, expected, atol=1e-3) + + +def test_merge_bounds_the_future_keys_by_the_right_window(): + """A right window smaller than the block hides block keys more than + window_right ahead of the query.""" + torch.manual_seed(1) + block_size, w = 7, 2 + query = torch.randn(1, block_size, 4, 16) + keys = [torch.randn(10 + block_size, 1, 16)] + values = [torch.randn(10 + block_size, 1, 16)] + out_causal, lse_causal = _oracle(query, keys, values, [10], w, -1, causal=True) + expected, _ = _oracle(query, keys, values, [10], w, w, causal=False) + merged = dflash_merge_future_block_attention( + out_causal, + lse_causal, + query, + keys[0][-block_size:].unsqueeze(0), + values[0][-block_size:].unsqueeze(0), + 16**-0.5, + w, + ) + torch.testing.assert_close(merged, expected, atol=1e-5, rtol=1e-5) + + +# --------------------------------------------------------------------------- +# Mask resolution (CPU) +# --------------------------------------------------------------------------- + + +def _wrapper(cls, config, layer_windows=(), backend="TRTLLM"): + wrapper = cls.__new__(cls) + torch.nn.Module.__init__(wrapper) + wrapper.config = config + wrapper._is_dflash2 = False + wrapper._uses_lilicorr = False + wrapper._sliding_layers_causal = False + wrapper.dflash_attention_backend = backend + wrapper._layer_windows = list(layer_windows) + return wrapper + + +def _minimax_like_config(num_layers=5, window=4096): + return SimpleNamespace( + num_hidden_layers=num_layers, + layer_types=["sliding_attention"] * num_layers, + sliding_window=window, + ) + + +def test_dspark_windowed_layers_are_non_causal_and_two_sided(): + """The MiniMax-M3.1 drafter: five sliding layers, window 4096, plus the + use_swa window the config flattening adds.""" + config = _minimax_like_config() + wrapper = _wrapper(GQADSparkForCausalLM, config, layer_windows=[(4095, 4095)] * 5) + for layer_idx in range(5): + assert wrapper._get_attention_mask_args(layer_idx) == (False, (4095, 4095)) + assert wrapper._resolve_block_attention(layer_idx) == (False, (4095, 4095)) + assert wrapper.validate_block_attention_windows() == ["causal decode + future-block merge"] * 5 + + +def test_dspark_unwindowed_layers_stay_non_causal(): + config = SimpleNamespace(num_hidden_layers=2, layer_types=["full_attention"] * 2) + wrapper = _wrapper(GQADSparkForCausalLM, config) + assert wrapper._resolve_block_attention(0) == (False, (-1, -1)) + assert wrapper.validate_block_attention_windows() == ["non-causal context"] * 2 + + +def test_plain_dflash_keeps_reading_a_windowed_layer_as_causal(): + """No-regression: the DSpark override does not reach DFlash drafters.""" + wrapper = _wrapper(DFlashForCausalLM, _minimax_like_config(num_layers=2, window=512)) + assert wrapper._get_attention_mask_args(0) == (True, (511, 0)) + + +# --------------------------------------------------------------------------- +# Full trtllm-gen path (SM100/SM103): paged cache, decode + merge vs oracle +# --------------------------------------------------------------------------- + + +def _trtllm_gen_ops_or_skip(): + 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() + + +class _PagedCase: + """Requests with their own context lengths in one shuffled page pool.""" + + def __init__(self, ops, ctx_lens, block_size, num_heads, num_kv_heads, head_dim, page_size): + dev = torch.device("cuda") + self.ops = ops + self.ctx_lens = ctx_lens + self.block_size = block_size + B = len(ctx_lens) + seq_lens = [c + block_size for c in ctx_lens] + pages_per_req = [math.ceil(s / page_size) for s in seq_lens] + self.max_pages = max(pages_per_req) + 1 + num_pages = sum(pages_per_req) + 3 + perm = torch.randperm(num_pages).tolist() + self.block_tables = torch.zeros(B, self.max_pages, dtype=torch.int32, device=dev) + cursor = 0 + for b, n in enumerate(pages_per_req): + self.block_tables[b, :n] = torch.tensor(perm[cursor : cursor + n], dtype=torch.int32) + cursor += n + self.page_size = page_size + self.cache = torch.zeros( + num_pages, 2, num_kv_heads, page_size, head_dim, dtype=torch.bfloat16, device=dev + ) + self.keys = [ + torch.randn(s, num_kv_heads, head_dim, device=dev).to(torch.bfloat16) for s in seq_lens + ] + self.values = [ + torch.randn(s, num_kv_heads, head_dim, device=dev).to(torch.bfloat16) for s in seq_lens + ] + # Context K/V: written straight into the pages. + for b, c in enumerate(ctx_lens): + for pos in range(0, c, page_size): + n = min(page_size, c - pos) + page = self.block_tables[b, pos // page_size] + self.cache[page, 0, :, :n] = self.keys[b][pos : pos + n].transpose(0, 1) + self.cache[page, 1, :, :n] = self.values[b][pos : pos + n].transpose(0, 1) + self.query = torch.randn(B, block_size, num_heads, head_dim, device=dev).to(torch.bfloat16) + self.block_key = torch.stack([k[-block_size:] for k in self.keys]) + self.block_value = torch.stack([v[-block_size:] for v in self.values]) + self.seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=dev) + self.workspace = torch.empty( + ops.get_workspace_size( + dtype=torch.bfloat16, + num_tokens=B * block_size, + num_gen_tokens=B * block_size, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_size=head_dim, + max_num_requests=B, + rotary_embedding_dim=0, + fp8_context_fmha=False, + ), + dtype=torch.uint8, + device=dev, + ) + sm_count = torch.cuda.get_device_properties(dev).multi_processor_count + self.counters = torch.zeros( + ops.get_multi_ctas_kv_counter_size(num_heads, B, sm_count), + dtype=torch.uint8, + device=dev, + ) + + def append_block(self): + """The block K/V go in through the forward's append op.""" + B, blk = len(self.ctx_lens), self.block_size + dev = self.query.device + last = ((self.seq_lens - 1) % self.page_size) + 1 + self.ops.append_paged_kv_cache( + append_key=self.block_key.reshape(B * blk, *self.block_key.shape[2:]).contiguous(), + append_value=self.block_value.reshape( + B * blk, *self.block_value.shape[2:] + ).contiguous(), + batch_indices=torch.arange(B, dtype=torch.int32, device=dev).repeat_interleave(blk), + positions=( + torch.tensor(self.ctx_lens, dtype=torch.int32, device=dev).view(-1, 1) + + torch.arange(blk, dtype=torch.int32, device=dev) + ).reshape(-1), + paged_kv_cache=self.cache, + kv_indices=self.block_tables.flatten(), + kv_indptr=torch.arange( + 0, (B + 1) * self.max_pages, self.max_pages, dtype=torch.int32, device=dev + ), + kv_last_page_len=last, + kv_layout="HND", + ) + + def run(self, causal, window_size, out=None): + B, blk, num_heads, head_dim = self.query.shape + out = ( + torch.empty(B * blk, num_heads, head_dim, dtype=torch.bfloat16, device="cuda") + if out is None + else out + ) + dflash_trtllm_windowed_block_attention( + self.ops, + self.query, + self.block_key, + self.block_value, + self.cache, + self.block_tables, + self.seq_lens, + self.max_pages * self.page_size, + causal, + window_size, + self.workspace, + self.counters, + out, + ) + return out + + +def _ctx_lens_around(window, block_size, page_size): + w = window + return sorted( + { + 1, + max(1, w - block_size - 1), + max(1, w - block_size), + w - 1, + w, + w + 1, + w + page_size - 1, + w + page_size, + 3 * w + 17, + } + ) + + +@pytest.mark.parametrize("page_size", [32, 64]) +@pytest.mark.parametrize("block_size", [1, 2, 4, 7]) +@pytest.mark.parametrize("heads", [(16, 1), (16, 2)]) +@pytest.mark.parametrize("window", [64, 4096]) +def test_trtllm_windowed_non_causal_matches_the_oracle(page_size, block_size, heads, window): + """Mixed context lengths in one batch, straddling the window and page + boundaries, through the forward's append + decode + merge.""" + ops = _trtllm_gen_ops_or_skip() + torch.manual_seed(block_size * 1000 + window + page_size) + num_heads, num_kv_heads = heads + ctx_lens = _ctx_lens_around(window, block_size, page_size) + case = _PagedCase(ops, ctx_lens, block_size, num_heads, num_kv_heads, 128, page_size) + case.append_block() + w = window - 1 + got = case.run(False, (w, w)).float() + expected, _ = _oracle(case.query, case.keys, case.values, ctx_lens, w, w, causal=False) + torch.testing.assert_close(got, expected, atol=2e-2, rtol=2e-2) + + # The stock causal path is the same kernel without the merge; with more + # than one block token it is a different attention. + causal_expected, _ = _oracle(case.query, case.keys, case.values, ctx_lens, w, -1, causal=True) + torch.testing.assert_close( + case.run(True, (w, 0)).float(), causal_expected, atol=2e-2, rtol=2e-2 + ) + + +def test_trtllm_windowed_non_causal_long_context(): + """Contexts far past the window: only the window may be read.""" + ops = _trtllm_gen_ops_or_skip() + torch.manual_seed(7) + ctx_lens = [40_000, 131_000, 4_095, 70] + case = _PagedCase(ops, ctx_lens, 7, 16, 1, 128, 64) + case.append_block() + expected, _ = _oracle(case.query, case.keys, case.values, ctx_lens, 4095, 4095, causal=False) + torch.testing.assert_close( + case.run(False, (4095, 4095)).float(), expected, atol=2e-2, rtol=2e-2 + ) + + +def test_trtllm_windowed_non_causal_replays_under_cuda_graph(): + """Captured once, replayed with new queries, block K/V and lengths.""" + ops = _trtllm_gen_ops_or_skip() + torch.manual_seed(11) + ctx_lens = [100, 4_090, 5_000] + case = _PagedCase(ops, ctx_lens, 7, 16, 1, 128, 64) + case.append_block() + out = torch.empty(len(ctx_lens) * 7, 16, 128, dtype=torch.bfloat16, device="cuda") + case.run(False, (4095, 4095), out) # warm up outside capture + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + case.run(False, (4095, 4095), out) + + # New data in the same buffers: shorter contexts, fresh queries and block. + new_ctx = [90, 4_089, 4_999] + case.ctx_lens[:] = new_ctx + case.seq_lens.copy_(torch.tensor([c + 7 for c in new_ctx], dtype=torch.int32)) + case.query.copy_(torch.randn_like(case.query, dtype=torch.float32).to(torch.bfloat16)) + for b, c in enumerate(new_ctx): + case.keys[b] = case.keys[b][: c + 7].clone() + case.values[b] = case.values[b][: c + 7].clone() + case.keys[b][-7:] = torch.randn_like(case.keys[b][-7:], dtype=torch.float32).to( + torch.bfloat16 + ) + case.values[b][-7:] = torch.randn_like(case.values[b][-7:], dtype=torch.float32).to( + torch.bfloat16 + ) + case.block_key.copy_(torch.stack([k[-7:] for k in case.keys])) + case.block_value.copy_(torch.stack([v[-7:] for v in case.values])) + case.append_block() + graph.replay() + expected, _ = _oracle(case.query, case.keys, case.values, new_ctx, 4095, 4095, causal=False) + torch.testing.assert_close(out.float(), expected, atol=2e-2, rtol=2e-2) From 11d6f38c493dbb7a1d60b12006e93d6249c32fc6 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Mon, 26 Oct 2026 19:08:00 -0700 Subject: [PATCH 3/5] [None][test] Run the DSpark non-causal window test on B200 The trtllm-gen cases skip on the CPU and H100 lists that collect hw_agnostic. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- tests/integration/test_lists/test-db/l0_b200.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 94cea2e5c098..8b1db5cd6f01 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -141,6 +141,8 @@ l0_b200: - unittest/_torch/modules/test_engram_history_uploads.py - unittest/_torch/modules/test_engram_replay.py - unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py + # trtllm-gen non-causal windowed DSpark block attention; skips off SM100/SM103. + - unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py - unittest/disaggregated/test_csa2_kv_transfer.py # ------------- Kimi K3 (KimiLinear) unit tests --------------- - unittest/_torch/modeling/test_kimi_k3_checkpoint.py TIMEOUT (15) From 99428dc3959bc4fde32f54d47b43e85dea8c6771 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Tue, 27 Oct 2026 10:57:31 -0700 Subject: [PATCH 4/5] [None][docs] List the MiniMax-M3.1-Flash DSpark drafter architecture DSparkMiniMaxDraftModel is the public architecture of the drafter in the MiniMax-M3.1-Flash checkpoint, so it joins the telemetry allowlist and the supported-models table. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- docs/source/models/supported-models.md | 3 ++- tensorrt_llm/usage/architecture_allowlist.py | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 654295a8bd14..be80d6999f50 100644 --- a/docs/source/models/supported-models.md +++ b/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: /dspark}`. # Encoder-Decoder Feature Support Matrix (PyTorch Backend) diff --git a/tensorrt_llm/usage/architecture_allowlist.py b/tensorrt_llm/usage/architecture_allowlist.py index eb71ba1beb1e..d36dfad92ebb 100644 --- a/tensorrt_llm/usage/architecture_allowlist.py +++ b/tensorrt_llm/usage/architecture_allowlist.py @@ -30,6 +30,7 @@ "Cohere2ForCausalLM", "Cosmos3EdgeForConditionalGeneration", "Cosmos3ForConditionalGeneration", + "DSparkMiniMaxDraftModel", "DeciLMForCausalLM", "DeepSeekV2ForCausalLM", "DeepseekV32ForCausalLM", From 2bd688e69cdddf18025ad05704ecdcb2ed77a733 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Tue, 27 Oct 2026 17:52:55 -0700 Subject: [PATCH 5/5] [TRTLLM-17389][fix] Honor window metadata, derive the context-kernel check from lowerings GQA DSpark keeps the DFlash resolution when the config declares a causal window; the TRTLLM runtime check reads the returned lowerings, so windowed layers skip the context-kernel head-size limit. MXFP8 scales are read as UE8M0 bytes, a quantized draft config is refused, the drafter tests carry cpu_only, and the 131k-prompt limitation is in the supported-models footnote. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- docs/source/models/supported-models.md | 2 +- tensorrt_llm/_torch/models/modeling_dflash.py | 11 +++- tensorrt_llm/_torch/models/modeling_dspark.py | 51 +++++++++++++++-- tensorrt_llm/_torch/speculative/dflash.py | 11 ++-- .../_torch/speculative/dflash_attention.py | 5 ++ .../hw_agnostic/test_dflash2_semantics.py | 2 + .../test_dspark_noncausal_window.py | 57 +++++++++++++++++++ .../hw_agnostic/test_minimax_dspark_draft.py | 27 +++++++++ 8 files changed, 152 insertions(+), 14 deletions(-) diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index be80d6999f50..c613ca70c75e 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -101,7 +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: /dspark}`. +[^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: /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) diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index ce061eb41fba..5b325d73ebc2 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -23,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, @@ -33,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 @@ -181,8 +187,7 @@ def dflash_trtllm_windowed_block_attention( multi_ctas_kv_counter_buffer=counters, ) if merge_future: - # trtllm-gen reports the log-sum-exp in base 2. - lse = decoded[1] * math.log(2.0) + 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 @@ -1515,7 +1520,7 @@ def validate_block_attention_windows(self) -> list[str]: elif window_size[0] >= 0: lowering = "causal decode + future-block merge" else: - lowering = "non-causal context" + lowering = DFLASH_CONTEXT_LOWERING logger.info_once( f"DFlash draft layer {layer_idx}: causal={causal}, " f"window={window_size}, TRTLLM lowering: {lowering}.", diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index a6a055b70928..b43054310cc2 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -122,7 +122,7 @@ class or the edge would become a cycle. _rename_deepseek_v4_attn_subkey, _rename_deepseek_v4_ffn_subkey, ) -from .modeling_dflash import DFlashForCausalLM, resolve_dspark_head_config +from .modeling_dflash import TRTLLM_GEN_LSE_LN2, DFlashForCausalLM, resolve_dspark_head_config from .modeling_speculative import ( DSparkConfidenceHead, SpecDecOneEngineForCausalLM, @@ -2466,8 +2466,6 @@ def __init__(self, model_config): "confidence_proj.bias": ("confidence_head.proj.bias", "confidence_proj.bias"), } - -_LN2 = math.log(2.0) # flashinfer sizes its own scratch from this; the MLA decode path is the only # consumer here and 256 MiB is what the DFlash TRTLLM buffers already reserve. _MLA_DECODE_WORKSPACE_BYTES = 256 * 1024 * 1024 @@ -2796,10 +2794,26 @@ def _get_attention_mask_args(self, layer_idx): :class:`_DSparkHeadMixin`), and a windowed layer bounds both sides by the window, as the HF flash path they are trained under does (see ``dspark_layer_window_size``). + + Explicit checkpoint metadata wins over this default: a config that + declares a causal window (``is_causal: true``, or the speculators + schema's ``sliding_window_non_causal: false``) keeps the DFlash + resolution, the causal windowed decode. """ - _, (window_left, _) = super()._get_attention_mask_args(layer_idx) + base = super()._get_attention_mask_args(layer_idx) + if self._declares_causal_window(): + return base + _, (window_left, _) = base return False, (window_left, window_left) + def _declares_causal_window(self) -> bool: + """Whether the drafter config states that its sliding window is causal.""" + is_causal = getattr(self.config, "is_causal", None) + if is_causal is not None: + return bool(is_causal) + non_causal = getattr(self.config, "sliding_window_non_causal", None) + return non_causal is not None and not bool(non_causal) + 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) @@ -3212,7 +3226,7 @@ def _mla_paged_attention(self, query, blk_kv, layer_cache, fixup, block_size, ma # --- log-sum-exp merge (flash-decoding combine), kernel LSE is base 2 --- # w_u / (w_l + w_u) == sigmoid(lse_u - lse_l): one sigmoid and one lerp, no # explicit max needed. The last query's triangle is empty, so sigmoid is 0. - lse_l = lse_lower.view(batch, block_size, num_heads).float() * _LN2 + lse_l = lse_lower.view(batch, block_size, num_heads).float() * TRTLLM_GEN_LSE_LN2 weight_upper = torch.sigmoid(lse_upper - lse_l).unsqueeze(-1) merged = torch.lerp(out_lower.float(), out_upper.float(), weight_upper) return merged.to(query.dtype) @@ -3641,6 +3655,15 @@ def dequantize_mxfp8_weight( """``weight`` (e4m3, ``[out, in]``) times ``2 ** (scale_inv - 127)`` per ``[1, 32]`` block (``scale_inv``: UE8M0 as uint8, ``[out, in / 32]``), in ``dtype``.""" + if scale_inv.dtype != torch.uint8: + # The exponent code is wanted, not the value it encodes: ``.to`` of a + # float8_e8m0fnu tensor would convert ``2 ** k``, not ``k``. + if scale_inv.element_size() != 1: + raise ValueError( + "MXFP8 weight_scale_inv must hold UE8M0 bytes (uint8 or " + f"float8_e8m0fnu), got {scale_inv.dtype}" + ) + scale_inv = scale_inv.view(torch.uint8) w = weight.to(torch.float32) exponent = scale_inv.to(torch.int32) - 127 scale = torch.exp2(exponent.to(torch.float32)) @@ -3656,6 +3679,9 @@ def prepare_minimax_dspark_weights(weights: Dict, dtype: torch.dtype) -> Dict: tensor and every Gemma norm weight incremented by one. """ renamed = {remap_minimax_dspark_weight_name(k): v for k, v in weights.items()} + # Every tensor is materialized whole on the host, and the dequantization + # goes through float32, before any TP slicing: sized for this five-layer + # drafter, not a pattern for a full backbone. out = {} for name, value in renamed.items(): if name.endswith(_MXFP8_SCALE_SUFFIX): @@ -3809,6 +3835,21 @@ class MiniMaxDSparkForCausalLM(GQADSparkForCausalLM): the DFlash loader sees it. See :func:`prepare_minimax_dspark_weights`. """ + def __init__(self, draft_config, *args: Any, **kwargs: Any) -> None: + # The backbone must be built unquantized: the loader hands it bf16 + # tensors. ``quantization_config`` is dropped from config.json by + # ``build_minimax_dspark_draft_config``, but a ``hf_quant_config.json`` + # next to the checkpoint would still quantize the Linears. + quant_algo = getattr(draft_config.quant_config, "quant_algo", None) + if quant_algo is not None: + raise ValueError( + "The MiniMax DSpark drafter dequantizes its MXFP8 weights at " + f"load and must be built unquantized; the draft config carries " + f"quant_algo={quant_algo}. Remove hf_quant_config.json from the " + "drafter directory." + ) + super().__init__(draft_config, *args, **kwargs) + def load_weights(self, weights: Dict, weight_mapper: Any = None, **kwargs: Any) -> Any: weights = prepare_minimax_dspark_weights(weights, self.config.torch_dtype) return super().load_weights(weights, weight_mapper=weight_mapper, **kwargs) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 6b7e6b517876..fb85eafd3cd8 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -36,6 +36,7 @@ from ..pyexecutor.resource_manager import BaseResourceManager from .accept_stats import maybe_create_recorder from .dflash_attention import ( + DFLASH_CONTEXT_LOWERING, get_dflash_paged_append, validate_dflash_fa4_runtime, validate_dflash_trtllm_gen_runtime, @@ -1337,11 +1338,11 @@ 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() - has_context_attention = any( - not draft_model._get_attention_mask_args(layer_idx)[0] for layer_idx in range(L) - ) + # like them. The context kernel's head-size limits apply only + # to layers lowered to it; a non-causal windowed layer runs + # the block decode plus the future-block merge. + lowerings = draft_model.validate_block_attention_windows() + has_context_attention = DFLASH_CONTEXT_LOWERING in lowerings validate_dflash_trtllm_gen_runtime( dtype=dtype, num_heads=nh, diff --git a/tensorrt_llm/_torch/speculative/dflash_attention.py b/tensorrt_llm/_torch/speculative/dflash_attention.py index 5759211742e1..fd37b5aa5c26 100644 --- a/tensorrt_llm/_torch/speculative/dflash_attention.py +++ b/tensorrt_llm/_torch/speculative/dflash_attention.py @@ -25,6 +25,11 @@ from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE from tensorrt_llm._utils import get_sm_version, is_sm_100f +# TRTLLM lowering ``validate_block_attention_windows`` reports for a draft +# layer that runs the context kernel; every other lowering is the block decode, +# so only this one is subject to the context kernel's head-size limits. +DFLASH_CONTEXT_LOWERING = "non-causal context" + @dataclass(frozen=True) class DFlashTrtllmGenOps: 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 2d7d1ee6f515..88915f2b5a71 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash2_semantics.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash2_semantics.py @@ -28,6 +28,8 @@ from tensorrt_llm._torch.speculative.dflash import DFlashWorker from tensorrt_llm.llmapi.llm_args import is_dflash2_checkpoint +pytestmark = pytest.mark.cpu_only + # --------------------------------------------------------------------------- # Two-tap dynamic depthwise convolution # --------------------------------------------------------------------------- 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 index 668266bca8b2..94b5c161a84c 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py @@ -13,6 +13,7 @@ import math from types import SimpleNamespace +from unittest.mock import patch import pytest import torch @@ -23,6 +24,10 @@ dflash_trtllm_windowed_block_attention, ) from tensorrt_llm._torch.models.modeling_dspark import GQADSparkForCausalLM +from tensorrt_llm._torch.speculative.dflash_attention import ( + DFLASH_CONTEXT_LOWERING, + validate_dflash_trtllm_gen_runtime, +) # --------------------------------------------------------------------------- # Dense oracle @@ -68,6 +73,7 @@ def _oracle(query, keys, values, ctx_lens, window_left, window_right, causal): # --------------------------------------------------------------------------- +@pytest.mark.cpu_only @pytest.mark.parametrize("block_size", [1, 2, 4, 7]) @pytest.mark.parametrize("window", [3, 8, 64]) @pytest.mark.parametrize("heads", [(4, 1), (8, 2)]) @@ -96,6 +102,7 @@ def test_merge_matches_the_non_causal_oracle(block_size, window, heads): assert not torch.allclose(out_causal, expected, atol=1e-3) +@pytest.mark.cpu_only def test_merge_bounds_the_future_keys_by_the_right_window(): """A right window smaller than the block hides block keys more than window_right ahead of the query.""" @@ -143,6 +150,7 @@ def _minimax_like_config(num_layers=5, window=4096): ) +@pytest.mark.cpu_only def test_dspark_windowed_layers_are_non_causal_and_two_sided(): """The MiniMax-M3.1 drafter: five sliding layers, window 4096, plus the use_swa window the config flattening adds.""" @@ -154,6 +162,7 @@ def test_dspark_windowed_layers_are_non_causal_and_two_sided(): assert wrapper.validate_block_attention_windows() == ["causal decode + future-block merge"] * 5 +@pytest.mark.cpu_only def test_dspark_unwindowed_layers_stay_non_causal(): config = SimpleNamespace(num_hidden_layers=2, layer_types=["full_attention"] * 2) wrapper = _wrapper(GQADSparkForCausalLM, config) @@ -161,12 +170,60 @@ def test_dspark_unwindowed_layers_stay_non_causal(): assert wrapper.validate_block_attention_windows() == ["non-causal context"] * 2 +@pytest.mark.cpu_only def test_plain_dflash_keeps_reading_a_windowed_layer_as_causal(): """No-regression: the DSpark override does not reach DFlash drafters.""" wrapper = _wrapper(DFlashForCausalLM, _minimax_like_config(num_layers=2, window=512)) assert wrapper._get_attention_mask_args(0) == (True, (511, 0)) +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "metadata, expected", + [ + ({"is_causal": True}, (True, (511, 0))), + ({"sliding_window_non_causal": False}, (True, (511, 0))), + ({"is_causal": False}, (False, (511, 511))), + ({"sliding_window_non_causal": True}, (False, (511, 511))), + ({"is_causal": False, "sliding_window_non_causal": False}, (False, (511, 511))), + ], +) +def test_dspark_honors_explicit_window_causality_metadata(metadata, expected): + """A config that states its window is causal keeps the DFlash resolution + (causal windowed decode); ``is_causal`` outranks the speculators spelling.""" + config = _minimax_like_config(num_layers=2, window=512) + for key, value in metadata.items(): + setattr(config, key, value) + wrapper = _wrapper(GQADSparkForCausalLM, config) + assert wrapper._get_attention_mask_args(0) == expected + assert wrapper._resolve_block_attention(0) == expected + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("head_dim", [72, 80, 96]) +def test_windowed_layers_are_not_held_to_the_context_kernel_head_sizes(head_dim): + """Regression: a windowed non-causal layer is lowered to the block decode, + so the context kernel's head-size limits must not reject it; an unwindowed + non-causal layer still is.""" + windowed = _wrapper(GQADSparkForCausalLM, _minimax_like_config(num_layers=2, window=512)) + lowerings = windowed.validate_block_attention_windows() + assert DFLASH_CONTEXT_LOWERING not in lowerings + unwindowed = _wrapper( + GQADSparkForCausalLM, + SimpleNamespace(num_hidden_layers=2, layer_types=["full_attention"] * 2), + ) + assert DFLASH_CONTEXT_LOWERING in unwindowed.validate_block_attention_windows() + shape = dict( + dtype=torch.bfloat16, num_heads=16, num_kv_heads=2, head_dim=head_dim, tokens_per_block=32 + ) + with patch("tensorrt_llm._torch.speculative.dflash_attention.get_dflash_trtllm_gen_ops"): + validate_dflash_trtllm_gen_runtime( + **shape, has_context_attention=DFLASH_CONTEXT_LOWERING in lowerings + ) + with pytest.raises(RuntimeError, match="head dimension"): + validate_dflash_trtllm_gen_runtime(**shape, has_context_attention=True) + + # --------------------------------------------------------------------------- # Full trtllm-gen path (SM100/SM103): paged cache, decode + merge vs oracle # --------------------------------------------------------------------------- diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_minimax_dspark_draft.py b/tests/unittest/_torch/speculative/hw_agnostic/test_minimax_dspark_draft.py index 3e562a454368..b631e0c8576d 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_minimax_dspark_draft.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_minimax_dspark_draft.py @@ -17,6 +17,7 @@ import json import pathlib +from types import SimpleNamespace import pytest import torch @@ -35,6 +36,10 @@ build_minimax_dspark_draft_config, load_pretrained_config, ) +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization import QuantAlgo + +pytestmark = pytest.mark.cpu_only # The shipped drafter config.json, reduced to the fields the loaders read. The # backbone geometry sits under text_config; the dspark_* / dflash_config knobs @@ -200,6 +205,28 @@ def test_mxfp8_dequantization_matches_per_block_reference() -> None: torch.testing.assert_close(got, expected, rtol=0, atol=0) +def test_mxfp8_dequantization_reads_the_scale_exponent_bytes() -> None: + """The scale is the UE8M0 code, whichever dtype safetensors hands back.""" + weight = torch.randn(4, 64).to(torch.float8_e4m3fn) + scale = torch.tensor([[127, 128], [126, 130], [127, 127], [120, 135]], dtype=torch.uint8) + expected = dequantize_mxfp8_weight(weight, scale, torch.bfloat16) + e8m0 = getattr(torch, "float8_e8m0fnu", None) + if e8m0 is not None: + torch.testing.assert_close( + dequantize_mxfp8_weight(weight, scale.view(e8m0), torch.bfloat16), expected + ) + with pytest.raises(ValueError, match="UE8M0"): + dequantize_mxfp8_weight(weight, scale.to(torch.float32), torch.bfloat16) + + +def test_minimax_drafter_refuses_a_quantized_draft_config() -> None: + """A hf_quant_config.json next to the checkpoint would quantize the + backbone the loader hands bf16 tensors to; the drafter refuses to build.""" + draft_config = SimpleNamespace(quant_config=QuantConfig(quant_algo=QuantAlgo.MXFP8)) + with pytest.raises(ValueError, match="hf_quant_config.json"): + MiniMaxDSparkForCausalLM(draft_config) + + def test_prepare_weights_dequantizes_and_folds_gemma_norms() -> None: prefix = "language_model.model.dspark." quantized = torch.ones(2, 64, dtype=torch.float8_e4m3fn)