From 0273e8203851eee899f08c64ac62c14f98f05d00 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Wed, 9 Sep 2026 00:46:32 -0700 Subject: [PATCH 01/25] [None][feat] MLA-backboned standalone DSpark drafter (Inferact/Kimi-K3-DSpark) Adds the MLA drafter path alongside the GQA one: it stores a single 576-wide latent per token per layer instead of 16 KV heads x 64 for K and V, which under attention-DP is 5760 vs 20480 bytes per token per rank. Ported from the internal rubin-advance branch, with the position table sized from the runtime ceiling the worker publishes rather than model_config. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/configs/__init__.py | 5 + tensorrt_llm/_torch/configs/k3_dspark.py | 24 + .../dspark_rmsnorm_rope_custom_op.py | 19 +- .../blackwell/dspark_rmsnorm_rope.py | 36 +- tensorrt_llm/_torch/models/__init__.py | 1 + tensorrt_llm/_torch/models/_arch_index.py | 2 + tensorrt_llm/_torch/models/modeling_dflash.py | 135 ++- tensorrt_llm/_torch/models/modeling_dspark.py | 1073 ++++++++++++++++- .../_torch/models/modeling_kimi_linear.py | 60 +- .../_torch/models/modeling_speculative.py | 12 + tensorrt_llm/_torch/models/modeling_utils.py | 16 +- .../_torch/pyexecutor/config_utils.py | 1 + tensorrt_llm/_torch/speculative/dflash.py | 176 ++- .../test_kimi_k3_dspark_semantics.py | 718 +++++++++++ .../test_dspark_cute_dsl_rmsnorm_rope.py | 45 + 15 files changed, 2213 insertions(+), 110 deletions(-) create mode 100644 tensorrt_llm/_torch/configs/k3_dspark.py diff --git a/tensorrt_llm/_torch/configs/__init__.py b/tensorrt_llm/_torch/configs/__init__.py index fee97e5b06e5..dac17fdd45a9 100644 --- a/tensorrt_llm/_torch/configs/__init__.py +++ b/tensorrt_llm/_torch/configs/__init__.py @@ -23,6 +23,7 @@ Gemma4UnifiedTextConfig, Gemma4UnifiedVisionConfig, ) +from tensorrt_llm._torch.configs.k3_dspark import K3DsparkConfig from tensorrt_llm._torch.configs.kimi_k3 import KimiK3Config, KimiK3VisionConfig from tensorrt_llm._torch.configs.kimi_linear import KimiLinearConfig from tensorrt_llm._torch.configs.laguna import LagunaConfig @@ -67,6 +68,9 @@ def _register_custom_configs_with_transformers() -> None: # sub-configs and multimodal is not disabled, and otherwise flattens to # the text config. Registering both here lets AutoConfig / AutoTokenizer # resolve them without trust_remote_code. + # The MLA DSpark drafter checkpoint ships no auto_map, so AutoConfig + # cannot resolve its model_type on its own. + "k3_dspark": K3DsparkConfig, "kimi_k3": KimiK3Config, "kimi_linear": KimiLinearConfig, "laguna": LagunaConfig, @@ -104,6 +108,7 @@ def _register_custom_configs_with_transformers() -> None: "Gemma4UnifiedConfig", "Gemma4UnifiedTextConfig", "Gemma4UnifiedVisionConfig", + "K3DsparkConfig", "KimiK3Config", "KimiK3VisionConfig", "KimiLinearConfig", diff --git a/tensorrt_llm/_torch/configs/k3_dspark.py b/tensorrt_llm/_torch/configs/k3_dspark.py new file mode 100644 index 000000000000..bfba65cd74a9 --- /dev/null +++ b/tensorrt_llm/_torch/configs/k3_dspark.py @@ -0,0 +1,24 @@ +# 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. + +from transformers.configuration_utils import PretrainedConfig + + +# The MLA-backboned DSpark drafter (Inferact/Kimi-K3-DSpark) ships a config.json +# with model_type "k3_dspark", no auto_map and no modeling code, so +# AutoConfig.from_pretrained cannot resolve it. Same workaround as LagunaConfig: +# the fields are plain attributes, and MLADSparkForCausalLM reads them directly. +class K3DsparkConfig(PretrainedConfig): + model_type = "k3_dspark" diff --git a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py index 7f341c4999d8..98ed0d4e3b75 100644 --- a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py @@ -41,6 +41,7 @@ def is_fused_dspark_rmsnorm_rope_supported( freqs: torch.Tensor, num_heads: int, rope_dim: int, + norm_dim: int | None = None, ) -> bool: """Return whether tensors satisfy the production fused-op contract.""" if _get_dspark_arch_str() is None or not all(t.is_cuda for t in (x, weight, freqs)): @@ -51,10 +52,15 @@ def is_fused_dspark_rmsnorm_rope_supported( return False if x.ndim < 2 or x.shape[-1] % 32 != 0: return False - if weight.shape != (x.shape[-1],): - return False if rope_dim < 0 or rope_dim > x.shape[-1] or rope_dim % 2 != 0: return False + # norm_dim defaults to the whole row; the only other supported value is the + # nope prefix, where the weight spans just that prefix. + effective_norm_dim = x.shape[-1] if norm_dim is None else norm_dim + if effective_norm_dim not in (x.shape[-1], x.shape[-1] - rope_dim): + return False + if effective_norm_dim % 32 != 0 or weight.shape != (effective_norm_dim,): + return False if (x.shape[-1] - rope_dim) % 32 != 0 or (rope_dim // 2) % 32 != 0: return False rows = x.numel() // x.shape[-1] @@ -144,6 +150,7 @@ def _compile_fused_dspark_rmsnorm_rope( apply_weight: bool, apply_rmsnorm: bool, inverse_rope: bool, + norm_dim: int, ): rows = cute.sym_int() freq_rows = cute.sym_int() @@ -151,7 +158,7 @@ def _compile_fused_dspark_rmsnorm_rope( cutlass.BFloat16, (rows, hidden_dim), stride_order=(1, 0) ) weight_fake = cute.runtime.make_fake_compact_tensor( - cutlass.BFloat16, (hidden_dim,), stride_order=(0,) + cutlass.BFloat16, (norm_dim,), stride_order=(0,) ) freqs_fake = cute.runtime.make_fake_compact_tensor( cutlass.Float32, @@ -170,6 +177,7 @@ def _compile_fused_dspark_rmsnorm_rope( apply_weight, apply_rmsnorm, inverse_rope, + norm_dim=norm_dim, ) return cute.compile( kernel, @@ -296,9 +304,10 @@ def cute_dsl_dspark_rmsnorm_rope( apply_weight: bool, apply_rmsnorm: bool, inverse_rope: bool, + norm_dim: int | None = None, ) -> torch.Tensor: """Apply fused RMSNorm and adjacent-pair RoPE to contiguous BF16 rows.""" - if not is_fused_dspark_rmsnorm_rope_supported(x, weight, freqs, num_heads, rope_dim): + if not is_fused_dspark_rmsnorm_rope_supported(x, weight, freqs, num_heads, rope_dim, norm_dim): raise ValueError( "cute_dsl_dspark_rmsnorm_rope requires contiguous BF16 tensors on " "an SM100 or SM103 GPU with a valid FP32 frequency view; " @@ -316,6 +325,7 @@ def cute_dsl_dspark_rmsnorm_rope( apply_weight, apply_rmsnorm, inverse_rope, + x.shape[-1] if norm_dim is None else norm_dim, ) compiled(x_flat, weight, freqs, output) return output.view(original_shape) @@ -332,6 +342,7 @@ def _( apply_weight: bool, apply_rmsnorm: bool, inverse_rope: bool, + norm_dim: int | None = None, ) -> torch.Tensor: return torch.empty_like(x) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark_rmsnorm_rope.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark_rmsnorm_rope.py index 71e3171543af..f14c51aa3650 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark_rmsnorm_rope.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark_rmsnorm_rope.py @@ -26,6 +26,7 @@ def __init__( apply_weight: bool, apply_rmsnorm: bool, inverse_rope: bool, + norm_dim: int | None = None, ): if hidden_dim % self.num_threads != 0: raise ValueError( @@ -44,6 +45,22 @@ def __init__( self.apply_weight = apply_weight self.apply_rmsnorm = apply_rmsnorm self.inverse_rope = inverse_rope + # Which prefix the RMS is taken over, and how far the norm scale and the + # weight reach. DSpark normalizes the whole row (the default, bit-identical + # to before this knob existed); DeepSeek-style MLA normalizes only the + # kv_lora_rank latent and leaves k_pe raw, which is nope_dim here. + self.norm_dim = hidden_dim if norm_dim is None else norm_dim + if self.norm_dim not in (hidden_dim, self.nope_dim): + raise ValueError( + f"norm_dim must be hidden_dim ({hidden_dim}) or nope_dim " + f"({self.nope_dim}); got {self.norm_dim}" + ) + self.norm_covers_rope = self.norm_dim == hidden_dim + if self.norm_dim % self.num_threads != 0: + raise ValueError( + f"norm_dim must be divisible by {self.num_threads}; got {self.norm_dim}" + ) + self.norm_elements_per_thread = self.norm_dim // self.num_threads if self.nope_dim % self.num_threads != 0: raise ValueError( f"nope_dim must be divisible by {self.num_threads}; got {self.nope_dim}" @@ -85,12 +102,12 @@ def kernel( inverse_rms = cutlass.Float32(1.0) if cutlass.const_expr(self.apply_rmsnorm): sum_sq = cutlass.Float32(0.0) - for item in cutlass.range_constexpr(self.elements_per_thread): + for item in cutlass.range_constexpr(self.norm_elements_per_thread): dim = tidx + item * self.num_threads value = cutlass.Float32(x[row, dim]) sum_sq += value * value sum_sq = cute.arch.warp_reduction_sum(sum_sq) - inverse_rms = cute.math.rsqrt(sum_sq / self.hidden_dim + self.eps) + inverse_rms = cute.math.rsqrt(sum_sq / self.norm_dim + self.eps) for item in cutlass.range_constexpr(self.nope_elements_per_thread): dim = tidx + item * self.num_threads @@ -105,11 +122,16 @@ def kernel( pair = tidx + item * self.num_threads real_dim = self.nope_dim + pair * 2 imag_dim = real_dim + 1 - real = cutlass.Float32(x[row, real_dim]) * inverse_rms - imag = cutlass.Float32(x[row, imag_dim]) * inverse_rms - if cutlass.const_expr(self.apply_weight): - real *= cutlass.Float32(weight[real_dim]) - imag *= cutlass.Float32(weight[imag_dim]) + real = cutlass.Float32(x[row, real_dim]) + imag = cutlass.Float32(x[row, imag_dim]) + # Outside norm_dim the rope lanes are passed through raw: no RMS + # scale, no weight. weight is only norm_dim long in that case. + if cutlass.const_expr(self.norm_covers_rope): + real *= inverse_rms + imag *= inverse_rms + if cutlass.const_expr(self.apply_weight): + real *= cutlass.Float32(weight[real_dim]) + imag *= cutlass.Float32(weight[imag_dim]) cos = cutlass.Float32(freqs[freq_row, pair, 0]) sin = cutlass.Float32(freqs[freq_row, pair, 1]) if cutlass.const_expr(self.inverse_rope): diff --git a/tensorrt_llm/_torch/models/__init__.py b/tensorrt_llm/_torch/models/__init__.py index e29b68b77a3c..945fdd2fc6ed 100644 --- a/tensorrt_llm/_torch/models/__init__.py +++ b/tensorrt_llm/_torch/models/__init__.py @@ -48,6 +48,7 @@ "HCXVisionForCausalLM", "HunYuanDenseV1ForCausalLM", "HunYuanMoEV1ForCausalLM", + "K3DsparkForCausalLM", "KimiK25ForConditionalGeneration", "KimiK3ForConditionalGeneration", "KimiLinearForCausalLM", diff --git a/tensorrt_llm/_torch/models/_arch_index.py b/tensorrt_llm/_torch/models/_arch_index.py index ca2892fbe509..0a698209c6a9 100644 --- a/tensorrt_llm/_torch/models/_arch_index.py +++ b/tensorrt_llm/_torch/models/_arch_index.py @@ -64,6 +64,7 @@ def is_builtin_zoo_module(module_name: str) -> bool: "HCXVisionModel": "modeling_hyperclovax", "HunYuanDenseV1ForCausalLM": "modeling_hunyuan_dense", "HunYuanMoEV1ForCausalLM": "modeling_hunyuan_moe", + "K3DsparkForCausalLM": "modeling_dspark", "KimiK25ForConditionalGeneration": "modeling_kimi_k25", "KimiK3ForConditionalGeneration": "modeling_kimi_k3_vl", "KimiLinearForCausalLM": "modeling_kimi_linear", @@ -144,6 +145,7 @@ def is_builtin_zoo_module(module_name: str) -> bool: "HCXVisionForCausalLM": "modeling_hyperclovax", "HunYuanDenseV1ForCausalLM": "modeling_hunyuan_dense", "HunYuanMoEV1ForCausalLM": "modeling_hunyuan_moe", + "K3DsparkForCausalLM": "modeling_dspark", "KimiK25ForConditionalGeneration": "modeling_kimi_k25", "KimiK3ForConditionalGeneration": "modeling_kimi_k3_vl", "KimiLinearForCausalLM": "modeling_kimi_linear", diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index 45a31bccce6a..5bd49cc08074 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -26,7 +26,7 @@ get_dflash_trtllm_gen_ops, ) from ..speculative.interface import SpeculativeDecodingMode -from .modeling_utils import get_model_architecture, register_draft_model +from .modeling_utils import FUSED_MODULE_COMPONENTS, get_model_architecture, register_draft_model def dspark_layer_window_size( @@ -323,6 +323,18 @@ class DFlashForCausalLM(nn.Module): Reference: https://arxiv.org/pdf/2602.06036 """ + # Whether ``dflash_attention_backend`` drives this drafter's block decode. + # Subclasses that bring their own attention set this False: neither backend + # can express every drafter shape, the ops behind them are optional + # dependencies, and the worker's per-backend shape checks do not apply. + _uses_worker_attention_backend = True + + # Whether the drafter's context KV lives in the draft KV cache manager's + # paged pool rather than a private arena dense in max_seq_len. Orthogonal + # to the attention backend: paging is about where the KV lives, the backend + # is about which kernel reads it. + _paged_ctx_cache = False + def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): """Build the draft model, resolving its architecture from the draft config (falling back to a model_type-derived name when the checkpoint uses a @@ -374,23 +386,32 @@ def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): "block_size", getattr(pretrained_config, "block_size", None) ) self.dflash_attention_backend = dflash_attention_backend + if self.dflash_attention_backend not in ("VANILLA", "TRTLLM", "FA4"): + raise ValueError( + "DFlash attention backend must be VANILLA, TRTLLM or FA4, got " + f"{self.dflash_attention_backend!r}." + ) # Each backend loads only its own ops; the rest stay None so the # shared paged prologue can read them unconditionally. + self._dflash_flash_attention = None self._dflash_trtllm_gen_ops = None self._dflash_fa4_fwd = None self._dflash_paged_append = None - if self.dflash_attention_backend == "VANILLA": + if not self._uses_worker_attention_backend: + # Still validated above so a typo fails here rather than silently, + # but no op set is loaded: this drafter calls none of them. + logger.info_once( + f"{type(self).__name__} brings its own block decode; " + f"attention_backend={self.dflash_attention_backend!r} is not used.", + key=f"dflash_own_attention_{type(self).__name__}", + ) + elif self.dflash_attention_backend == "VANILLA": self._dflash_flash_attention = get_dflash_flash_attention() elif self.dflash_attention_backend == "TRTLLM": self._dflash_trtllm_gen_ops = get_dflash_trtllm_gen_ops() - elif self.dflash_attention_backend == "FA4": + else: self._dflash_fa4_fwd = get_dflash_fa4_fwd() self._dflash_paged_append = get_dflash_paged_append() - else: - raise ValueError( - "DFlash attention backend must be VANILLA, TRTLLM or FA4, got " - f"{self.dflash_attention_backend!r}." - ) self._dflash_trtllm_gen_workspace = None self._dflash_trtllm_gen_counters = None self.register_buffer("_dflash_batch_indices", None, persistent=False) @@ -469,6 +490,11 @@ def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): self._num_heads = 0 self._head_dim = 0 self._num_kv_heads = 0 + # Cache halves per token per layer: K and V for a GQA drafter, one MLA + # latent (and no V) for an MLA one. The worker sizes its context arena + # and validates the managed pool against this, so a drafter that stores + # a single tensor must override it (see MLADSparkForCausalLM). + self._kv_factor = 2 self._has_qk_norm = False self._use_fused_qk_norm_rope = False # Laguna-specific draft-layer behaviors, disabled by default so generic @@ -727,6 +753,21 @@ def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): else: remapped[key] = value + # Wrapper-owned and built FROM the checkpoint, so they are never in + # draft_model_full's module tree and _assert_backbone_complete cannot + # see them: a module conjured from data cannot be reported missing by + # walking modules. One tuple drives both the extraction below and this + # check, so the two cannot drift. Without fc the drafter has no capture + # projection, has_target_features stays False, _ctx_len never advances + # and it drafts from an empty context forever (the `hasattr` guards on + # that path are degradation, not a supported mode). + wrapper_missing = [k for k in self.WRAPPER_OWNED_WEIGHTS if k not in remapped] + if wrapper_missing: + raise ValueError( + f"{type(self).__name__}: checkpoint is missing {wrapper_missing}, " + "which this wrapper owns and builds from the checkpoint." + ) + # Load DFlash-specific weights directly if "fc.weight" in remapped: self.fc = nn.Linear( @@ -751,9 +792,10 @@ def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): self.hidden_norm.weight.data.copy_(remapped["hidden_norm.weight"]) del remapped["hidden_norm.weight"] - # Load remaining weights into the draft model. - # DFlash checkpoints don't include embed_tokens or lm_head, so allow partial loading - # since those modules won't find matching weights. + # allow_partial_loading is what lets the shared modules below be absent. + # It is also why a truncated checkpoint loads clean, so gate it on the + # subclass's own declaration of what may legitimately be missing. + self._assert_backbone_complete(remapped, weight_mapper) self.draft_model_full.load_weights( weights=remapped, weight_mapper=weight_mapper, allow_partial_loading=True ) @@ -804,6 +846,77 @@ def take(key: str) -> torch.Tensor: self.mlp_convs = mlp_convs self.candidate_selector = selector return {k: v for k, v in weights.items() if k not in consumed} + #: Tensors the wrapper itself owns: not in draft_model_full, built from the + #: checkpoint in load_weights, and required. + WRAPPER_OWNED_WEIGHTS = ("fc.weight", "hidden_norm.weight") + + #: Parameter-name prefixes this drafter takes from the target instead of its + #: own checkpoint. Everything else in ``draft_model_full`` must be provided. + #: GQA DFlash checkpoints ship neither embedding nor head; an MLA drafter + #: ships its own embedding and overrides this. + WEIGHTS_SHARED_WITH_TARGET = ("embed_tokens", "lm_head") + + def _assert_backbone_complete(self, weights: Dict, weight_mapper=None) -> None: + """Fail on a checkpoint missing weights the drafter does not share. + + The truth is the constructed module tree, not a hand-kept list that + rots as the backbone changes. + + Checked at module granularity. For a PLAIN module that is the hole the + flag cannot close: the loader skips one whose subtree filters to + nothing (modeling_utils.py `if module_weights:`) whatever the flag + says. For a FUSED module `allow_partial_loading=False` would catch it + (linear.py asserts all three shards) -- but the flag has to stay True + for the target-shared modules, so the check covers that case here + instead, and requires every component rather than any. + + Missing parameters INSIDE a present component stay tolerated: a + checkpoint with all three weights but only `q_proj.bias` takes the same + per-shard copy and leaves the rest at `torch.empty`. Module-granular + checking cannot see that and does not pretend to. + """ + provided = set(weights) + + # Whichever fusion table the load below will actually use, not a third + # copy: with a mapper modeling_utils dispatches to _load_weights_impl_v2 + # and the mapper's own table applies; without one it falls back to + # _load_weights_impl, whose table is FUSED_MODULE_COMPONENTS. An empty + # `mapping` means init_model_and_config has not run, so the mapper has + # no table to offer yet and the constant is still the right answer. + fusion = dict(getattr(weight_mapper, "mapping", None) or FUSED_MODULE_COMPONENTS) + + def _has(prefix: str) -> bool: + return any(k == prefix or k.startswith(prefix + ".") for k in provided) + + def _supplied(module_name: str) -> bool: + # A fused module is named once here and stored unfused in the + # checkpoint. ALL components must be present, not any: the fused + # load path is happy with a subset under allow_partial_loading + # (linear.py load_weights_fused_qkv_helper) and leaves the absent + # shards at torch.empty -- uninitialised device memory, not zeros. + if _has(module_name): + return True + for fused, parts in fusion.items(): + if fused in module_name: + return all(_has(module_name.replace(fused, p)) for p in parts) + return False + + missing = sorted( + { + name.rsplit(".", 1)[0] + for name, _ in self.draft_model_full.named_parameters() + if not any(part in name for part in self.WEIGHTS_SHARED_WITH_TARGET) + and not _supplied(name.rsplit(".", 1)[0]) + } + ) + if missing: + raise ValueError( + f"{type(self).__name__}: checkpoint provides no weights for " + f"{missing[:8]}{' ...' if len(missing) > 8 else ''}. These are " + f"not in WEIGHTS_SHARED_WITH_TARGET " + f"({', '.join(self.WEIGHTS_SHARED_WITH_TARGET)}), so loading " + "would leave them randomly initialized." + ) def load_weights_from_target_model(self, target_model: torch.nn.Module) -> None: """Share embed_tokens and lm_head from the target model.""" diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 65fbd52b4f1f..0fbdb8553bff 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -79,25 +79,32 @@ class or the edge would become a cycle. import copy import json +import math import os import re from functools import lru_cache -from typing import Dict, List, Optional +from typing import Dict, List, NamedTuple, Optional import torch import torch.nn.functional as F from torch import nn +from transformers import PretrainedConfig from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping from tensorrt_llm.quantization.mode import QuantAlgo from ..._utils import is_sm_100f from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from ..distributed import AllReduceParams from ..model_config import ModelConfig -from ..modules.linear import Linear +from ..modules.decoder_layer import DecoderLayer +from ..modules.embedding import Embedding +from ..modules.gated_mlp import GatedMLP +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 ..speculative.interface import SpeculativeDecodingMode from ..utils import AuxStreamType from .modeling_deepseekv4 import ( @@ -112,11 +119,12 @@ class or the edge would become a cycle. from .modeling_dflash import DFlashForCausalLM, resolve_dspark_head_config from .modeling_speculative import ( DSparkConfidenceHead, + SpecDecOneEngineForCausalLM, build_markov_head, confident_prefix_length, dspark_markov_chain_logits, ) -from .modeling_utils import register_draft_model +from .modeling_utils import DecoderModel, register_auto_model, register_draft_model if IS_CUTLASS_DSL_AVAILABLE: from ..custom_ops.dspark_attention_custom_op import ( @@ -386,11 +394,21 @@ def _rmsnorm_rope_batched( apply_weight: bool = True, apply_rmsnorm: bool = True, inverse_rope: bool = False, + norm_dim: Optional[int] = None, ) -> torch.Tensor: - """Fuse DSpark RMSNorm and last-dimension RoPE when supported.""" + """Fuse DSpark RMSNorm and last-dimension RoPE when supported. + + ``norm_dim`` bounds the RMSNorm (and the weight) to a prefix of the row. + None means the whole row, which is DSpark's own convention. The other + supported value is ``t.shape[-1] - rope_head_dim``: DeepSeek-style MLA + normalizes the kv_lora_rank latent and leaves k_pe raw, and then ``weight`` + spans only that latent. + """ if IS_CUTLASS_DSL_AVAILABLE and is_sm_100f(): freqs_real = torch.view_as_real(freqs_cis).reshape(-1, freqs_cis.shape[-1], 2) - if is_fused_dspark_rmsnorm_rope_supported(t, weight, freqs_real, num_heads, rope_head_dim): + if is_fused_dspark_rmsnorm_rope_supported( + t, weight, freqs_real, num_heads, rope_head_dim, norm_dim + ): return cute_dsl_dspark_rmsnorm_rope( t, weight, @@ -401,15 +419,19 @@ def _rmsnorm_rope_batched( apply_weight, apply_rmsnorm, inverse_rope, + norm_dim, ) + split_norm = norm_dim is not None and norm_dim != t.shape[-1] + head, tail = (t[..., :norm_dim], t[..., norm_dim:]) if split_norm else (t, None) if apply_rmsnorm: if apply_weight: - t = _rmsnorm(t, weight, eps) + head = _rmsnorm(head, weight, eps) else: - t = t * torch.rsqrt(t.square().mean(-1, keepdim=True) + eps) + head = head * torch.rsqrt(head.square().mean(-1, keepdim=True) + eps) elif apply_weight: - t = (t.float() * weight.float()).to(t.dtype) + head = (head.float() * weight.float()).to(head.dtype) + t = head if tail is None else torch.cat([head, tail], dim=-1) if rope_head_dim > 0: t = _rope_last_dims_batched(t, rope_head_dim, freqs_cis, inverse=inverse_rope) return t @@ -1042,6 +1064,27 @@ def has_heads(self) -> bool: return self.stage_id == self.num_stages - 1 +def _runtime_position_cap(model_config, pretrained_config, slack: int = 0) -> int: + """Positions a RoPE table must span: the served length, not the advertised one. + + Both drafter families build a position table, and both were sized from the + checkpoint's max_position_embeddings. K3 advertises 1,048,576, which costs + ~256 MiB per rank as complex64 -- built via full-size fp32 cos/sin, so + ~768 MiB transient -- while the context cache is bounded by the runtime + max_seq_len. Only the table length is affected; YaRN's correction range is + computed from original_max_position_embeddings and does not move. + """ + runtime = getattr(model_config, "max_seq_len", None) + return int(runtime or getattr(pretrained_config, "max_position_embeddings", 163840)) + slack + + +# NOTE on slack: this is a construction-time cap only. A drafter driven by +# DFlashWorker gets its real bound at runtime from dflash_position_ceiling +# (published as _runtime_position_ceiling), because the engine raises +# max_seq_len past model_config's value and never writes it back. The DSv4 site +# keeps its own slack because it is not driven by that worker. + + class DSv4DSparkDraftModel(nn.Module): """The ``n_mtp_layers``-stage DSpark draft stacked on a DeepSeek-V4 target. @@ -1159,9 +1202,7 @@ def __init__( # batched paths. It is built once per device and gathered/sliced by the # runtime decode positions, so the cache does not grow with sequence # length and the batched consuming op's shape remains static. - self._freqs_cap = ( - int(getattr(config, "max_position_embeddings", 163840)) + self.block_size + 2 - ) + self._freqs_cap = _runtime_position_cap(model_config, config, self.block_size + 2) self._freqs_table_cache: Dict = {} def post_load_weights(self) -> None: @@ -1968,6 +2009,327 @@ def load_weights_from_target_model(self, target_model): self.dspark_model.lm_head = target_model.lm_head +# ---------------------------------------------------------------------------- +# MLA-shaped standalone drafter backbone (Inferact/Kimi-K3-DSpark). +# +# A weight container, not a servable model: ``DFlashForCausalLM`` builds this +# through the registry only to own the drafter's modules, then runs its own +# hand-written block decode over them (see ``MLADSparkForCausalLM``). Nothing +# here registers with the attention backend or the KV cache manager, so the +# layers deliberately have no working ``forward``. +# +# Not a ``DeepseekV3`` reuse: that model carries MoE branches, a fused +# ``qkv_a_proj`` switch, quantization and KV-manager registration, none of +# which this bf16 five-layer drafter has. +# ---------------------------------------------------------------------------- + + +def _yarn_get_mscale(scale: float, mscale: float = 1.0) -> float: + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +def _yarn_find_correction_dim( + num_rotations: float, dim: int, base: float, max_position_embeddings: int +) -> float: + return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / ( + 2 * math.log(base) + ) + + +def _yarn_find_correction_range( + low_rot: float, high_rot: float, dim: int, base: float, max_position_embeddings: int +) -> tuple[int, int]: + low = math.floor(_yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings)) + high = math.ceil(_yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings)) + return max(low, 0), min(high, dim - 1) + + +def _yarn_linear_ramp_mask(minimum: float, maximum: float, dim: int) -> torch.Tensor: + if minimum == maximum: + maximum += 0.001 # Prevent singularity. + linear_func = (torch.arange(dim, dtype=torch.float32) - minimum) / (maximum - minimum) + return torch.clamp(linear_func, 0, 1) + + +def build_dspark_mla_yarn_rope( + *, + dim: int, + base: float, + scaling_factor: float, + original_max_position_embeddings: int, + max_position_embeddings: int, + beta_fast: float = 32.0, + beta_slow: float = 1.0, + mscale: float = 1.0, + mscale_all_dim: float = 0.0, + device: str = "cuda", +) -> tuple[torch.Tensor, torch.Tensor]: + """YaRN cos/sin tables for an MLA drafter's ``qk_rope_head_dim`` slice. + + Transcribed from HF ``DeepseekV3YarnRotaryEmbedding``: the drafters are + distilled under HF/vLLM numerics, so the table is built here rather than + routed through ``RopeParams`` (whose MLA convention carries the fused + kernel's ``duplicate_data`` / GPT-J packing). + + Returns ``(cos, sin)``, both ``[max_position_embeddings, dim]`` fp32, with + the frequency half duplicated so ``rotate_half`` applies. + """ + freq_extra = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + freq_inter = 1.0 / ( + scaling_factor * base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim) + ) + low, high = _yarn_find_correction_range( + beta_fast, beta_slow, dim, base, original_max_position_embeddings + ) + inv_freq_mask = 1.0 - _yarn_linear_ramp_mask(low, high, dim // 2) + inv_freq = freq_inter * (1 - inv_freq_mask) + freq_extra * inv_freq_mask + + t = torch.arange(max_position_embeddings, dtype=torch.float32) + freqs = torch.outer(t, inv_freq) + # HF divides the two mscales; both published drafters set them equal, so + # this is 1.0 there. Kept general because the checkpoint declares both. + scale = _yarn_get_mscale(scaling_factor, mscale) / _yarn_get_mscale( + scaling_factor, mscale_all_dim + ) + emb = torch.cat((freqs, freqs), dim=-1) + return (emb.cos() * scale).to(device), (emb.sin() * scale).to(device) + + +def build_dspark_mla_yarn_freqs_cis(*, device: str = "cuda", **kwargs) -> torch.Tensor: + """The same YaRN table as :func:`build_dspark_mla_yarn_rope`, adjacent-pair. + + Returns ``[max_position_embeddings, dim // 2]`` complex64: the convention the + fused ``cute_dsl_dspark_rmsnorm_rope`` kernel consumes, where the last dim of + the rotated slice is read as (re, im) pairs. + + HF DeepSeek instead de-interleaves and then rotates halves. The two produce + the *same values in a different order*: with even/odd the pair members, + HF writes ``even*cos - odd*sin`` to lane i and ``odd*cos + even*sin`` to lane + i + dim/2, while this one writes them to lanes 2i and 2i+1. That permutation + of the rope slice cancels inside ``q_rope . k_rope``, so the drafter is free + to use either -- but every producer of a rope slice must use the SAME one. + Mixing them changes the scores silently and only shows up as lower AL. + + ``mscale`` rides in the modulus rather than the phase, exactly as the + cos/sin builder folds it into the table. + """ + cos, sin = build_dspark_mla_yarn_rope(device=device, **kwargs) + half = cos.shape[-1] // 2 + return torch.complex(cos[..., :half], sin[..., :half]).contiguous() + + +def apply_dspark_mla_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + """Rotate the trailing ``qk_rope_head_dim`` of ``x`` the HF DeepSeek way. + + ``x`` is ``[..., dim]`` and ``cos``/``sin`` are ``[..., dim]`` already + gathered at the query positions. HF interleaves the rope slice, so it first + reshapes ``(d//2, 2) -> transpose -> (d,)`` and only then applies the + half-split ``rotate_half``; the net rotation is GPT-J style. + """ + d = x.shape[-1] + x = x.reshape(*x.shape[:-1], d // 2, 2).transpose(-1, -2).reshape(*x.shape[:-1], d) + x1, x2 = x[..., : d // 2], x[..., d // 2 :] + rotated = torch.cat((-x2, x1), dim=-1) + return x * cos + rotated * sin + + +class K3DsparkMLA(nn.Module): + """MLA projections for one drafter layer. No attention, no KV registration. + + Head sharding follows :class:`~tensorrt_llm._torch.modules.mla.MLA`: the + low-rank ``*_a`` projections are replicated (every head reads the same + latent), ``q_b_proj`` / ``kv_b_proj`` are column-sharded by head and + ``o_proj`` is row-sharded. Under attention-DP ``tp_size`` collapses to 1, + which is also why the drafter's latent cache is unsharded. + """ + + def __init__(self, model_config, layer_idx: int): + super().__init__() + config = model_config.pretrained_config + self.layer_idx = layer_idx + dtype = config.torch_dtype + + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.q_lora_rank = config.q_lora_rank + self.kv_lora_rank = config.kv_lora_rank + self.qk_nope_head_dim = config.qk_nope_head_dim + self.qk_rope_head_dim = config.qk_rope_head_dim + self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + self.v_head_dim = config.v_head_dim + # What the drafter's KV cache stores per token: one MLA latent, no V. + self.kv_cache_head_dim = self.kv_lora_rank + self.qk_rope_head_dim + + mapping = model_config.mapping + if mapping.cp_size > 1: + # MLA's own module folds cp_size into the head split and the o_proj + # mapping; this drafter does not, and a silently wrong world_size + # would mis-shard the head projections. + raise NotImplementedError( + f"The MLA DSpark drafter does not support context parallelism " + f"(cp_size={mapping.cp_size})." + ) + tp_size = mapping.tp_size + dp_size = 1 + if mapping.enable_attention_dp: + dp_size = tp_size + tp_size = 1 + if self.num_heads % tp_size != 0: + raise ValueError( + f"DSpark MLA drafter has {self.num_heads} heads, not divisible by tp_size {tp_size}." + ) + self.num_heads_tp = self.num_heads // tp_size + attn_mapping = Mapping( + world_size=mapping.pp_size * dp_size * tp_size, + tp_size=tp_size, + pp_size=mapping.pp_size * dp_size, + rank=mapping.rank, + gpus_per_node=mapping.gpus_per_node, + enable_attention_dp=mapping.enable_attention_dp, + ) + self.mapping = attn_mapping + # Row-parallel o_proj reduces only when the heads are actually split. + reduce_output = not mapping.enable_attention_dp and mapping.tp_size > 1 + + skip_init = model_config.skip_create_weights_in_init + self.q_a_proj = Linear( + self.hidden_size, + self.q_lora_rank, + bias=False, + dtype=dtype, + skip_create_weights_in_init=skip_init, + ) + self.q_a_layernorm = RMSNorm( + hidden_size=self.q_lora_rank, eps=config.rms_norm_eps, dtype=dtype + ) + self.q_b_proj = Linear( + self.q_lora_rank, + self.num_heads * self.qk_head_dim, + bias=False, + dtype=dtype, + mapping=attn_mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + skip_create_weights_in_init=skip_init, + allreduce_strategy=model_config.allreduce_strategy, + ) + self.kv_a_proj_with_mqa = Linear( + self.hidden_size, + self.kv_cache_head_dim, + bias=False, + dtype=dtype, + skip_create_weights_in_init=skip_init, + ) + self.kv_a_layernorm = RMSNorm( + hidden_size=self.kv_lora_rank, eps=config.rms_norm_eps, dtype=dtype + ) + self.kv_b_proj = Linear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + dtype=dtype, + mapping=attn_mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + skip_create_weights_in_init=skip_init, + allreduce_strategy=model_config.allreduce_strategy, + ) + self.o_proj = Linear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + dtype=dtype, + mapping=attn_mapping, + tensor_parallel_mode=TensorParallelMode.ROW, + reduce_output=reduce_output, + skip_create_weights_in_init=skip_init, + allreduce_strategy=model_config.allreduce_strategy, + ) + + def forward(self, *args, **kwargs): + raise NotImplementedError( + "K3DsparkMLA holds the drafter's MLA weights for MLADSparkForCausalLM's " + "block decode; it has no standalone attention path." + ) + + +class K3DsparkDecoderLayer(DecoderLayer): + """One drafter layer: MLA projections + dense gated MLP + the two norms.""" + + def __init__(self, model_config, layer_idx: int): + super().__init__() + config = model_config.pretrained_config + self.layer_idx = layer_idx + self.self_attn = K3DsparkMLA(model_config, layer_idx) + self.mlp = GatedMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + bias=False, + dtype=config.torch_dtype, + overridden_tp_size=1 if model_config.mapping.enable_attention_dp else None, + config=model_config, + 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, **kwargs): + raise NotImplementedError( + "K3DsparkDecoderLayer is driven by MLADSparkForCausalLM.dflash_forward, " + "not by the generic decoder loop." + ) + + +class K3DsparkModel(DecoderModel): + def __init__(self, model_config): + super().__init__(model_config) + config = model_config.pretrained_config + # The drafter ships its own trained embedding rather than sharing the + # target's -- verified different weights on Inferact/Kimi-K3-DSpark -- + # so this is loaded from the drafter checkpoint, not overwritten. + 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( + [ + K3DsparkDecoderLayer(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, **kwargs): + raise NotImplementedError( + "K3DsparkModel is a weight container for the DSpark block decode; " + "see MLADSparkForCausalLM.dflash_forward." + ) + + +@register_auto_model("K3DsparkForCausalLM") +class K3DsparkForCausalLM(SpecDecOneEngineForCausalLM[K3DsparkModel, PretrainedConfig]): + """Registry entry for the MLA DSpark drafter checkpoint. + + The name is the one ``DFlashForCausalLM.__init__`` derives from the + checkpoint's ``model_type`` ("k3_dspark"), but MLADSparkForCausalLM pins + ``architectures`` explicitly rather than relying on that derivation -- + the Laguna precedent in ``modeling_dflash.py``. + """ + + def __init__(self, model_config): + super().__init__(K3DsparkModel(model_config), model_config) + + # ---------------------------------------------------------------------------- # Standalone DSpark drafters. # @@ -1991,36 +2353,108 @@ def load_weights_from_target_model(self, target_model): } -class GQADSparkForCausalLM(DFlashForCausalLM): - """DSpark drafter on a GQA-shaped backbone, from a standalone checkpoint. +_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 - Adds the DSpark head set on top of the DFlash block decode: - - the vanilla Markov intra-block logit bias, applied by ``DSparkWorker`` - through :meth:`apply_markov_chain_logits`; - - the ``shift_label`` output convention (the hidden state at block slot j - predicts draft token j+1, so slot 0 holds the anchor token); - - the confidence head weights. +class _MLABlockFixup(NamedTuple): + """Layer-invariant state for the MLA block-decode fixup. - Confidence-scheduled verification is not implemented yet: ``confidence_proj`` - is loaded but unused, and drafting always proposes the full K tokens. + The block's own latents are written into the pool at ``ctx_len + j`` first. + Those slots belong to this request -- the draft KV manager adds + ``max(draft_len, max_total_draft_tokens)`` tokens to every generation + request each step (resource_manager.py:1060) -- and the accepted tokens + overwrite them next step, so the write is transient rather than a claim on + the cache. It is what the GQA drafter's TRTLLM backend already does + (modeling_dflash.py append_paged_kv_cache). - Named for the attention shape, not for a model: the backbone is whatever - the drafter config resolves to through the model registry, and the - inherited block decode works for every GQA family the DFlash drafters - already cover (qwen3, llama, gpt_oss, ...). A per-model subclass would be - empty. The GQA precondition is inherited, not introduced here -- see - ``DFlashForCausalLM._validate_gqa_shape``. An MLA-backboned drafter needs - its own block decode and becomes a sibling, ``MLADSparkForCausalLM``, not a - subclass of this. + With ``seq_lens = ctx_len + block_size`` the kernel's in-block causality + then leaves query j seeing block keys 0..j, so all that is left to fix up is + the block's strict upper triangle -- ``blk_kv`` itself, already in hand. The + earlier form kept the pool read-only and paid for it twice: a gather of the + hidden context tail, and a 15-key rather than 8-key fixup. - Reference: arXiv 2607.05147; deepseek-ai/DeepSpec. + Every field depends only on ctx_len, the page table and the block geometry, + so it is built once per forward rather than once per layer. """ - def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): - super().__init__(draft_config, dflash_attention_backend=dflash_attention_backend) + pages: torch.Tensor # [B, block] page holding each block position + offsets: torch.Tensor # [B, block] offset of that position inside its page + valid: torch.Tensor # [1, block, 1, block] strict upper triangle + page_tables_i32: torch.Tensor # kernel operands, cast once + seq_lens_i32: torch.Tensor # ctx_len + block_size + + +def _build_mla_block_fixup(ctx_len, page_tables, block_size, page_size) -> _MLABlockFixup: + """Where the block's latents go, and the upper-triangle mask, for all layers.""" + device = ctx_len.device + pos = ctx_len.view(-1, 1) + torch.arange(block_size, device=device) + idx = torch.arange(block_size, device=device) + return _MLABlockFixup( + pages=page_tables.gather(1, pos // page_size), + offsets=pos % page_size, + valid=(idx.view(1, block_size) > idx.view(block_size, 1)).view( + 1, block_size, 1, block_size + ), + page_tables_i32=page_tables.to(torch.int32), + seq_lens_i32=(ctx_len + block_size).to(torch.int32), + ) - cfg = draft_config.pretrained_config + +def _resolve_dspark_mla_rope_params(pretrained_config, model_config=None, slack: int = 0) -> Dict: + """RoPE settings for an MLA drafter, from either HF spelling. + + Transformers v5 checkpoints (TorchSpec) carry ``rope_parameters``; older + ones carry ``rope_scaling``. Only YaRN is accepted: both published MLA + drafters use it, and silently treating an unknown scaling as plain RoPE + would cost acceptance without an error. + """ + scaling = ( + getattr(pretrained_config, "rope_parameters", None) + or getattr(pretrained_config, "rope_scaling", None) + or {} + ) + rope_type = scaling.get("rope_type") or scaling.get("type") or "default" + if rope_type != "yarn": + raise ValueError( + f"MLA DSpark drafter declares rope_type={rope_type!r}; only 'yarn' is " + "supported (the block decode builds its own YaRN table)." + ) + theta = scaling.get("rope_theta") or getattr(pretrained_config, "rope_theta", 10000.0) + return { + "theta": float(theta), + "scaling_factor": float(scaling["factor"]), + "original_max_positions": int(scaling["original_max_position_embeddings"]), + # Table length, not YaRN math: the correction range below uses + # original_max_positions, so capping this at the served length is free. + # K3 advertises 1,048,576, which would build a ~256 MiB complex64 table + # per rank -- through full-size fp32 cos/sin, so ~768 MiB transient -- + # for a drafter whose point is 5760 B/token instead of 20480. + "max_positions": _runtime_position_cap(model_config, pretrained_config, slack), + "beta_fast": float(scaling.get("beta_fast", 32.0)), + "beta_slow": float(scaling.get("beta_slow", 1.0)), + "mscale": float(scaling.get("mscale", 1.0)), + "mscale_all_dim": float(scaling.get("mscale_all_dim", 0.0)), + } + + +class _DSparkHeadMixin: + """The DSpark head set, shared by the GQA and MLA drafter wrappers. + + DSpark is DFlash plus three things, and only these three are common to both + backbone shapes: the vanilla Markov intra-block logit bias, the + ``shift_label`` slot convention (block slot j predicts draft token j+1, so + slot 0 holds the anchor) and the confidence head. The block decode itself + has nothing in common between the shapes, which is why this is a mixin and + not a base class. + + Confidence-scheduled verification is not implemented yet: ``confidence_proj`` + is loaded but unused, and drafting always proposes the full K tokens. + """ + + def _init_dspark_heads(self, cfg) -> None: # Defaults on, unlike the DFlash base: the shift_label slot layout is # part of what DSpark *is*, and both published drafters set # block_size == max_draft_len, where the DFlash layout (slots 1..K) @@ -2096,11 +2530,12 @@ def apply_markov_chain_logits( base_logits, first_prev_tokens, self.markov_w1, markov_w2, argmax_fn=argmax_fn ) - def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): - """Take the DSpark head weights, then hand the rest to DFlash. + def _take_dspark_head_weights(self, weights: Dict) -> tuple[Dict, Dict]: + """Pull the head tensors out of ``weights`` and load them onto self. The head keys are pulled out before the backbone remap: left in, they would pick up a ``model.`` prefix and be dropped by partial loading. + Returns the remaining weights and the extracted head weights. """ dspark_weights = {} consumed = set() @@ -2144,9 +2579,552 @@ def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): self.confidence_proj_weight = dspark_weights["confidence_proj.weight"].to("cuda") if "confidence_proj.bias" in dspark_weights: self.confidence_proj_bias = dspark_weights["confidence_proj.bias"].to("cuda") + return weights, dspark_weights + + +class GQADSparkForCausalLM(_DSparkHeadMixin, DFlashForCausalLM): + """DSpark drafter on a GQA-shaped backbone, from a standalone checkpoint. + + Adds the DSpark head set (see :class:`_DSparkHeadMixin`) on top of the + DFlash block decode. + + Named for the attention shape, not for a model: the backbone is whatever + the drafter config resolves to through the model registry, and the + inherited block decode works for every GQA family the DFlash drafters + already cover (qwen3, llama, gpt_oss, ...). A per-model subclass would be + empty. The GQA precondition is inherited, not introduced here -- see + ``DFlashForCausalLM._validate_gqa_shape``. The MLA-backboned drafter is the + sibling :class:`MLADSparkForCausalLM`, not a subclass of this. + + Reference: arXiv 2607.05147; deepseek-ai/DeepSpec. + """ + + def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): + super().__init__(draft_config, dflash_attention_backend=dflash_attention_backend) + self._init_dspark_heads(draft_config.pretrained_config) + + 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) return super().load_weights(weights, weight_mapper=weight_mapper, **kwargs) +class MLADSparkForCausalLM(_DSparkHeadMixin, DFlashForCausalLM): + """DSpark drafter on an MLA-shaped backbone, from a standalone checkpoint. + + The sibling ``modeling_dspark``'s GQA class could not be: DFlash's block + decode splits a fused ``qkv_proj`` by (num_heads, num_kv_heads, head_dim) + and fuses per-head K/V across layers, and MLA has neither. What it keeps is + the DFlash *contract* -- dual-source KV (context from the projected target + hidden, block from the draft hidden), non-causal block attention, one + precomputed context cache -- and the DSpark head set on top. + + The cache stores one ``kv_lora_rank + qk_rope_head_dim`` latent per token + per layer and no V, which is the point: 5 x 576 x 2B = 5760 B/token/rank + against the GQA drafter's 20480 under attention-DP, where KV heads are + unsharded. + + Block attention runs on flashinfer's trtllm-gen absorbed-MLA paged decode + (``_mla_paged_attention``); the eager torch path stays only as the + reference, for builds without flashinfer and for the unpaged arena. + ``spec_config.attention_backend`` does NOT select between those two -- it + picks the *GQA* drafter's backend in the shared DFlash worker, and this + class overrides the attention entirely. + + Reference: arXiv 2607.05147; deepseek-ai/DeepSpec. + """ + + # Attention is _mla_paged_attention, so neither worker backend applies and + # neither backend's shape checks bind -- the absorbed shape (64 heads : 1 KV + # head, head_dim 576) would fail both. + _uses_worker_attention_backend = False + # The whole point of the MLA drafter: its context KV comes out of the + # manager's pool (5760 B/token/rank) instead of an arena dense in + # max_seq_len that free_gpu_memory_fraction never bounds. Independent of the + # backend field, which is why either value of it works here. + _paged_ctx_cache = True + + def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): + cfg = draft_config.pretrained_config + # Pin the backbone instead of relying on the model_type-derived name + # (the Laguna precedent in modeling_dflash.py): the checkpoint labels + # itself "K3DSparkModel", which is registered nowhere. + cfg.architectures = ["K3DsparkForCausalLM"] + super().__init__(draft_config, dflash_attention_backend=dflash_attention_backend) + self._init_dspark_heads(cfg) + + attn0 = self.model.layers[0].self_attn + self.qk_nope_head_dim = attn0.qk_nope_head_dim + self.qk_rope_head_dim = attn0.qk_rope_head_dim + self.kv_lora_rank = attn0.kv_lora_rank + self.v_head_dim = attn0.v_head_dim + # The cache contract the worker sizes its arena from. Config-derived, so + # it is published here rather than waiting for the lazy buffer build: + # one MLA latent per token per layer, no V half. + self._num_attn_layers = len(self.model.layers) + self._num_heads = attn0.num_heads_tp + self._num_kv_heads = 1 + self._head_dim = attn0.kv_cache_head_dim + self._kv_factor = 1 + + # No slack: the lookahead lives in dflash_position_ceiling, which the + # worker publishes as _runtime_position_ceiling before any forward and + # _mla_freqs_cis builds the table from. This config-derived cap is only + # the fallback for direct construction, where no worker runs and the + # positions asked for are whatever the caller passes. + rope = _resolve_dspark_mla_rope_params(cfg, draft_config) + self._mla_rope_params = rope + self._mla_workspace = None + self._mla_freqs = None + self._mla_noop_weight = None + # HF DeepSeek folds YaRN's attention scaling into the softmax scale: + # softmax_scale = mscale^2 / sqrt(qk_nope + qk_rope). NOT 1/sqrt(576) -- + # the absorbed product reproduces the 192-dim un-absorbed score. + mscale = _yarn_get_mscale(rope["scaling_factor"], rope["mscale_all_dim"]) + self.softmax_scale = (mscale * mscale) / math.sqrt( + self.qk_nope_head_dim + self.qk_rope_head_dim + ) + + # -- shape / buffers --------------------------------------------------- + + def _validate_gqa_shape(self): + """Replace the base's GQA precondition with the MLA one.""" + layers = getattr(self.model, "layers", None) + if not layers: + return + attn0 = layers[0].self_attn + keys = ( + "num_heads_tp", + "qk_nope_head_dim", + "qk_rope_head_dim", + "kv_lora_rank", + "v_head_dim", + ) + base = tuple(getattr(attn0, k) for k in keys) + for idx, layer in enumerate(layers[1:], start=1): + attn = getattr(layer, "self_attn", None) + if attn is None or not hasattr(attn, "kv_a_proj_with_mqa"): + raise ValueError( + f"MLA DSpark block decode needs self_attn.kv_a_proj_with_mqa on every " + f"draft layer, but layer {idx} of {type(self.config).__name__} has none." + ) + if tuple(getattr(attn, k) for k in keys) != base: + raise ValueError( + "MLA DSpark fuses the latent projection across layers and needs one " + f"uniform MLA shape, but layer {idx} differs from layer 0 " + f"({dict(zip(keys, base))})." + ) + + def _init_rope(self): + """No shared RoPE cache: the MLA path rotates only the rope slice.""" + self._rope_initialized = True + + @staticmethod + @lru_cache(maxsize=1) + def _mla_decode_op(): + """flashinfer's absorbed-MLA paged decode, or None when unavailable. + + Loaded lazily so a build without flashinfer still runs the eager path. + """ + try: + import flashinfer + + return flashinfer.mla.trtllm_batch_decode_with_kv_cache_mla + except (ImportError, AttributeError): + return None + + def _mla_rope_noop_weight(self, like: torch.Tensor) -> torch.Tensor: + """Placeholder weight for a rope-only call; never read by the kernel.""" + if self._mla_noop_weight is None: + self._mla_noop_weight = torch.ones(like.shape[-1], device=like.device, dtype=like.dtype) + return self._mla_noop_weight + + def _mla_decode_workspace(self, device): + if self._mla_workspace is None: + self._mla_workspace = torch.zeros( + _MLA_DECODE_WORKSPACE_BYTES, dtype=torch.uint8, device=device + ) + return self._mla_workspace + + def _mla_paged_attention(self, query, blk_kv, layer_cache, fixup, block_size, max_kv_len): + """Absorbed-MLA block attention: paged kernel + a block-local fixup. + + The kernel is causal across the query block and takes no mask (measured, + not assumed: it matches a causal reference to 1.6e-4 against 1.5e-2 for + full visibility). Its LSE is base 2 (also measured: exact against + ``log2(sum exp)``, 1.84 off against the natural log). + + The block's own latents go into the pool first, at ``ctx_len + j`` -- + slots the draft KV manager already reserves for this request each step + (see _MLABlockFixup). With ``seq_lens = ctx_len + block_size`` the kernel + then covers context AND block, and in-block causality leaves only the + block's strict upper triangle: 8 keys, no context-tail gather. Those are + attended eagerly against ``blk_kv``, already in hand, and merged by + log-sum-exp. + + Writing into the pool is safe because dflash.py bounds the advertised + context at ``allocated - block_size``; without that the first block + position lands on the first unallocated page, whose table entry was + clamped to physical block 0 -- another request's. + + Returns the latent-space output ``[B, block, heads, kv_lora_rank]``. + """ + decode = self._mla_decode_op() + batch, _, num_heads, _ = query.shape + device = query.device + + # --- the block's own latents into the pool, then one causal pass --- + pool = layer_cache[:, 0, 0] # [pages, page_size, head_dim] + pool[fixup.pages.reshape(-1), fixup.offsets.reshape(-1)] = blk_kv.reshape( + -1, blk_kv.shape[-1] + ) + out_lower = torch.empty( + batch, block_size, num_heads, self.kv_lora_rank, dtype=query.dtype, device=device + ) + lse_lower = torch.empty(batch * block_size, num_heads, dtype=torch.float32, device=device) + decode( + query, + layer_cache[:, 0], # [pages, 1, page, head_dim] + self._mla_decode_workspace(device), + self.qk_nope_head_dim, + self.kv_lora_rank, + self.qk_rope_head_dim, + fixup.page_tables_i32, + fixup.seq_lens_i32, + max_kv_len, + 0, + out_lower, + self.softmax_scale, + 1.0, + None, + None, + None, + backend="trtllm-gen", + is_var_seq=True, + uses_shared_paged_kv_idx=True, + lse=lse_lower, + return_lse=True, + ) + + # --- fixup: the block's strict upper triangle, the only thing the + # kernel's in-block causality hides once the block is in the pool --- + scores = torch.einsum("bshc,btc->bsht", query, blk_kv).float() * self.softmax_scale + scores = scores.masked_fill(~fixup.valid, float("-inf")) + lse_upper = scores.logsumexp(-1) + probs = torch.softmax(scores, dim=-1).nan_to_num_().to(blk_kv.dtype) + out_upper = torch.einsum("bsht,btc->bshc", probs, blk_kv[..., : self.kv_lora_rank]) + + # --- log-sum-exp merge (flash-decoding combine), kernel LSE is base 2 --- + # The weight pair collapses: w_u / (w_l + w_u) == sigmoid(lse_u - lse_l), + # so the merge is one sigmoid and one lerp instead of a peak subtraction, + # two exps and a division -- and it needs no explicit max for stability. + # The last query has an empty triangle, hence lse_upper -inf, sigmoid 0, + # pure kernel output. + lse_l = lse_lower.view(batch, block_size, num_heads).float() * _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) + + def _mla_freqs_cis(self, device): + """Adjacent-pair YaRN table, cached. See build_dspark_mla_yarn_freqs_cis. + + Sized from the ceiling the worker publishes once it knows the runtime + one (_runtime_position_ceiling, set in DFlashDrafter._lazy_init_ctx_buffers + before any forward), because that is the value ctx_len is clamped to and + it is strictly above model_config.max_seq_len whenever spec decoding is + on. The config-derived cap is the fallback for direct construction in + tests, where no worker runs. + """ + if self._mla_freqs is None: + rope = dict(self._mla_rope_params) + runtime_cap = getattr(self, "_runtime_position_ceiling", None) + if runtime_cap is not None: + # Outright, not max(): the published ceiling is already + # min(runtime, max_position_embeddings) + lookahead, so it is + # both sufficient and the only thing that keeps the table small + # when max_seq_len is unset and the config cap is the + # checkpoint's advertised 1,048,576. + rope["max_positions"] = int(runtime_cap) + self._mla_freqs = build_dspark_mla_yarn_freqs_cis( + dim=self.qk_rope_head_dim, + base=rope["theta"], + scaling_factor=rope["scaling_factor"], + original_max_position_embeddings=rope["original_max_positions"], + max_position_embeddings=rope["max_positions"], + beta_fast=rope["beta_fast"], + beta_slow=rope["beta_slow"], + mscale=rope["mscale"], + mscale_all_dim=rope["mscale_all_dim"], + device=device, + ) + return self._mla_freqs + + def _build_fused_kv_buffers(self) -> None: + """Stack the per-layer latent projection and its RMSNorm weight. + + The MLA analogue of the base's fused K/V GEMM: one ``[L*576, hidden]`` + weight, plus the ``[L, kv_lora_rank]`` layernorm scales applied after. + """ + if self._fused_kv_weight is not None: + return + layers_attn = [layer.self_attn for layer in self.model.layers] + attn0 = layers_attn[0] + self._fused_kv_weight = torch.cat( + [a.kv_a_proj_with_mqa.weight for a in layers_attn], dim=0 + ).contiguous() + self._fused_kv_bias = None + self._kv_a_norm_stacked = torch.stack([a.kv_a_layernorm.weight.data for a in layers_attn]) + self._kv_a_norm_eps = attn0.kv_a_layernorm.variance_epsilon + self._input_ln_eps = None + self._has_qk_norm = False + self._use_fused_qk_norm_rope = False + + # Split kv_b_proj into the absorption operands once. K absorbs into the + # query, V un-absorbs the attention output. + nh = attn0.num_heads_tp + for a in layers_attn: + w = a.kv_b_proj.weight.view(nh, self.qk_nope_head_dim + self.v_head_dim, -1) + a._k_b_proj = w[:, : self.qk_nope_head_dim].contiguous() + a._v_b_proj = w[:, self.qk_nope_head_dim :].contiguous() + + logger.debug( + f"MLA DSpark: fused latent projection built for {self._num_attn_layers} layers " + f"(weight={tuple(self._fused_kv_weight.shape)}, head_dim={self._head_dim})" + ) + + # -- context cache ----------------------------------------------------- + + def precompute_context_kv(self, projected_hidden, positions): + """Post-norm / post-RoPE MLA latent for ALL drafter layers in one GEMM. + + Returns ``(latent, None)``: ``[N, L, 1, kv_lora_rank + qk_rope_head_dim]`` + and no V half, which is what ``_kv_factor == 1`` means to the worker. + """ + if self._fused_kv_weight is None: + self._build_fused_kv_buffers() + n = projected_hidden.shape[0] + num_layers = self._num_attn_layers + weight_dtype = self._fused_kv_weight.dtype + if projected_hidden.dtype != weight_dtype: + projected_hidden = projected_hidden.to(weight_dtype) + + latent = F.linear(projected_hidden, self._fused_kv_weight) + latent = latent.view(n, num_layers, self._head_dim) + + # One fused RMSNorm+RoPE per layer, not one eager chain for all of them: + # the kernel takes a single norm weight per call and the layers do not + # share one. Same op the block path uses, so the two halves of an + # attention stay numerically consistent -- which is the point here, not + # speed (this runs per request, not per decode step). + freqs = self._mla_freqs_cis(latent.device)[positions.view(-1).long()] + out = torch.empty_like(latent) + for layer_idx in range(num_layers): + # 3D throughout: the eager fallback's rotary indexes [G, s, rd]. + out[:, layer_idx] = _rmsnorm_rope_batched( + latent[:, layer_idx].unsqueeze(0).contiguous(), + self._kv_a_norm_stacked[layer_idx], + self._kv_a_norm_eps, + self.qk_rope_head_dim, + freqs.unsqueeze(0), + norm_dim=self.kv_lora_rank, + ).squeeze(0) + return out.unsqueeze(2).contiguous(), None + + # -- block decode ------------------------------------------------------ + + def dflash_forward( + self, + noise_embedding, + query_positions, + num_ctx_per_req, + ctx_k_cache, + ctx_v_cache, + ctx_cache_batch_idx, + ctx_kv_cache=None, + ctx_page_table=None, + ): + """Eager absorbed-MLA block decode over the drafter's latent cache. + + Args mirror the base: ``ctx_k_cache`` is + ``[pool_slots, L, max_ctx+block_size, 1, kv_cache_head_dim]`` and + ``ctx_v_cache`` is unused (``None``). + """ + if self._fused_kv_weight is None: + self._build_fused_kv_buffers() + + batch, block_size = noise_embedding.shape[0], noise_embedding.shape[1] + device = noise_embedding.device + num_heads = self._num_heads + latent_dim = self.kv_lora_rank + + # [B, blk, rope/2] complex: one phase per draft position, shared by the + # query and the block latent of every layer. + blk_freqs = self._mla_freqs_cis(device)[query_positions.long()] + + slots = ctx_cache_batch_idx.long() + # Context pages for this batch, resolved once: rows are batch positions + # when bound to the manager's per-request block table and slots on the + # private arena, which is exactly what ctx_cache_batch_idx already is. + page_tables = ( + None if ctx_page_table is None else ctx_page_table.index_select(0, slots).long() + ) + + # Valid-key mask, shared by every layer: context slots below this + # request's length, then the whole block (DSpark is non-causal). Pages + # past a request's allocation hold another request's data, so the mask + # is what keeps them out -- not the block table. + page_size = None if page_tables is None else ctx_kv_cache[0].shape[-2] + capacity = ctx_k_cache.shape[2] if page_tables is None else page_tables.shape[1] * page_size + # The paged kernel is the fast path; the eager gather stays as the + # reference and covers builds without flashinfer and the unpaged arena. + use_kernel = page_tables is not None and self._mla_decode_op() is not None + # Kernel path: index and mask operands are layer-invariant, so build + # them once here rather than five times inside the decode loop. + fixup = ( + _build_mla_block_fixup(num_ctx_per_req[:batch], page_tables, block_size, page_size) + if use_kernel + else None + ) + logger.info_once( + "MLA DSpark block decode: " + + ( + "trtllm-gen paged kernel, block in pool + upper-triangle fixup" + if use_kernel + else f"eager (paged={page_tables is not None}, " + f"flashinfer={self._mla_decode_op() is not None})" + ), + key="mla_dspark_block_decode_variant", + ) + # Eager path only: dense key mask over the gathered context. + if not use_kernel: + ctx_valid = torch.arange(capacity, device=device).unsqueeze(0) < num_ctx_per_req[ + :batch + ].view(-1, 1) + key_valid = torch.cat( + [ctx_valid, torch.ones(batch, block_size, dtype=torch.bool, device=device)], dim=1 + ) + hidden_states = noise_embedding + residual = None + for layer_idx, layer in enumerate(self.model.layers): + attn = layer.self_attn + hs_flat = hidden_states.reshape(-1, hidden_states.shape[-1]) + if residual is None: + residual = hidden_states.clone() + hs_normed = layer.input_layernorm(hs_flat) + else: + res_flat = residual.reshape(-1, residual.shape[-1]) + hs_normed, res_flat = layer.input_layernorm(hs_flat, res_flat) + residual = res_flat.reshape(batch, block_size, -1) + + q = attn.q_b_proj(attn.q_a_layernorm(attn.q_a_proj(hs_normed))) + q = q.view(batch, block_size, num_heads, self.qk_nope_head_dim + self.qk_rope_head_dim) + # Rotate the trailing rope slice of every head in one kernel. No + # norm and no weight here -- q_a_layernorm already ran -- but the + # fused op still shape-checks the weight, hence the dummy. + q = _rmsnorm_rope_batched( + q, + self._mla_rope_noop_weight(q), + 0.0, + self.qk_rope_head_dim, + blk_freqs, + num_heads=num_heads, + apply_weight=False, + apply_rmsnorm=False, + ) + q_nope, q_rope = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + # Absorb the nope half into latent space so the block attends the + # stored latent directly (MQA), instead of expanding it per head. + q_absorbed = torch.einsum("bshd,hdc->bshc", q_nope, attn._k_b_proj.to(q_nope.dtype)) + query = torch.cat([q_absorbed, q_rope], dim=-1) + + # The block's own latent comes from the draft hidden; the context's + # from the projected target hidden (precompute_context_kv). That + # dual-source split is the DFlash contract, not an MLA detail. + blk_latent = attn.kv_a_proj_with_mqa(hs_normed).view(batch, block_size, -1) + # RMSNorm over the latent, RoPE over k_pe, one kernel. norm_dim keeps + # k_pe out of the reduction, which is what DeepSeek-style MLA wants + # and what DSpark's own convention (whole row) does not do. + blk_kv = _rmsnorm_rope_batched( + blk_latent, + attn.kv_a_layernorm.weight, + attn.kv_a_layernorm.variance_epsilon, + self.qk_rope_head_dim, + blk_freqs, + norm_dim=latent_dim, + ) + if use_kernel: + ctx_out = self._mla_paged_attention( + query, + blk_kv, + ctx_kv_cache[layer_idx], + fixup, + block_size, + capacity, + ) + else: + if page_tables is None: + ctx = ctx_k_cache[slots, layer_idx, :, 0, :] + else: + # [pages, kv_factor=1, nkv=1, page, hd] -> gather this + # batch's pages and flatten back to a per-request context. + ctx = ctx_kv_cache[layer_idx][:, 0, 0][page_tables].reshape(batch, capacity, -1) + kv_full = torch.cat([ctx, blk_kv], dim=1) + + scores = torch.einsum("bshc,btc->bsht", query, kv_full).float() * self.softmax_scale + scores = scores.masked_fill(~key_valid.view(batch, 1, 1, -1), float("-inf")) + probs = torch.softmax(scores, dim=-1).to(kv_full.dtype) + ctx_out = torch.einsum("bsht,btc->bshc", probs, kv_full[..., :latent_dim]) + attn_out = torch.einsum("bshc,hvc->bshv", ctx_out, attn._v_b_proj.to(ctx_out.dtype)) + + hidden_out = attn.o_proj(attn_out.reshape(batch * block_size, -1)) + res_flat = residual.reshape(-1, residual.shape[-1]) + hidden_out, res_flat = layer.post_attention_layernorm(hidden_out, res_flat) + hidden_out = layer.mlp(hidden_out) + hidden_states = hidden_out.reshape(batch, block_size, -1) + residual = res_flat.reshape(batch, block_size, -1) + + hidden_states_out, _ = self.model.norm( + hidden_states.reshape(-1, hidden_states.shape[-1]), + residual.reshape(-1, residual.shape[-1]), + ) + return hidden_states_out + + # -- weights ----------------------------------------------------------- + + #: Only the head is shared. This drafter ships its own embed_tokens whose + #: values differ from the target's -- see load_weights_from_target_model -- + #: so a checkpoint missing it must fail rather than run on random weights. + WEIGHTS_SHARED_WITH_TARGET = ("lm_head",) + + def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): + """Take the DSpark heads, rename the TorchSpec spellings, then load. + + The TorchSpec drafters name the capture projection ``context_proj`` / + ``context_norm`` and the final norm ``final_norm``, where the SpecForge + ones use ``fc`` / ``hidden_norm`` / ``norm``. Everything below that is + already HF-standard MLA. + """ + weights, _ = self._take_dspark_head_weights(weights) + renames = { + "context_proj.weight": "fc.weight", + "context_norm.weight": "hidden_norm.weight", + "final_norm.weight": "norm.weight", + } + weights = {renames.get(k, k): v for k, v in weights.items()} + return super().load_weights(weights, weight_mapper=weight_mapper, **kwargs) + + def load_weights_from_target_model(self, target_model: torch.nn.Module) -> None: + """Share only ``lm_head``; the drafter has its own trained embedding. + + Inferact/Kimi-K3-DSpark ships ``embed_tokens`` whose values differ from + the target's, so adopting the target's would feed the block decode + embeddings it was not distilled against -- silently, at some acceptance + cost. The GQA drafters ship no embedding and keep the base behavior. + """ + self.draft_model_full.lm_head = target_model.lm_head + self.lm_head = target_model.lm_head + + def draft_is_embedded_in_target(model_config) -> bool: """True when the DSpark draft weights live inside the target checkpoint. @@ -2195,15 +3173,15 @@ def _build_dspark_draft(model_config, draft_config, lm_head, model): draft_moe_backend=getattr(model_config.spec_config, "moe_backend", None), ) - # No per-model_type table here. ``DFlashForCausalLM.__init__`` already - # resolves the backbone from the drafter config through the model registry, - # so keying on model_type a second time would only duplicate that dispatch - # and force a new entry for every GQA family that already works. What the - # table really guarded was the block decode's GQA precondition, which is now - # checked where it belongs, in the DFlash base. An MLA-backboned drafter - # (e.g. Inferact/Kimi-K3-DSpark) fails that check with a clear message until - # ``MLADSparkForCausalLM`` lands as a sibling. - return GQADSparkForCausalLM( + # Dispatch on the attention shape, not on model_type: the block decode is + # 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. + drafter_cls = ( + MLADSparkForCausalLM if is_mla(draft_config.pretrained_config) else GQADSparkForCausalLM + ) + return drafter_cls( draft_config, dflash_attention_backend=model_config.spec_config.attention_backend, ) @@ -2216,6 +3194,13 @@ def _build_dspark_draft(model_config, draft_config, lm_head, model): "DSv4DSparkForCausalLM", # Standalone flavour. "GQADSparkForCausalLM", + "MLADSparkForCausalLM", + # MLA drafter backbone (a weight container for the block decode). + "K3DsparkForCausalLM", + "K3DsparkModel", + "build_dspark_mla_yarn_rope", + "build_dspark_mla_yarn_freqs_cis", + "apply_dspark_mla_rope", "draft_is_embedded_in_target", "validate_dspark_eplb_layer_base", "validate_dspark_eplb_stage_layers", diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 38dad076d57f..c771cf2c9e0c 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -352,6 +352,44 @@ def _is_mla_layer(cfg, layer_idx: int) -> bool: # --------------------------------------------------------------------------- +KIMI_K3_AUX_ATTN_RES_STREAM_ENV = "KIMI_K3_AUX_ATTN_RES_STREAM" +"""Which residual-stream value the DFlash/DSpark hidden-state tap captures. + +``1`` (default) captures the pre-norm attn_res mixture -- the value the next +consumer actually reads. ``0`` captures the raw running prefix sum instead. + +Both conventions exist in the wild and a drafter distilled against one scores +lower on the other with nothing raised, so this is a property of the DRAFTER +checkpoint, not a performance knob. SGLang (and therefore RadixArk/Kimi-K3-DSpark) +uses the mixture: ``kimi_k3.py _dspark_capture_stream`` -> ``attn_residual.py +aggregate_stream``. vLLM implements both and defaults to the prefix +(``VLLM_KIMI_K3_AUX_ATTN_RES_STREAM=0``, ``models/kimi_k3/nvidia/model.py +_capture_aux_hidden_stream``), which is what a TorchSpec-distilled drafter may +have been trained against. Measured cost of getting it wrong on K3 + RadixArk: +AR 71.4% -> 66.9%. + +Per-checkpoint measurements on K3, GSM8K AL, n=200, TEP8: + +=================== ================== ============ ====== +drafter stream (default 1) prefix (0) delta +=================== ================== ============ ====== +RadixArk (GQA) 71.8% -- -- +Inferact (MLA) 65.7% 66.6% +0.9pt +=================== ================== ============ ====== + +So the default is right for RadixArk. For Inferact the prefix convention +matches its vLLM/TorchSpec lineage and measures better, but +0.9pt at n=200 is +inside this harness's noise band (it treats RadixArk's own 71-73% spread as +noise), so this is a direction, not a settled requirement -- unlike the 4.5pt +RadixArk case above, which was unambiguous. Both Inferact acc_len values (5.60 +and 5.66) sit on its model card's 5.64, so neither convention is grossly wrong +for it. + +Deriving this from checkpoint metadata is not possible today: neither published +drafter's config records which capture convention it was distilled against.""" + +_AUX_ATTN_RES_STREAM_ENABLED = os.environ.get(KIMI_K3_AUX_ATTN_RES_STREAM_ENV, "1") == "1" + KIMI_K3_FUSED_ATTN_RES_ENV = "KIMI_K3_FUSED_ATTN_RES" """Set to ``0`` to disable the in-tree fused Torch op ``trtllm::attn_res_fwd`` (Blackwell only). Default: fused with fallback.""" @@ -1675,7 +1713,15 @@ def forward( self.self_attention_res_norm, ) if capture is not None: - capture[0].maybe_capture_hidden_states(capture[1], hidden_states, None) + # Which residual value the drafter was distilled against is a + # property of the DRAFTER checkpoint, not a tuning knob: a mismatch + # only lowers acceptance, silently. ``hidden_states`` here is the + # pre-norm attn_res mixture (SGLang's aggregate_stream); prefix_only + # wants the incoming running prefix, which vLLM builds as + # ``prefix_sum + pending_mlp_out`` on its deferred-add path and is + # already in hand as ``prefix_sum``. + tapped = hidden_states if _AUX_ATTN_RES_STREAM_ENABLED else prefix_sum + capture[0].maybe_capture_hidden_states(capture[1], tapped, None) if self.layer_idx % self.attn_res_block_size == 0: block_residual[num_snapshots].copy_(prefix_sum) @@ -1766,6 +1812,16 @@ def __init__(self, model_config: ModelConfig): cfg.num_hidden_layers + cfg.attn_res_block_size - 1 ) // cfg.attn_res_block_size + # Which convention the drafter tap is on is not recoverable from the + # served output -- a mismatch only lowers acceptance -- so state it once + # at construction rather than leaving it to be inferred from an AL. + logger.info_once( + "Kimi K3 aux hidden capture: mode=" + f"{'attn_res_stream' if _AUX_ATTN_RES_STREAM_ENABLED else 'prefix_only'} " + f"({KIMI_K3_AUX_ATTN_RES_STREAM_ENV}={int(_AUX_ATTN_RES_STREAM_ENABLED)})", + key="kimi_k3_aux_capture_mode", + ) + def forward( self, attn_metadata: AttentionMetadata, @@ -1830,7 +1886,7 @@ def forward( self.output_attn_res_proj, self.output_attn_res_norm, ) - if num_snapshots > 0 + if num_snapshots > 0 and _AUX_ATTN_RES_STREAM_ENABLED else hidden_states ) spec_metadata.maybe_capture_hidden_states(last.layer_idx, tail, None) diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 5af73e0add7c..4d80300d2567 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1405,6 +1405,18 @@ def external_drafter_config_kwargs(model_config, spec_config) -> dict: spec_config=None, # Avoid recursive spec-dec max_num_tokens=model_config.max_num_tokens, moe_max_num_tokens=model_config.moe_max_num_tokens, + # Bounds the drafter's position tables. Without it the field stays + # None and they fall back to the checkpoint's advertised + # max_position_embeddings -- 1,048,576 for K3, a ~256 MiB complex64 + # table per rank for a context the runtime bounds far below that. + # + # The user's value, NOT the engine's. py_executor_creator raises + # model_engine_max_seq_len past this and never writes it back, so a + # drafter that indexes absolute positions must read the raised value at + # runtime (DFlashDrafter publishes it as _runtime_position_ceiling) + # rather than have this line predict it -- reproducing that arithmetic + # here is what let the two drift apart in the first place. + max_seq_len=model_config.max_seq_len, ) # Only the embedded DSpark draft shares the target's EPLB namespace (its # stages are target decoder blocks registered into the target's balancer). diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 0c16cb0972f3..9f19293deae8 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -1512,6 +1512,17 @@ def run_concurrently(func, raise +#: Modules stored unfused in a checkpoint and fused in the module tree, mapped +#: to the components they are built from. This is the table the mapper-less +#: load path below uses; HfWeightMapper.map_weights installs the same pairs for +#: the mapper path. Anything reasoning about what a checkpoint must provide has +#: to read one of the two rather than keep a third copy. +FUSED_MODULE_COMPONENTS = { + 'qkv_proj': ['q_proj', 'k_proj', 'v_proj'], + 'gate_up_proj': ['gate_proj', 'up_proj'], +} + + def _load_weights_impl(model: Union[nn.Module, DecoderModelForCausalLM], weights: Dict, skip_modules: List[str] = [], @@ -1538,10 +1549,7 @@ def _load_weights_impl(model: Union[nn.Module, DecoderModelForCausalLM], model.config, 'num_key_value_heads' ) and model.config.num_key_value_heads is not None else model.config.num_attention_heads - params_map = { - 'qkv_proj': ['q_proj', 'k_proj', 'v_proj'], - 'gate_up_proj': ['gate_proj', 'up_proj'] - } + params_map = dict(FUSED_MODULE_COMPONENTS) device_id = local_mpi_rank() def load_single_module(name, module): diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index 57f70ff2ca22..067ca2ca0a23 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -794,6 +794,7 @@ def __getitem__(self, key): deepseek_v32="DeepseekV3Config", kimi_k2="DeepseekV3Config", glm_moe_dsa="DeepseekV3Config", + k3_dspark="K3DsparkConfig", laguna="LagunaConfig", ) # NOTE: HF config.json uses deepseek_v32 as model_type but with same DSV3 config class diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index dcf39fa4df67..c5ac3689b302 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -71,6 +71,24 @@ def dflash_draft_slot_ids( return (request_bases.unsqueeze(1) + first_slot + offsets.unsqueeze(0)).flatten() +def dflash_position_ceiling(max_ctx: int, block_size: int, max_draft_len: int) -> int: + """Positions a drafter that indexes absolute positions must be able to encode. + + The forward reads two position sets past the accumulated context, both + starting from a ctx_len clamped to ``max_ctx``: + + * block decode ``ctx_len + num_accepted + j``, ``j in [0, block_size)`` + * context KV ``ctx_len + [0, max_draft_len]`` + + so the largest index is ``max_ctx + max(block_size, max_draft_len + 1) - 1`` + and a table needs one entry more than that. Derived from the runtime + ``max_ctx`` rather than from model_config: py_executor_creator raises the + engine's max_seq_len past the configured value and never writes it back, so + any config-derived reconstruction drifts from what ctx_len is clamped to. + """ + return int(max_ctx) + max(int(block_size), int(max_draft_len) + 1) + + @dataclass class DFlashSpecMetadata(SpecMetadata): """Metadata for DFlash speculative decoding. @@ -259,6 +277,9 @@ def __init__( self._max_ctx = 0 self._ctx_k_buf = None # [max_batch+1, L, max_ctx+block, nkv, hd] self._ctx_v_buf = None + # Set by _lazy_init_ctx_buffers; every context write switches on the + # cache layout rather than on which kernel reads it. + self._ctx_paged = False # [L, pages, K/V, nkv, page, hd] when privately allocated, or the draft # KV cache manager's per-layer pool views when bound to it. self._ctx_kv_buf = None @@ -335,7 +356,7 @@ def set_draft_model(self, draft_model) -> None: super().set_draft_model(draft_model) self._validate_draft_attention_backend(draft_model) - def _check_ctx_arena_fits(self, capacity, num_slots, L, nkv, hd, dtype): + def _check_ctx_arena_fits(self, capacity, num_slots, L, nkv, hd, dtype, kv_factor=2): """Fail with the arithmetic before allocating the drafter context arena. The arena is dense in max_seq_len and lands *after* the KV cache manager @@ -345,12 +366,12 @@ def _check_ctx_arena_fits(self, capacity, num_slots, L, nkv, hd, dtype): neither the buffer nor the knob that controls it. """ itemsize = torch.tensor([], dtype=dtype).element_size() - if self._dflash_attention_backend in _PAGED_ATTENTION_BACKENDS: + if self._ctx_paged: page_size = self._ctx_page_size pages_per_slot = (capacity + page_size - 1) // page_size - arena = L * num_slots * pages_per_slot * 2 * nkv * page_size * hd * itemsize + arena = L * num_slots * pages_per_slot * kv_factor * nkv * page_size * hd * itemsize else: - arena = 2 * num_slots * L * capacity * nkv * hd * itemsize + arena = kv_factor * num_slots * L * capacity * nkv * hd * itemsize # mem_get_info() excludes the caching allocator's reserved-but-unused # segments, which torch.zeros() below can serve from. @@ -370,7 +391,7 @@ def _check_ctx_arena_fits(self, capacity, num_slots, L, nkv, hd, dtype): f"DFlash drafter context arena needs {arena / gib:.2f} GiB but only " f"{free / gib:.2f} GiB is free. It is dense in sequence length: " f"slots={num_slots} x layers={L} x capacity={capacity} x " - f"kv_heads={nkv} x head_dim={hd} x {itemsize}B, K and V. " + f"kv_heads={nkv} x head_dim={hd} x {itemsize}B x {kv_factor} halves. " f"capacity comes from max_seq_len (+{self._resolved_block_size} " f"block slack), so pass an explicit max_seq_len matching the " f"sequences you actually serve instead of letting it fall back to " @@ -378,7 +399,7 @@ def _check_ctx_arena_fits(self, capacity, num_slots, L, nkv, hd, dtype): f"scales it down linearly too." ) - def _managed_ctx_pool(self, draft_kv_cache_manager, L, nkv, hd, dtype): + def _managed_ctx_pool(self, draft_kv_cache_manager, L, nkv, hd, dtype, kv_factor=2): """Per-layer views of the draft KV cache manager's pool, or None. The manager already funds a pool sized from the drafter's own config @@ -397,13 +418,15 @@ def _managed_ctx_pool(self, draft_kv_cache_manager, L, nkv, hd, dtype): checked here and the index space is settled in _init_ctx_block_tables. """ if draft_kv_cache_manager is None: - # No separate draft KV cache: attention DP disables it - # (_util.py:_should_create_separate_draft_kv_cache), and the - # two-model paths never build one. + # No separate draft KV cache. Attention DP alone does NOT disable + # it: _util.py skips that bail for an external drafter, which is + # what DSpark is, so this drafter does get a manager under DP. The + # reachable cases are the two-model paths, which never build one. return None layers = [draft_kv_cache_manager.get_buffers(i, kv_layout="HND") for i in range(L)] base = layers[0] - expected = (2, nkv, base.size(-2), hd) + # kv_factor 1 is the MLA drafter's SELFKONLY pool: one latent, no V. + expected = (kv_factor, nkv, base.size(-2), hd) for i, layer in enumerate(layers): if tuple(layer.shape[1:]) != expected or layer.dtype != dtype: # Fall back rather than fail: the private arena is what every @@ -600,18 +623,44 @@ def _lazy_init_ctx_buffers( nh = draft_model._num_heads nkv = draft_model._num_kv_heads hd = draft_model._head_dim + kv_factor = getattr(draft_model, "_kv_factor", 2) capacity = self._max_ctx + self._compute_block_size - if self._dflash_attention_backend in _PAGED_ATTENTION_BACKENDS: + # The one number a drafter that builds an absolute-position table needs, + # published from the only place that knows it: _max_ctx comes from + # attn_metadata.max_seq_len, the engine's value, which + # py_executor_creator raises past model_config.max_seq_len and never + # writes back. Tables are built lazily on first forward and this runs + # before any of them, so a drafter may read it in place of its + # config-derived cap. _compute_block_size, not _resolved_block_size: + # the block decode's j runs over the slots the forward computes. + draft_model._runtime_position_ceiling = dflash_position_ceiling( + self._max_ctx, self._compute_block_size, self.max_draft_len + ) + # TRTLLM and FA4 need pages because their kernels are paged; a drafter + # can also ask for them on its own (the MLA one does, to get its + # footprint under free_gpu_memory_fraction) while still reading them + # eagerly. + use_paged = self._dflash_attention_backend in _PAGED_ATTENTION_BACKENDS or getattr( + draft_model, "_paged_ctx_cache", False + ) + self._ctx_paged = use_paged + if use_paged: + # FA4 pages, but on the private arena with its own kernel. pool = ( - self._managed_ctx_pool(draft_kv_cache_manager, L, nkv, hd, dtype) - if self._dflash_attention_backend == "TRTLLM" - else None + None + if self._dflash_attention_backend == "FA4" + else self._managed_ctx_pool(draft_kv_cache_manager, L, nkv, hd, dtype, kv_factor) ) # The manager's page size wins when bound to it: its pool is already # carved, so the drafter adopts the geometry rather than imposing one. page_size = self._ctx_page_size if pool is None else pool[0].size(-2) self._ctx_page_size = page_size - if self._dflash_attention_backend == "TRTLLM": + if self._dflash_attention_backend == "TRTLLM" and getattr( + draft_model, "_uses_worker_attention_backend", True + ): + # Shape limits of the trtllm-gen FMHA ops -- so they bind only + # on a drafter those ops actually run. One that brings its own + # block decode pages the same way but is not shaped like them. has_context_attention = any( not draft_model._get_attention_mask_args(layer_idx)[0] for layer_idx in range(L) ) @@ -623,7 +672,7 @@ def _lazy_init_ctx_buffers( tokens_per_block=page_size, has_context_attention=has_context_attention, ) - else: # FA4 stays on the private arena, with its own kernel. + elif self._dflash_attention_backend == "FA4": validate_dflash_fa4_runtime(dtype=dtype, head_dim=hd) # Settle the block table before committing to the pool: it is the # last thing that can rule the pool out, and falling back after @@ -637,11 +686,11 @@ def _lazy_init_ctx_buffers( # takes its pages from the manager one iteration at a time. self._ctx_pages_per_slot = (capacity + page_size - 1) // page_size total_pages = num_slots * self._ctx_pages_per_slot - self._check_ctx_arena_fits(capacity, num_slots, L, nkv, hd, dtype) + self._check_ctx_arena_fits(capacity, num_slots, L, nkv, hd, dtype, kv_factor) # HND layout consumed by both FlashInfer's paged append and the # TRTLLM-Gen launcher: [L, pages, K/V, Hkv, page, D]. self._ctx_kv_buf = torch.zeros( - (L, total_pages, 2, nkv, page_size, hd), + (L, total_pages, kv_factor, nkv, page_size, hd), dtype=dtype, device="cuda", ) @@ -663,11 +712,15 @@ def _lazy_init_ctx_buffers( self._ctx_kv_last_page_len = torch.full( (num_slots,), page_size, dtype=torch.int32, device="cuda" ) - else: # VANILLA DFlash backend (FlashAttention) - self._check_ctx_arena_fits(capacity, num_slots, L, nkv, hd, dtype) + else: # dense private arena (VANILLA, unpaged) + self._check_ctx_arena_fits(capacity, num_slots, L, nkv, hd, dtype, kv_factor) kv_shape = (num_slots, L, capacity, nkv, hd) self._ctx_k_buf = torch.zeros(kv_shape, dtype=dtype, device="cuda") - self._ctx_v_buf = torch.zeros(kv_shape, dtype=dtype, device="cuda") + # An MLA drafter stores one latent per token; leaving _ctx_v_buf as + # None is what tells the write paths below there is no V half. + self._ctx_v_buf = ( + torch.zeros(kv_shape, dtype=dtype, device="cuda") if kv_factor == 2 else None + ) self._ctx_buf_inited = True logger.info( @@ -733,6 +786,21 @@ def _store_context_kv_paged( rows: torch.Tensor, positions: torch.Tensor, ) -> None: + if v is None: + # Single-latent (MLA) cache: one vector per token per layer, so + # there is no K/V interleave for append_paged_kv_cache to do and the + # write is a plain scatter into the page the block table names. + table = ( + self._ctx_block_tables + if self._ctx_block_tables is not None + else self._ctx_page_table + ) + page_size = self._ctx_page_size + pages = table[rows.long(), positions.long() // page_size].long() + offsets = positions.long() % page_size + for layer_idx in range(k.size(1)): + self._ctx_kv_buf[layer_idx][pages, 0, 0, offsets] = k[:, layer_idx, 0] + return append_paged_kv_cache = self._get_ctx_paged_append() # Convert at most once and reuse for every layer. Calling ``.to()`` in @@ -911,9 +979,7 @@ def _store_prefill_context( self._req_ctx_pos[req_id] = first_pos + slen if actual > 0: cache_dtype = ( - self._ctx_kv_buf[0].dtype - if self._dflash_attention_backend in _PAGED_ATTENTION_BACKENDS - else self._ctx_k_buf.dtype + self._ctx_kv_buf[0].dtype if self._ctx_paged else self._ctx_k_buf.dtype ) chunk_proj_cast = chunk_proj[:actual].to(cache_dtype) ctx_len_updates[slot] = end @@ -923,7 +989,7 @@ def _store_prefill_context( chunk_proj_cast, chunk_pos[:actual] ) # chunk_k/v: [actual, L, nkv, hd] → [L, actual, nkv, hd] - if self._dflash_attention_backend in _PAGED_ATTENTION_BACKENDS: + if self._ctx_paged: # Manager block tables are keyed by batch position, the # private arena's by slot. See _ctx_paged_index_args. row = i if self._ctx_block_tables is not None else slot @@ -935,7 +1001,8 @@ def _store_prefill_context( ) else: # VANILLA DFlash backend (FlashAttention) self._ctx_k_buf[slot, :, cur:end] = chunk_k.permute(1, 0, 2, 3) - self._ctx_v_buf[slot, :, cur:end] = chunk_v.permute(1, 0, 2, 3) + if chunk_v is not None: + self._ctx_v_buf[slot, :, cur:end] = chunk_v.permute(1, 0, 2, 3) offset += slen self._write_ctx_len(ctx_len_updates) @@ -998,7 +1065,17 @@ def _forward_impl( ) spec_metadata._dflash_worker = self # Before any store: prefill and decode both address pages through it. - self._refresh_ctx_block_tables(attn_metadata, batch_size) + refreshed = self._refresh_ctx_block_tables(attn_metadata, batch_size) + if self._ctx_block_tables is not None and not refreshed: + # False is normal only while no paged pool is bound (the private + # arena addresses slots directly). Once one IS bound, the counts + # drive the advertised context length, so a missed refresh leaves + # them at zero and every drafter attends an empty context -- no + # error, just acceptance quietly collapsing. + raise RuntimeError( + "DFlash: draft block tables are bound to a paged pool but this " + "iteration's attn_metadata carried no draft_kv_cache_block_offsets." + ) # Save context lengths so both warmup and a failed forward can roll # back the in-place _ctx_len updates made during drafting. @@ -1420,11 +1497,13 @@ def prepare_1st_drafter_inputs( j_block = torch.arange(query_tokens_per_req, dtype=torch.long, device="cuda") offsets_kp1 = torch.arange(K_plus_1, dtype=torch.long, device="cuda") - query_position_ids = ( - ctx_len_gen.unsqueeze(1) - + gen_num_accepted.long().unsqueeze(1) - + j_block.unsqueeze(0) - ) + # _ctx_len is clamped to _max_ctx only AFTER this step's accepted + # tokens are folded in (see the update below), so the running length + # used here has to be clamped on its own -- otherwise a request that + # already sits at the ceiling indexes num_accepted positions past + # any position the sequence can legitimately reach. + ctx_len_now = (ctx_len_gen + gen_num_accepted.long()).clamp_(max=self._max_ctx) + query_position_ids = ctx_len_now.unsqueeze(1) + j_block.unsqueeze(0) ctx_position_ids = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0) # Go through embed_tokens.forward (NOT .weight[...]) so TP-sharded @@ -1469,19 +1548,18 @@ def prepare_1st_drafter_inputs( # Fast path: store the pre-projected/pre-RoPE'd K/V. # dflash_forward reads these directly via cache_batch_idx. cache_dtype = ( - self._ctx_kv_buf[0].dtype - if self._dflash_attention_backend in _PAGED_ATTENTION_BACKENDS - else self._ctx_k_buf.dtype + self._ctx_kv_buf[0].dtype if self._ctx_paged else self._ctx_k_buf.dtype ) k_new, v_new = draft_model.precompute_context_kv( proj_flat.to(cache_dtype), pos_flat ) mask_bc = mask_1d.view(-1, 1, 1, 1).to(k_new.dtype) k_new.mul_(mask_bc) - v_new.mul_(mask_bc) + if v_new is not None: + v_new.mul_(mask_bc) slot_long = slot_flat.long() col_long = col_flat.long() - if self._dflash_attention_backend in _PAGED_ATTENTION_BACKENDS: + if self._ctx_paged: if self._ctx_block_tables is not None: # Batch positions of the gen requests, matching the # per-request block table's row order. @@ -1491,12 +1569,34 @@ def prepare_1st_drafter_inputs( self._store_context_kv_paged(k_new, v_new, rows_long, col_long) else: # VANILLA DFlash backend (FlashAttention) self._ctx_k_buf[slot_long, :, col_long] = k_new - self._ctx_v_buf[slot_long, :, col_long] = v_new + if v_new is not None: + self._ctx_v_buf[slot_long, :, col_long] = v_new self._ctx_len[slots] += gen_num_accepted_long self._ctx_len.clamp_(max=self._max_ctx) num_ctx_per_req_t = self._ctx_len[slots] + if self._ctx_block_tables is not None: + # ctx_len tracks the target sequence, but the write above clamps + # columns to what the manager allocated for this request, so a + # context that outruns its allocation has its tail written on + # top of the last valid slot. Reading past the allocation would + # then attend either that clobbered value or -- for pages the + # block table left clamped to 0 -- another request's data. + # Truncating the advertised length keeps the two consistent: + # the drafter attends a short context instead of a wrong one. + # + # Leave the block room too. The MLA path writes its own latents + # at ctx_len..ctx_len+block_size (_build_mla_block_fixup), so + # stopping at `allocated` puts the first block position on the + # first UNallocated page -- whose block-table entry _refresh + # clamped from its negative placeholder to 0, i.e. another + # request's block. That is a silent cross-request write, not an + # out-of-range fault. + allocated = (self._ctx_block_counts * self._ctx_page_size - block_size).clamp_( + min=0 + ) + num_ctx_per_req_t = torch.minimum(num_ctx_per_req_t, allocated[gen_rows_out]) noise_embedding = noise_embed_2d query_positions = query_position_ids.long() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py index ff3af9a76b90..01f4f60f40ac 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py @@ -19,6 +19,7 @@ confidence_proj weights load without being used. """ +import math import re from types import SimpleNamespace @@ -34,6 +35,32 @@ ) from tensorrt_llm._torch.speculative.dflash import dflash_draft_slot_ids +needs_cuda = pytest.mark.skipif( + not torch.cuda.is_available(), reason="drafter module construction needs CUDA" +) + + +def _has_absorbed_mla_kernel() -> bool: + """SM100-family only: trtllm-gen has no absorbed-MLA decode below it. + + This directory is mapped to the CPU and H100 stages, so a plain + torch.cuda.is_available() gate sends the paged variant to sm90, where the + kernel does not exist -- it either raises inside flashinfer or falls back + to eager and trips the branch assertion. Either way the failure is about + the stage, not the code. + """ + if not torch.cuda.is_available(): + return False + from tensorrt_llm._utils import is_sm_100f + + return is_sm_100f() + + +needs_absorbed_mla = pytest.mark.skipif( + not _has_absorbed_mla_kernel(), + reason="paged absorbed-MLA decode needs an SM100-family GPU", +) + # --------------------------------------------------------------------------- # Reference oracle: line-for-line port of DeepSpec VanillaMarkov # (deepspec/modeling/dspark/markov_head.py) at temperature 0. @@ -442,6 +469,7 @@ def test_published_drafter_spelling_activates_the_heads(): @needs_gpu +@needs_cuda def test_head_weights_without_a_resolvable_rank_raise(): """The inverse of the missing-weights check. @@ -638,3 +666,693 @@ def test_plain_dflash_block_decode_matches_full_attention_oracle(): ) diff = (out - oracle).abs().max().item() assert diff < 0.02, f"plain DFlash parity failed: max abs diff {diff}" + + +# --------------------------------------------------------------------------- +# MLA-backboned DSpark drafter (Inferact/Kimi-K3-DSpark shape). +# +# The GQA block decode cannot express it: there is no per-head K/V to fuse, and +# the cache holds one MLA latent per token instead of a K and a V half. What +# can break silently here is the absorption and the YaRN rope the drafter was +# distilled under, so both are checked against references written in the +# un-absorbed / HF form rather than against the same math again. +# --------------------------------------------------------------------------- + +MLA_HIDDEN = 64 +MLA_INTERMEDIATE = 32 +MLA_HEADS = 4 +MLA_Q_LORA = 16 +# Real MLA head geometry, even though everything else here is tiny. Shrinking +# these silently disables both kernel paths: flashinfer's MLA templates are +# specialised on (kv_lora, rope) = (512, 64) and fail to compile at 16/4, and +# the fused RMSNorm/RoPE op requires rope_dim/2 % 32 == 0. With the old toy +# dims every "kernel" variant fell back to eager and the test proved nothing. +MLA_KV_LORA = 512 +MLA_NOPE = 128 +MLA_ROPE = 64 +MLA_V_DIM = 128 +MLA_LAYERS = 2 +MLA_BLOCK = 3 +MLA_LATENT = MLA_KV_LORA + MLA_ROPE +MLA_THETA = 50000.0 +MLA_FACTOR = 4.0 +MLA_ORIG_MAX = 64 + + +def _tiny_mla_config(): + """Tiny MLA drafter config in the published (TorchSpec) spelling. + + Mirrors Inferact/Kimi-K3-DSpark: an architecture label no registry knows, + ``model_type`` "k3_dspark", head switches and ``target_layer_ids`` at the + top level, and YaRN under the transformers-v5 ``rope_parameters`` key. + """ + from transformers import PretrainedConfig + + cfg = PretrainedConfig( + architectures=["K3DSparkModel"], + model_type="k3_dspark", + hidden_size=MLA_HIDDEN, + intermediate_size=MLA_INTERMEDIATE, + num_hidden_layers=MLA_LAYERS, + num_attention_heads=MLA_HEADS, + num_key_value_heads=MLA_HEADS, + q_lora_rank=MLA_Q_LORA, + kv_lora_rank=MLA_KV_LORA, + qk_nope_head_dim=MLA_NOPE, + qk_rope_head_dim=MLA_ROPE, + v_head_dim=MLA_V_DIM, + vocab_size=VOCAB, + rms_norm_eps=1e-6, + max_position_embeddings=256, + rope_theta=MLA_THETA, + block_size=MLA_BLOCK, + mask_token_id=VOCAB - 3, + target_layer_ids=[0, 1], + markov_rank=RANK, + markov_head_type="vanilla", + enable_confidence_head=True, + confidence_head_with_markov=True, + tie_word_embeddings=False, + rope_parameters={ + "rope_type": "yarn", + "factor": MLA_FACTOR, + "original_max_position_embeddings": MLA_ORIG_MAX, + "rope_theta": MLA_THETA, + "beta_fast": 32, + "beta_slow": 1, + "mscale": 1.0, + "mscale_all_dim": 1.0, + }, + ) + cfg.dflash_config = {"mask_token_id": VOCAB - 3, "target_layer_ids": [0, 1]} + cfg.torch_dtype = torch.bfloat16 + return cfg + + +def _tiny_mla_weights(seed=11): + g = torch.Generator().manual_seed(seed) + + def rnd(*shape): + return (torch.randn(*shape, generator=g) * 0.05).to(torch.bfloat16) + + qk_head_dim = MLA_NOPE + MLA_ROPE + w = { + # TorchSpec spellings, deliberately not SpecForge's fc/hidden_norm/norm. + "context_proj.weight": rnd(MLA_HIDDEN, MLA_HIDDEN * NUM_CAPTURE), + "context_norm.weight": rnd(MLA_HIDDEN).abs() + 1.0, + "final_norm.weight": rnd(MLA_HIDDEN).abs() + 1.0, + "embed_tokens.weight": rnd(VOCAB, MLA_HIDDEN), + "markov_head.markov_w1.weight": rnd(VOCAB, RANK), + "markov_head.markov_w2.weight": rnd(VOCAB, RANK), + "confidence_head.proj.weight": rnd(1, MLA_HIDDEN + RANK), + "confidence_head.proj.bias": rnd(1), + } + for i in range(MLA_LAYERS): + p = f"layers.{i}." + w[p + "input_layernorm.weight"] = rnd(MLA_HIDDEN).abs() + 1.0 + w[p + "post_attention_layernorm.weight"] = rnd(MLA_HIDDEN).abs() + 1.0 + w[p + "self_attn.q_a_proj.weight"] = rnd(MLA_Q_LORA, MLA_HIDDEN) + w[p + "self_attn.q_a_layernorm.weight"] = rnd(MLA_Q_LORA).abs() + 1.0 + w[p + "self_attn.q_b_proj.weight"] = rnd(MLA_HEADS * qk_head_dim, MLA_Q_LORA) + w[p + "self_attn.kv_a_proj_with_mqa.weight"] = rnd(MLA_LATENT, MLA_HIDDEN) + w[p + "self_attn.kv_a_layernorm.weight"] = rnd(MLA_KV_LORA).abs() + 1.0 + w[p + "self_attn.kv_b_proj.weight"] = rnd(MLA_HEADS * (MLA_NOPE + MLA_V_DIM), MLA_KV_LORA) + w[p + "self_attn.o_proj.weight"] = rnd(MLA_HIDDEN, MLA_HEADS * MLA_V_DIM) + w[p + "mlp.gate_proj.weight"] = rnd(MLA_INTERMEDIATE, MLA_HIDDEN) + w[p + "mlp.up_proj.weight"] = rnd(MLA_INTERMEDIATE, MLA_HIDDEN) + w[p + "mlp.down_proj.weight"] = rnd(MLA_HIDDEN, MLA_INTERMEDIATE) + return w + + +def _build_mla_drafter(weights, *, dflash_attention_backend="VANILLA"): + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models.modeling_dspark import MLADSparkForCausalLM + + model_config = ModelConfig(pretrained_config=_tiny_mla_config(), attn_backend="TRTLLM") + drafter = MLADSparkForCausalLM( + model_config, dflash_attention_backend=dflash_attention_backend + ).to("cuda") + drafter.load_weights(dict(weights)) + return drafter + + +def _hf_yarn_reference(dim, base, factor, original_max, max_pos, beta_fast=32.0, beta_slow=1.0): + """Line-for-line port of HF ``DeepseekV3YarnRotaryEmbedding`` (mscale 1).""" + + def find_dim(num_rotations): + return (dim * math.log(original_max / (num_rotations * 2 * math.pi))) / (2 * math.log(base)) + + low = max(math.floor(find_dim(beta_fast)), 0) + high = min(math.ceil(find_dim(beta_slow)), dim - 1) + if low == high: + high += 0.001 + freq_extra = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + freq_inter = 1.0 / (factor * base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + ramp = torch.clamp((torch.arange(dim // 2, dtype=torch.float32) - low) / (high - low), 0, 1) + inv_freq = freq_inter * ramp + freq_extra * (1 - ramp) + freqs = torch.outer(torch.arange(max_pos, dtype=torch.float32), inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + return emb.cos(), emb.sin() + + +def _hf_apply_rope(x, cos, sin): + """HF ``DeepseekV3`` rope: de-interleave, then the half-split rotation.""" + d = x.shape[-1] + x = x.reshape(*x.shape[:-1], d // 2, 2).transpose(-1, -2).reshape(*x.shape[:-1], d) + x1, x2 = x[..., : d // 2], x[..., d // 2 :] + return x * cos + torch.cat((-x2, x1), dim=-1) * sin + + +def test_mla_dspark_yarn_rope_matches_hf_reference(): + """The drafter's YaRN table and rotation are HF DeepSeek's. + + Both published MLA drafters are distilled under HF/vLLM numerics. Routing + through ``RopeParams`` instead picks up the fused MLA kernel's + ``duplicate_data`` / GPT-J packing convention, which rotates differently and + costs acceptance with nothing raised. + """ + from tensorrt_llm._torch.models.modeling_dspark import ( + apply_dspark_mla_rope, + build_dspark_mla_yarn_rope, + ) + + dim, max_pos = 8, 32 + cos, sin = build_dspark_mla_yarn_rope( + dim=dim, + base=MLA_THETA, + scaling_factor=MLA_FACTOR, + original_max_position_embeddings=MLA_ORIG_MAX, + max_position_embeddings=max_pos, + beta_fast=32.0, + beta_slow=1.0, + mscale=1.0, + mscale_all_dim=1.0, + device="cpu", + ) + ref_cos, ref_sin = _hf_yarn_reference(dim, MLA_THETA, MLA_FACTOR, MLA_ORIG_MAX, max_pos) + torch.testing.assert_close(cos, ref_cos) + torch.testing.assert_close(sin, ref_sin) + + g = torch.Generator().manual_seed(3) + x = torch.randn(5, 3, dim, generator=g) + pos = torch.tensor([0, 7, 31]) + torch.testing.assert_close( + apply_dspark_mla_rope(x, cos[pos], sin[pos]), + _hf_apply_rope(x, ref_cos[pos], ref_sin[pos]), + ) + + +def _oracle_mla_block_decode(weights, captured, noise_embed, *, swap_kv_b=False): + """fp32 eager MLA block decode, written in the UN-ABSORBED form. + + The drafter attends the stored latent directly (``q_nope`` folded through + ``kv_b_proj``'s K half); this expands the latent into per-head K/V and runs + plain attention instead, so agreement is evidence about the absorption + rather than a restatement of it. ``swap_kv_b`` is the negative control: + exchanging the K and V halves must break parity. + """ + from tensorrt_llm._torch.models.modeling_dspark import apply_dspark_mla_rope + + w = {k: v.float() for k, v in weights.items()} + ctx, blk = captured.shape[0], noise_embed.shape[0] + ctx_pos = torch.arange(ctx, dtype=torch.long) + q_pos = torch.arange(ctx, ctx + blk, dtype=torch.long) + cos, sin = _hf_yarn_reference( + MLA_ROPE, MLA_THETA, MLA_FACTOR, MLA_ORIG_MAX, MLA_ORIG_MAX * int(MLA_FACTOR) + ) + scale = _yarn_mscale_sq() / (MLA_NOPE + MLA_ROPE) ** 0.5 + + # Context features: context_norm(context_proj(captured)); constant across + # layers and NOT passed through input_layernorm (the DFlash contract). + ctx_feat = _rms(captured.float() @ w["context_proj.weight"].T, w["context_norm.weight"]) + + hs = noise_embed.float() + for i in range(MLA_LAYERS): + p = f"layers.{i}." + h = _rms(hs, w[p + "input_layernorm.weight"]) + + q = _rms(h @ w[p + "self_attn.q_a_proj.weight"].T, w[p + "self_attn.q_a_layernorm.weight"]) + q = (q @ w[p + "self_attn.q_b_proj.weight"].T).view(blk, MLA_HEADS, MLA_NOPE + MLA_ROPE) + q_nope, q_rope = q.split([MLA_NOPE, MLA_ROPE], dim=-1) + q_rope = apply_dspark_mla_rope(q_rope, cos[q_pos].unsqueeze(1), sin[q_pos].unsqueeze(1)) + + def latent_of(x, pos): + lat = x @ w[p + "self_attn.kv_a_proj_with_mqa.weight"].T + return torch.cat( + [ + _rms(lat[..., :MLA_KV_LORA], w[p + "self_attn.kv_a_layernorm.weight"]), + apply_dspark_mla_rope(lat[..., MLA_KV_LORA:], cos[pos], sin[pos]), + ], + dim=-1, + ) + + kv = torch.cat([latent_of(ctx_feat, ctx_pos), latent_of(h, q_pos)], dim=0) + kv_b = w[p + "self_attn.kv_b_proj.weight"].view( + MLA_HEADS, MLA_NOPE + MLA_V_DIM, MLA_KV_LORA + ) + k_b, v_b = kv_b[:, :MLA_NOPE], kv_b[:, MLA_NOPE:] + if swap_kv_b: + k_b, v_b = v_b, k_b + c_kv, k_pe = kv[:, :MLA_KV_LORA], kv[:, MLA_KV_LORA:] + k = torch.cat( + [ + torch.einsum("tc,hdc->thd", c_kv, k_b), + k_pe.unsqueeze(1).expand(-1, MLA_HEADS, -1), + ], + dim=-1, + ) + v = torch.einsum("tc,hvc->thv", c_kv, v_b) + + scores = torch.einsum("qhd,khd->hqk", torch.cat([q_nope, q_rope], -1), k) * scale + attn = torch.softmax(scores, dim=-1) + o = torch.einsum("hqk,khv->qhv", attn, v).reshape(blk, MLA_HEADS * MLA_V_DIM) + + hs = hs + o @ w[p + "self_attn.o_proj.weight"].T + h2 = _rms(hs, w[p + "post_attention_layernorm.weight"]) + gate = h2 @ w[p + "mlp.gate_proj.weight"].T + up = h2 @ w[p + "mlp.up_proj.weight"].T + hs = hs + (F.silu(gate) * up) @ w[p + "mlp.down_proj.weight"].T + return _rms(hs, w["final_norm.weight"]) + + +def _yarn_mscale_sq(): + m = 0.1 * 1.0 * math.log(MLA_FACTOR) + 1.0 + return m * m + + +# trtllm-gen accepts only 32 or 64 (production uses 64); the wrapper path is +# looser, so a smaller value silently tests one variant and not the other. +MLA_PAGE = 32 + + +def _run_mla_block_decode(drafter, captured, noise_embed, paged=False): + """Drive one block decode, optionally through the paged cache. + + ``paged=False`` keeps the dense arena the eager reference path uses. + ``paged=True`` hands ``dflash_forward`` a page table, which is what selects + the kernel paths -- without it ``use_kernel`` is False and every kernel + variant silently falls back to the same eager branch. + """ + dev = "cuda" + proj = drafter.project_target_hidden(captured.to(dev, torch.bfloat16)) + ctx_pos = torch.arange(CTX_LEN, device=dev) + latent, v = drafter.precompute_context_kv(proj, ctx_pos) + assert v is None, "an MLA drafter stores no V half" + assert latent.shape == (CTX_LEN, MLA_LAYERS, 1, MLA_LATENT) + + kwargs = dict( + noise_embedding=noise_embed.to(dev, torch.bfloat16).unsqueeze(0), + query_positions=torch.arange(CTX_LEN, CTX_LEN + MLA_BLOCK, device=dev).unsqueeze(0), + num_ctx_per_req=torch.tensor([CTX_LEN], device=dev), + ctx_v_cache=None, + ctx_cache_batch_idx=torch.tensor([0], device=dev), + ) + if not paged: + pool = torch.zeros( + 1, MLA_LAYERS, CTX_LEN + MLA_BLOCK, 1, MLA_LATENT, dtype=torch.bfloat16, device=dev + ) + pool[0, :, :CTX_LEN] = latent.permute(1, 0, 2, 3) + out = drafter.dflash_forward(ctx_k_cache=pool, **kwargs) + return out.float().cpu() + + # [pages, kv_factor=1, nkv=1, page_size, head_dim] per layer, plus a block's + # worth of slack: the manager reserves it, and the non-causal path writes + # the draft block into it. + npages = -(-(CTX_LEN + MLA_BLOCK) // MLA_PAGE) + pool = [ + torch.zeros(npages, 1, 1, MLA_PAGE, MLA_LATENT, dtype=torch.bfloat16, device=dev) + for _ in range(MLA_LAYERS) + ] + rows, cols = ctx_pos // MLA_PAGE, ctx_pos % MLA_PAGE + for layer_idx in range(MLA_LAYERS): + pool[layer_idx][rows, 0, 0, cols] = latent[:, layer_idx, 0] + page_table = torch.arange(npages, device=dev, dtype=torch.int32).unsqueeze(0) + out = drafter.dflash_forward( + ctx_k_cache=pool[0], ctx_kv_cache=pool, ctx_page_table=page_table, **kwargs + ) + return out.float().cpu() + + +@needs_cuda +@pytest.mark.parametrize( + "paged", + [False, pytest.param(True, marks=needs_absorbed_mla)], + ids=["eager", "kernel_fixup"], +) +def test_mla_dspark_block_decode_matches_unabsorbed_oracle(monkeypatch, paged): + """Absorbed block decode == un-absorbed eager MLA, and the halves are not + interchangeable (the negative control keeps this from being a tautology). + + The paged variant writes the draft block into the pool and reads it back + through the kernel, so a bug in that write shows up here as a parity + failure rather than as lower acceptance length. + """ + import tensorrt_llm._torch.models.modeling_dspark as md + + # Count the builder instead of trusting the parametrisation: the paged + # variant has to enter the kernel branch, and a fixture that quietly falls + # back to eager is exactly how this test used to pass without testing + # anything (no page table, then MLA dims and a page size both kernels + # reject). + calls = {"fixup": 0} + _original = md._build_mla_block_fixup + + def _counted(*a, **kw): + calls["fixup"] += 1 + return _original(*a, **kw) + + monkeypatch.setattr(md, "_build_mla_block_fixup", _counted) + torch.manual_seed(0) + weights = _tiny_mla_weights() + drafter = _build_mla_drafter(weights) + + g = torch.Generator().manual_seed(42) + captured = torch.randn(CTX_LEN, MLA_HIDDEN * NUM_CAPTURE, generator=g) * 0.5 + noise_embed = torch.randn(MLA_BLOCK, MLA_HIDDEN, generator=g) * 0.5 + + out = _run_mla_block_decode(drafter, captured, noise_embed, paged=paged) + captured_q, noise_q = captured.to(torch.bfloat16), noise_embed.to(torch.bfloat16) + oracle = _oracle_mla_block_decode(weights, captured_q, noise_q) + swapped = _oracle_mla_block_decode(weights, captured_q, noise_q, swap_kv_b=True) + + assert calls == {"fixup": 1 if paged else 0}, f"paged={paged} took the wrong branch: {calls}" + + diff = (out - oracle).abs().max().item() + diff_swapped = (out - swapped).abs().max().item() + assert diff < 0.02, f"MLA parity failed: max abs diff {diff}" + assert diff_swapped > 4 * max(diff, 1e-4), ( + f"negative control failed: kv_b K/V halves swapped is too close " + f"({diff_swapped} vs {diff}) -- the absorption may not be exercised" + ) + + +@needs_cuda +def test_mla_dspark_loads_the_torchspec_spelling_and_keeps_its_embedding(): + """Loads unconverted, and does not adopt the target's embedding. + + Two silent failures guarded here: the TorchSpec capture projection is + ``context_proj`` / ``context_norm`` / ``final_norm`` where the DFlash base + expects ``fc`` / ``hidden_norm`` / ``norm``, and the drafter ships a trained + ``embed_tokens`` whose values differ from the target's, so taking the + target's would feed the block decode embeddings it was not distilled on. + """ + weights = _tiny_mla_weights() + drafter = _build_mla_drafter(weights) + + assert drafter._dspark_shift_label + assert drafter.has_markov_head + assert drafter._kv_factor == 1, "the MLA cache stores one latent, not K and V" + assert (drafter._num_kv_heads, drafter._head_dim) == (1, MLA_LATENT) + torch.testing.assert_close(drafter.markov_w1.cpu(), weights["markov_head.markov_w1.weight"]) + torch.testing.assert_close( + drafter.fc.weight.cpu().float(), weights["context_proj.weight"].float() + ) + + target = SimpleNamespace( + model=SimpleNamespace(embed_tokens=torch.nn.Embedding(VOCAB, MLA_HIDDEN)), + lm_head=torch.nn.Linear(MLA_HIDDEN, VOCAB, bias=False), + ) + drafter.load_weights_from_target_model(target) + assert drafter.lm_head is target.lm_head + torch.testing.assert_close( + drafter.model.embed_tokens.weight.cpu().float(), weights["embed_tokens.weight"].float() + ) + + +@needs_cuda +@pytest.mark.parametrize("backend", ["VANILLA", "TRTLLM"]) +def test_mla_dspark_ignores_the_worker_attention_backend(backend): + """The MLA drafter runs _mla_paged_attention, so the field selects nothing. + + Both values must construct, and neither worker op set may be loaded -- + otherwise a drafter that never calls them drags in an optional dependency, + and the worker's per-backend shape checks (which the absorbed 64:1 / + head_dim 576 shape fails) would bind on a path that does not use them. + """ + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models.modeling_dspark import MLADSparkForCausalLM + + model_config = ModelConfig(pretrained_config=_tiny_mla_config(), attn_backend=backend) + drafter = MLADSparkForCausalLM(model_config, dflash_attention_backend=backend) + assert drafter._uses_worker_attention_backend is False + assert drafter.dflash_attention_backend == backend + assert drafter._dflash_flash_attention is None + assert drafter._dflash_trtllm_gen_ops is None + + +@needs_cuda +def test_mla_dspark_still_rejects_an_unknown_backend(): + """Dropping the VANILLA-only guard must not drop the typo check.""" + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models.modeling_dspark import MLADSparkForCausalLM + + model_config = ModelConfig(pretrained_config=_tiny_mla_config(), attn_backend="VANILLA") + with pytest.raises(ValueError, match="must be VANILLA or TRTLLM"): + MLADSparkForCausalLM(model_config, dflash_attention_backend="FLASHINFER") + + +def test_mla_dspark_rope_conventions_agree_on_scores(): + """The drafter's adjacent-pair RoPE must be the HF one up to a lane swap. + + The block decode moved off the HF layout onto the layout the fused + RMSNorm/RoPE kernel consumes. That is only sound because the permutation + cancels inside q_rope . k_rope -- and only if EVERY producer switches. The + negative control is the half-migrated state, which is silent in accuracy and + shows up solely as lower acceptance length. + """ + from tensorrt_llm._torch.models.modeling_dspark import ( + apply_dspark_rotary_batched, + build_dspark_mla_yarn_freqs_cis, + build_dspark_mla_yarn_rope, + ) + + torch.manual_seed(0) + dim, batch, blk, heads = 64, 3, 8, 4 + params = dict( + dim=dim, + base=50000.0, + scaling_factor=32.0, + original_max_position_embeddings=32768, + max_position_embeddings=256, + beta_fast=32.0, + beta_slow=1.0, + mscale=1.0, + mscale_all_dim=1.0, + device="cpu", + ) + cos, sin = build_dspark_mla_yarn_rope(**params) + freqs = build_dspark_mla_yarn_freqs_cis(**params) + + pos = torch.arange(blk) + q = torch.randn(batch, blk, heads, dim) + k = torch.randn(batch, blk, dim) + fc = freqs[pos].unsqueeze(0).expand(batch, blk, dim // 2) + + hf_q = _hf_apply_rope(q, cos[pos].view(1, blk, 1, dim), sin[pos].view(1, blk, 1, dim)) + hf_k = _hf_apply_rope(k, cos[pos].view(1, blk, dim), sin[pos].view(1, blk, dim)) + ap_q = apply_dspark_rotary_batched(q, fc) + ap_k = apply_dspark_rotary_batched(k, fc) + + perm = torch.empty(dim, dtype=torch.long) + perm[: dim // 2] = torch.arange(0, dim, 2) + perm[dim // 2 :] = torch.arange(1, dim, 2) + torch.testing.assert_close(hf_q, ap_q[..., perm], atol=1e-5, rtol=1e-5) + + hf_score = torch.einsum("bshd,bsd->bsh", hf_q, hf_k) + ap_score = torch.einsum("bshd,bsd->bsh", ap_q, ap_k) + torch.testing.assert_close(hf_score, ap_score, atol=1e-4, rtol=1e-4) + + # Negative control: only q migrated. + mixed = torch.einsum("bshd,bsd->bsh", ap_q, hf_k) + assert (mixed - hf_score).abs().max() > 1e-2, ( + "mixing the two RoPE conventions must change the scores; if it does not, " + "this test cannot catch a half-finished migration" + ) + + +def test_mla_block_fixup_stays_inside_the_allocation(): + """A context that fills its allocation must still leave the block room. + + dflash.py truncates ctx_len to what the draft KV manager allocated, but the + MLA path then writes the block's own latents at ctx_len..ctx_len+block_size. + Stopping at exactly `allocated` puts the first block position on the first + unallocated page, whose block-table entry _refresh_ctx_block_tables clamped + from a negative placeholder to 0 -- another request's block. Silent + cross-request corruption, not a fault, so only an explicit bound catches it. + """ + from tensorrt_llm._torch.models.modeling_dspark import _build_mla_block_fixup + + page_size, block_size, allocated_pages = 8, 3, 2 + allocated = allocated_pages * page_size + # Entries past the allocation are the clamped placeholders: physical 0. + page_tables = torch.tensor([[41, 42, 0, 0]]) + ctx_len = torch.tensor([max(allocated - block_size, 0)]) + + fixup = _build_mla_block_fixup(ctx_len, page_tables, block_size, page_size) + + allocated_ids = page_tables[0, :allocated_pages].tolist() + assert set(fixup.pages.flatten().tolist()) <= set(allocated_ids) + assert int(fixup.seq_lens_i32[0]) <= allocated + + +def test_mla_rope_table_follows_runtime_max_seq_len(): + """The position table is sized by what is served, not what is advertised. + + K3 declares max_position_embeddings = 1,048,576; at complex64 that table is + ~256 MiB per rank, built through full-size fp32 cos/sin. The served + max_seq_len is what the context cache is bounded by, so the table follows + it. Only the length moves -- YaRN's correction range is computed from + original_max_position_embeddings. + """ + from tensorrt_llm._torch.models.modeling_dspark import _resolve_dspark_mla_rope_params + + cfg = _tiny_mla_config() + cfg.max_position_embeddings = 1 << 20 + + advertised = _resolve_dspark_mla_rope_params(cfg) + served = _resolve_dspark_mla_rope_params(cfg, SimpleNamespace(max_seq_len=4096)) + + assert advertised["max_positions"] == 1 << 20 + assert served["max_positions"] == 4096 + # the YaRN inputs that decide the rotation must not move with it + for key in ("theta", "scaling_factor", "original_max_positions", "beta_fast", "beta_slow"): + assert advertised[key] == served[key] + + +@pytest.mark.parametrize( + "max_ctx,checkpoint_block_size,max_draft_len", + [ + (8205, None, 7), # the shipped recipe: no block_size, fallback K+1 = 8 + (8205, 7, 7), # a checkpoint declaring a window NARROWER than the verify group + (4104, 16, 3), # ... and one declaring a WIDER window + (1026, 2, 1), + ], +) +def test_dflash_position_ceiling_covers_both_index_paths( + max_ctx, checkpoint_block_size, max_draft_len +): + """The published ceiling must exceed every index the forward can produce. + + Two independent position sets read the same table, and they do not have the + same reach: the block decode runs to ctx_len + block_size - 1 while the + context path runs to ctx_len + max_draft_len. Sizing for either one alone is + short whenever the other is wider, which is why the ceiling takes a max. + + max_ctx here is the RUNTIME value (attn_metadata.max_seq_len, what ctx_len + is clamped to), not model_config.max_seq_len -- 8205 is the measured value + for a configured 8192 with max_draft_len 7. + """ + from tensorrt_llm._torch.speculative.dflash import dflash_position_ceiling + + block_size = checkpoint_block_size or (max_draft_len + 1) + ceiling = dflash_position_ceiling(max_ctx, block_size, max_draft_len) + + # dflash.py: query_position_ids = clamp(ctx_len + num_accepted) + j + worst_block = max_ctx + block_size - 1 + # dflash.py: ctx_position_ids = ctx_len + [0, max_draft_len] + worst_ctx = max_ctx + max_draft_len + + assert ceiling > worst_block, ( + f"ceiling {ceiling} does not cover block decode index {worst_block} " + f"(max_ctx {max_ctx}, block_size {block_size})" + ) + assert ceiling > worst_ctx, ( + f"ceiling {ceiling} does not cover context index {worst_ctx} " + f"(max_ctx {max_ctx}, max_draft_len {max_draft_len})" + ) + # and no more than one position of slack, so it cannot quietly grow + assert ceiling == max(worst_block, worst_ctx) + 1 + + +@needs_cuda +def test_mla_rope_table_prefers_the_runtime_ceiling_over_the_config_cap(): + """A worker-published ceiling replaces the config-derived cap outright. + + Not max() of the two: with max_seq_len unset the config cap is the + checkpoint's advertised max_position_embeddings (1,048,576 for K3), and + taking the larger would rebuild exactly the ~256 MiB table the cap exists to + avoid. Exercises the real _mla_freqs_cis, so it also pins that the table is + still lazy enough to see an attribute set after construction. + """ + drafter = _build_mla_drafter(_tiny_mla_weights()) + # construction-time cap: the plain config value, no lookahead baked in + assert drafter._mla_rope_params["max_positions"] == _tiny_mla_config().max_position_embeddings + + drafter._runtime_position_ceiling = 8213 + table = drafter._mla_freqs_cis(torch.device("cuda")) + assert table.shape[0] == 8213, ( + f"table has {table.shape[0]} positions, expected the published ceiling 8213" + ) + + +@needs_cuda +def test_mla_drafter_rejects_a_checkpoint_missing_backbone_weights(): + """Only lm_head may be absent; a truncated backbone must not load quietly. + + The generic loader skips a module whose whole subtree filters to nothing, + so allow_partial_loading=False would not catch this either -- the check has + to come from the drafter's own module tree. + """ + weights = _tiny_mla_weights() + drafter = _build_mla_drafter(weights) + + # A whole module, not one tensor: the check is module-granular on purpose, + # since a fused module (gate_up_proj) is stored unfused and dropping half of + # it is a parameter-level gap the loader's own naming cannot distinguish. + # Derive the prefix from the fixture rather than spelling it: DFlash + # checkpoints name layers without the `model.` prefix, so a hardcoded key + # silently matches nothing and the test passes by not truncating anything. + victim = next(k for k in weights if k.endswith("self_attn.o_proj.weight")) + truncated = {k: v for k, v in weights.items() if k != victim} + assert len(truncated) == len(weights) - 1 + + with pytest.raises(ValueError, match="WEIGHTS_SHARED_WITH_TARGET"): + drafter.load_weights(truncated) + + +@needs_cuda +@pytest.mark.parametrize("victim", ["context_proj.weight", "context_norm.weight"]) +def test_mla_drafter_requires_the_wrapper_owned_tensors(victim): + """fc / hidden_norm are built FROM the checkpoint, so no module walk sees them. + + They live on the wrapper, not on draft_model_full, and the extraction is + `if in remapped`. Without an explicit requirement a checkpoint that + omits them loads clean and the drafter runs with no capture projection: + has_target_features stays False, _ctx_len never advances, and it drafts + from an empty context for the life of the process. + """ + weights = _tiny_mla_weights() + drafter = _build_mla_drafter(weights) + + assert victim in weights, "fixture no longer ships the TorchSpec spelling" + truncated = {k: v for k, v in weights.items() if k != victim} + + with pytest.raises(ValueError, match="wrapper owns"): + drafter.load_weights(truncated) + + +@needs_cuda +def test_mla_drafter_rejects_a_partial_fused_component_set(): + """One of q/k/v present is not the module being present. + + The fused load path accepts a subset under allow_partial_loading and + leaves the absent shards at torch.empty -- uninitialised device memory, + not zeros -- so the check has to require every component, not any. + """ + weights = _tiny_mla_weights() + drafter = _build_mla_drafter(weights) + + # gate_proj/up_proj, NOT an unfused module: the checkpoint stores them + # separately and the module tree names them once as mlp.gate_up_proj, so + # dropping one is the only way to reach the fused branch of _supplied. + # (kv_b_proj would exercise the plain _has path and duplicate the + # missing-backbone test above.) + victim = next(k for k in weights if k.endswith("mlp.gate_proj.weight")) + assert any(k.endswith("mlp.up_proj.weight") for k in weights), ( + "the surviving half of the fused pair must be present, or this is not " + "testing the fused branch" + ) + truncated = {k: v for k, v in weights.items() if k != victim} + + with pytest.raises(ValueError, match="WEIGHTS_SHARED_WITH_TARGET"): + drafter.load_weights(truncated) diff --git a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py index 4a805374bcfc..7fa00613c361 100644 --- a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py +++ b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py @@ -201,3 +201,48 @@ def test_fused_dspark_rmsnorm_rope_compiles_once_across_batches(): cache_info = _compile_fused_dspark_rmsnorm_rope.cache_info() assert cache_info.misses == 1 assert cache_info.hits == 1 + + +@pytest.mark.parametrize("split_norm", [False, True]) +def test_fused_dspark_rmsnorm_rope_norm_dim(split_norm): + """norm_dim bounds the RMS reduction; the default must not move DSv4. + + split_norm=False is the DSv4 regression gate (whole-row RMS, weight over the + whole row). split_norm=True is what DeepSeek-style MLA needs: normalize the + latent, leave k_pe raw. + + Driven through ``_rmsnorm_rope_batched`` rather than the custom op, because + the plumbing under test is the dispatcher's -- it has to forward norm_dim to + the support predicate and to the kernel. The predicate is asserted first so + a fused path that silently stopped applying never reads as a pass; that is + what made the original home of this test (hw_agnostic, mapped to CPU and + H100 only) vacuous, since is_sm_100f() is false there. + """ + from tensorrt_llm._torch.custom_ops.dspark_rmsnorm_rope_custom_op import ( + is_fused_dspark_rmsnorm_rope_supported, + ) + from tensorrt_llm._torch.models.modeling_dspark import _rmsnorm, _rmsnorm_rope_batched + + torch.manual_seed(0) + hidden, rope_dim, rows = 576, 64, 16 + nope = hidden - rope_dim + norm_dim = nope if split_norm else hidden + eps = 1e-5 + x = torch.randn(1, rows, hidden, dtype=torch.bfloat16, device="cuda") + w = torch.randn(norm_dim, dtype=torch.bfloat16, device="cuda").abs() + 0.5 + ang = torch.rand(1, rows, rope_dim // 2, device="cuda", dtype=torch.float32) * 6.28 + freqs = torch.complex(ang.cos(), ang.sin()) + + freqs_real = torch.view_as_real(freqs).reshape(-1, freqs.shape[-1], 2) + assert is_fused_dspark_rmsnorm_rope_supported(x, w, freqs_real, 1, rope_dim, norm_dim) + + got = _rmsnorm_rope_batched(x, w, eps, rope_dim, freqs, norm_dim=norm_dim) + + head = x[..., :norm_dim] + normed = _rmsnorm(head, w, eps) + ref = torch.cat([normed, x[..., norm_dim:]], dim=-1) if split_norm else normed + xc = torch.view_as_complex(ref[..., nope:].float().unflatten(-1, (-1, 2))) + rot = torch.view_as_real(xc * freqs.view(1, rows, 1, rope_dim // 2).squeeze(2)).flatten(-2) + ref = torch.cat([ref[..., :nope], rot.to(ref.dtype)], dim=-1) + + torch.testing.assert_close(got.float(), ref.float(), atol=2e-2, rtol=2e-2) From 9918cf73e2861e5e2e5b46912aa79de747f3c05f Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Thu, 10 Sep 2026 18:30:07 -0700 Subject: [PATCH 02/25] [None][doc] Correct the draft-mirror saturation comments Skipping a context request whose draft mirror found no IndexMapper slot does not defer it to the next iteration: copy_batch_block_offsets runs later in the same iteration and IndexMapper::getCopyIndex TLLM_CHECKs on the unmapped id. Behaviour is unchanged; only the claim about it was wrong. Signed-off-by: Zhenhuan Chen --- .../kv_cache/kv_cache_manager_v2.py | 26 +++++++++++++------ 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 57a3140595ce..a51e0dc2d44d 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -3664,10 +3664,14 @@ def _mirror_draft_kv_cache(self, req: LlmRequest): see py_executor._adp_dummy_is_gen), and a disaggregated generation worker receives prompt KV rather than prefilling it. - Returns None when the IndexMapper is saturated. A context request can - retry next iteration; a generation request cannot, and skipping it only - defers the failure to copy_batch_block_offsets(), which asserts in C++ - on the unmapped request ID. + Returns None when the IndexMapper is saturated, which neither branch + survives. copy_batch_block_offsets() runs later in the SAME iteration + and feeds every id in the batch -- context ids included, see + IndexMapper::getCopyIndex -- to getIndex(), which TLLM_CHECKs on an + unmapped id (kvCacheManagerV2Utils.cpp). So the caller's `continue` on + the context path does not defer the request to a later iteration; + there is no later iteration. That predates the mirror refactor and is + left as is here; only the claim about it is corrected. """ kv_cache = self.kv_cache_map.get(req.py_request_id) if kv_cache is not None: @@ -3698,12 +3702,18 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): for req in scheduled_batch.context_requests: kv_cache = self._mirror_draft_kv_cache(req) if kv_cache is None: - # Retryable here, unlike the generation loop below: a - # context request has not drafted yet, so the next - # iteration can mirror it once slots free up. + # Pre-existing behaviour, kept deliberately: skipping here + # does NOT buy a retry, because copy_batch_block_offsets() + # asserts on this id later in the same iteration (see + # _mirror_draft_kv_cache). Saturation needs a disagg worker + # whose cancelled / retired-session requests pile up past + # the 2x slack, since the normal path frees the slot before + # start_transfer -- so this is a loud symptom of that, not a + # recoverable state. logger.warning( f"Draft KV cache mirror has no free IndexMapper slot for " - f"context request {req.py_request_id}; retrying next iteration." + f"context request {req.py_request_id}; this iteration " + f"will fail in copy_batch_block_offsets." ) continue if not self._resume_and_restore(req.py_request_id, kv_cache): From 81dff1ff0b861b148b05718fffaab84577ac8472 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Thu, 10 Sep 2026 18:35:44 -0700 Subject: [PATCH 03/25] [None][fix] Charge the external drafter's KV budget at its allocated dtype The budget split resolved the draft config without stripping the target's inherited fp8 KV algo, charging 1 byte/element for a pool allocated at 2 -- 2880 vs 5760 B/token on a K3 DEP16 GEN worker. Route both the cost and the allocation through the same helper. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/pyexecutor/_util.py | 57 +++++++- .../kv_cache/kv_cache_manager_v2.py | 25 ++++ .../kv_cache/test_kv_cache_budget_split.py | 136 ++++++++++++++++++ 3 files changed, 215 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 50576b75e62b..cf9b07464a5e 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -925,7 +925,10 @@ def _get_draft_cache_cost( # (e.g. EAGLE3: config says 1, runtime uses 4). # For PP, draft layers are only on the last rank (see # get_pp_layers), so only that rank should include draft cost. - effective_draft_config = self._get_effective_draft_config() + # _get_draft_kv_model_config(), not _get_effective_draft_config(): + # the cost charged here must be the cost of the pool that + # _create_one_model_draft_kv_cache_manager actually allocates. + effective_draft_config = self._get_draft_kv_model_config() draft_kv_cache_config = self._get_one_model_draft_kv_cache_config( kv_cache_config, self._max_seq_len) if self._speculative_config.spec_dec_mode.is_external_drafter(): @@ -1672,6 +1675,51 @@ def _get_effective_draft_config(self) -> ModelConfig: # layers as well. return self._model_engine.model.model_config + def _get_draft_kv_model_config(self) -> ModelConfig: + """The draft ModelConfig describing the KV pool as it is ALLOCATED. + + The args-level ``kv_cache_config.dtype`` sync stamps the TARGET's fp8 + KV algo onto every loaded model, including an external drafter. The + drafter stores and reads its pool in its weights dtype (DFlash + validates a bf16 pool and otherwise falls back to the max_seq_len-dense + private arena, which OOMs at long context), so the pool dtype must + follow the drafter, not the target. + + Every consumer that reasons about draft KV bytes must go through here. + Construction and the target/draft budget split reading different dtypes + is not cosmetic: the split then charges the draft fp8 bytes for a pool + allocated bf16, so the draft manager receives exactly HALF the tokens + the target gets. Because the capacity scheduler admits on the target + pool alone, the draft pool cannot backpressure -- it can only raise + "Draft KV cache context resize failed" out of + ``KVCacheManagerV2._prepare_draft_resources``, which is fatal to every + rank, as soon as resident context passes ~50% of target utilization. + """ + effective_draft_config = self._get_effective_draft_config() + if not self._speculative_config.spec_dec_mode.is_external_drafter(): + return effective_draft_config + quant_config = getattr(effective_draft_config, "quant_config", None) + if quant_config is None or not quant_config.quant_mode.has_fp8_kv_cache( + ): + return effective_draft_config + logger.info( + "External drafter KV pool keeps the drafter dtype; dropping " + "the fp8 KV quant algo inherited from the target.") + neutral_quant = copy.copy(quant_config) + neutral_quant.kv_cache_quant_algo = None + # QuantConfig.quant_mode and .layer_quant_mode are both cached_property + # and the copy carries the already-computed caches, so BOTH must be + # dropped for the mutation to take. Leaving layer_quant_mode stale + # silently no-ops the `draft_kv_config.dtype -> "auto"` guard in + # _create_one_model_draft_kv_cache_manager, which reads it. + neutral_quant.__dict__.pop("quant_mode", None) + neutral_quant.__dict__.pop("layer_quant_mode", None) + effective_draft_config = copy.copy(effective_draft_config) + effective_draft_config._frozen = False + effective_draft_config.quant_config = neutral_quant + effective_draft_config._frozen = True + return effective_draft_config + def _get_num_draft_layers(self) -> int: """Return the actual number of draft KV cache layers. @@ -1728,8 +1776,11 @@ def _create_one_model_draft_kv_cache_manager( spec_dec_layer_mask = self._get_one_model_draft_layer_mask() # Get the effective draft config (explicit draft_config if available, - # otherwise fall back to target model config for MTP). - effective_draft_config = self._get_effective_draft_config() + # otherwise fall back to target model config for MTP), with the + # target's inherited fp8 KV algo dropped for an external drafter. The + # budget split in _get_kv_size_per_token resolves it through the SAME + # helper, so the bytes/token it charges match the pool allocated here. + effective_draft_config = self._get_draft_kv_model_config() kv_cache_config = (kv_cache_config_override if kv_cache_config_override is not None else self._kv_cache_config) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index a51e0dc2d44d..5ee10cb09c99 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -3688,6 +3688,29 @@ def _mirror_draft_kv_cache(self, req: LlmRequest): kv_cache.stop_committing() return kv_cache + def _draft_pool_diagnostic(self) -> str: + """Draft-pool occupancy, for the resize-failure messages below. + + The draft manager mirrors the target's tokens but is sized from its own + byte budget, and the capacity scheduler admits on the TARGET pool alone + (`scheduler_v2` touches `draft_kv_cache_manager` only to suspend/free). + A draft pool smaller in tokens than the target therefore cannot + backpressure -- it can only raise, and the raise kills every rank. When + that happens the first question is always "how big was the draft pool + and how full was it", so answer it in the message rather than leaving + it to post-hoc arithmetic over the budget-split log line. + + Best-effort: never let a diagnostic mask the failure it describes. + """ + try: + live = sum(c.capacity for c in self.kv_cache_map.values()) + return ( + f" [draft pool: {len(self.kv_cache_map)} live caches holding " + f"{live} tokens, gpu_max_tokens={self._gpu_max_tokens}]" + ) + except Exception: # noqa: BLE001 - diagnostic only + return "" + def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): """Create/resize KV caches in the draft V2 manager for scheduled requests. @@ -3731,6 +3754,7 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): raise RuntimeError( f"Draft KV cache context resize failed for request " f"{req.py_request_id}: could not resize to {capacity} tokens" + f"{self._draft_pool_diagnostic()}" ) for req in scheduled_batch.generation_requests: @@ -3753,6 +3777,7 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): raise RuntimeError( f"Draft KV cache generation resize failed for request " f"{req.py_request_id}: could not resize to {new_cap} tokens" + f"{self._draft_pool_diagnostic()}" ) def _reuse_token_source(self, req: LlmRequest) -> Sequence[int]: diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py index 4f9554364d1c..ecde9f7376cd 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py @@ -1027,3 +1027,139 @@ def test_two_model_keeps_the_gpu_budget_whole(self, tmp_path): for call in calls: config = call.kwargs["kv_cache_config_override"] assert config.max_gpu_total_bytes == 10 * GB + + +class TestExternalDrafterKvDtype: + """The draft budget must be charged at the dtype the draft pool is ALLOCATED in. + + ``kv_cache_config.dtype: fp8`` stamps the target's fp8 KV algo onto an + external drafter's ModelConfig, but the drafter keeps a bf16 pool (dflash + rejects an fp8 pool and falls back to a max_seq_len-dense private arena). + The allocation path dropped the inherited algo; the cost path did not, so + the split charged 2880 B/token for a pool costing 5760, handed the draft + manager half the tokens the target got, and the GEN worker died in + ``_prepare_draft_resources`` at ~50% target utilization -- fatal to every + rank, because the capacity scheduler admits on the target pool alone. + """ + + # MLA drafter: kv_lora_rank 512 + qk_rope_head_dim 64 = 576, kv_factor 1, + # 5 layers. These are the numbers from the run that exposed the bug. + DRAFT_LAYERS = 5 + HEAD_DIM = 576 + FP8_SLOPE = DRAFT_LAYERS * HEAD_DIM # 2880 -- what the split wrongly charged + BF16_SLOPE = FP8_SLOPE * 2 # 5760 -- what the pool actually costs + + def _creator(self, mocker, draft_quant_config): + class DraftModelConfig: + quant_config = draft_quant_config + pretrained_config = SimpleNamespace( + num_hidden_layers=TestExternalDrafterKvDtype.DRAFT_LAYERS, + hidden_size=32, + num_attention_heads=4, + num_key_value_heads=1, + kv_lora_rank=512, + qk_rope_head_dim=64, + torch_dtype=None, + ) + + def get_num_attention_layers(self): + return TestExternalDrafterKvDtype.DRAFT_LAYERS + + target_model_config = SimpleNamespace(is_encoder_decoder=False) + draft_model_config = DraftModelConfig() + seen_draft_model_configs = [] + + class ProbeKVCacheManager(KVCacheManagerV2): + @staticmethod + def get_cache_size_per_token(model_config, *args, **kwargs): + if model_config is target_model_config: + return 0 + seen_draft_model_configs.append(model_config) + return KVCacheManagerV2.get_cache_size_per_token(model_config, *args, **kwargs) + + c = object.__new__(KvCacheCreator) + c._kv_cache_config = KvCacheConfig(dtype="fp8") + c._tokens_per_block = 64 + c._max_seq_len = 16384 + c._max_batch_size = 1 + # Read by _build_managers on the draft path (_util.py); the cost + # assertions below are per-token, so the value only has to exist. + c._max_num_tokens = 8192 + c._mapping = Mock(enable_attention_dp=True, tp_size=1) + c._mapping.has_cp_helix.return_value = False + c._mapping.pp_layers.return_value = list(range(self.DRAFT_LAYERS)) + c._mapping.is_last_pp_rank.return_value = True + # The real enum, not Mock(): a bare Mock answers True to EVERY + # predicate, so use_one_engine() and is_mtp_vanilla() both fire and the + # code walks branches an external drafter never takes. DSPARK satisfies + # is_external_drafter() via is_parallel_draft() and nothing else. + c._speculative_config = SimpleNamespace( + spec_dec_mode=SpeculativeDecodingMode.DSPARK, + max_draft_len=4, + ) + c._model_engine = SimpleNamespace(model=SimpleNamespace(model_config=target_model_config)) + c._draft_model_engine = None + c._draft_config = draft_model_config + c._kv_cache_manager_cls = ProbeKVCacheManager + c._is_disagg = True + c._is_encoder_decoder = Mock(return_value=False) + c._should_create_separate_draft_kv_cache = Mock(return_value=True) + c._get_num_draft_layers = Mock(return_value=self.DRAFT_LAYERS) + mocker.patch( + "tensorrt_llm._torch.pyexecutor._util.get_kv_cache_manager_cls", + return_value=ProbeKVCacheManager, + ) + return c, draft_model_config, seen_draft_model_configs + + def test_inherited_fp8_kv_algo_is_dropped_from_the_draft_cost(self, mocker): + from tensorrt_llm.models.modeling_utils import QuantConfig + from tensorrt_llm.quantization.mode import QuantAlgo + + # What the args-level kv_cache_config.dtype sync puts on the drafter. + inherited = QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) + assert inherited.quant_mode.has_fp8_kv_cache() + c, draft_model_config, seen = self._creator(mocker, inherited) + + cost = c._get_kv_size_per_token() + + # The draft pool is bf16, so the split must charge bf16 bytes/token. + assert cost.slope == self.BF16_SLOPE, ( + f"draft charged {cost.slope} B/token; the bf16 pool costs " + f"{self.BF16_SLOPE}. Charging {self.FP8_SLOPE} gives the draft " + f"manager half the tokens the target gets." + ) + assert len(seen) == 1 + assert not seen[0].quant_config.quant_mode.has_fp8_kv_cache() + # Both cached_property caches must be invalidated, or the + # draft_kv_config.dtype guard in the allocation path silently no-ops. + assert not seen[0].quant_config.layer_quant_mode.has_fp8_kv_cache() + # The drafter's own ModelConfig must not be mutated in place. + assert draft_model_config.quant_config is inherited + assert inherited.quant_mode.has_fp8_kv_cache() + + def test_cost_and_allocation_paths_agree(self, mocker): + """Both call sites must resolve the draft config through one helper.""" + from tensorrt_llm.models.modeling_utils import QuantConfig + from tensorrt_llm.quantization.mode import QuantAlgo + + c, _, seen = self._creator(mocker, QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8)) + c._get_kv_size_per_token() + + allocation_config = c._get_draft_kv_model_config() + assert ( + seen[0].quant_config.quant_mode.has_fp8_kv_cache() + == allocation_config.quant_config.quant_mode.has_fp8_kv_cache() + ) + + def test_non_external_drafter_is_untouched(self, mocker): + """MTP/Eagle3 share the target's layout, so nothing is neutralized.""" + from tensorrt_llm.models.modeling_utils import QuantConfig + from tensorrt_llm.quantization.mode import QuantAlgo + + inherited = QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) + c, draft_model_config, _ = self._creator(mocker, inherited) + # A real non-external mode, not a patched predicate: MTP_EAGLE shares the + # target's KV layout, which is exactly the case this asserts is untouched. + c._speculative_config.spec_dec_mode = SpeculativeDecodingMode.MTP_EAGLE + + assert c._get_draft_kv_model_config() is draft_model_config From a248176f8d4c3344399cccb3b9470c17d84f6029 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Thu, 10 Sep 2026 21:24:54 -0700 Subject: [PATCH 04/25] [None][fix] Bound the standalone drafter's draft-KV writes per request Upstream's #18343 clamps to the block-table width and drops _ctx_block_counts, which the MLA port's read-length truncation consumes -- so the managed-pool generation path raised AttributeError. Restore the count and bound both writes. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/speculative/dflash.py | 77 ++++++++++++------- .../test_kimi_k3_dspark_semantics.py | 26 ++++--- 2 files changed, 66 insertions(+), 37 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index c5ac3689b302..a1a56471d0e9 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -89,6 +89,22 @@ def dflash_position_ceiling(max_ctx: int, block_size: int, max_draft_len: int) - return int(max_ctx) + max(int(block_size), int(max_draft_len) + 1) +def dflash_allocated_ctx_limit( + block_counts: torch.Tensor, page_size: int, block_size: int +) -> torch.Tensor: + """Context length a request may advertise, given the pages it holds. + + Leaves ``block_size`` positions of room on purpose. The MLA path writes the + block's own latents at ``ctx_len .. ctx_len + block_size``, so stopping at + exactly the allocation puts the first of them on the first UNallocated page + -- whose block-table entry ``_refresh_ctx_block_tables`` clamped from a + negative placeholder to 0, i.e. another request's block. That is a silent + cross-request write, not an out-of-range fault, so the room has to be left + here rather than caught downstream. + """ + return (block_counts * page_size - block_size).clamp(min=0) + + @dataclass class DFlashSpecMetadata(SpecMetadata): """Metadata for DFlash speculative decoding. @@ -290,6 +306,9 @@ def __init__( self._ctx_page_table = None self._ctx_block_tables = None # [max_batch+1, max_blocks_per_seq] int32 self._ctx_block_indptr = None + # Pages the manager has actually handed this request, per batch row. + # Recovered from the negative placeholders before _refresh clamps them. + self._ctx_block_counts = None self._ctx_pool_idx = 0 self._ctx_block_divisor = 1 # encoded offset -> raw pool block index self._ctx_kv_indptr = None @@ -489,6 +508,7 @@ def _init_ctx_block_tables( self._ctx_block_indptr = torch.arange( 0, (num_slots + 1) * max_blocks, max_blocks, dtype=torch.int32, device="cuda" ) + self._ctx_block_counts = torch.zeros(num_slots, dtype=torch.long, device="cuda") logger.info( f"DFlash: ctx block tables {tuple(self._ctx_block_tables.shape)}, " f"pool_idx={self._ctx_pool_idx}, divisor={self._ctx_block_divisor}, " @@ -519,6 +539,10 @@ def _refresh_ctx_block_tables(self, attn_metadata, num_seqs: int) -> bool: # clone(): .to() is a no-op for an int64 source, and the in-place ops # below would then edit attn_metadata's own tensor. encoded = src[self._ctx_pool_idx, :num_seqs, 0].to(torch.int64).clone() + # Count the placeholders BEFORE the clamp below erases them. Unlike the + # private arena, whose slots each owned a fixed page range, a write past + # a request's allocation lands in whatever block 0 belongs to. + self._ctx_block_counts[:num_seqs].copy_((encoded >= 0).sum(dim=1)) decoded = encoded.clamp_(min=0).div_(self._ctx_block_divisor, rounding_mode="floor") self._ctx_block_tables[:num_seqs].copy_(decoded.to(torch.int32)) return True @@ -551,6 +575,7 @@ def _lazy_init_ctx_buffers( # pass falling back to the arena still rebinds once the real pool lands. self._ctx_kv_manager = draft_kv_cache_manager self._ctx_block_tables = None + self._ctx_block_counts = None max_batch = spec_metadata.max_num_requests # ctx_len is 1:1 with the target's positions, so max_seq_len bounds it. @@ -930,6 +955,15 @@ def _store_prefill_context( ctx_len_updates = {} offset = 0 num_contexts = attn_metadata.num_contexts + # A context write addresses pages through row i of the block table, so + # it has to fit that row's allocation as well as the arena's length. + # One tolist() for the whole loop: the lengths themselves come from + # _ctx_len_host, so the loop body below stays sync-free. + ctx_alloc = ( + (self._ctx_block_counts[:num_contexts] * self._ctx_page_size).tolist() + if self._ctx_block_tables is not None + else None + ) for i in range(num_contexts): req_id = spec_metadata.request_ids[i] slen = int(attn_metadata._seq_lens[i]) @@ -956,13 +990,14 @@ def _store_prefill_context( slot = self._req_to_slot[req_id] cur = ctx_len_updates.get(slot, self._ctx_len_host[slot]) - if cur + slen > self._max_ctx: + cap = self._max_ctx if ctx_alloc is None else min(self._max_ctx, ctx_alloc[i]) + if cur + slen > cap: # Request-level, like the no-free-slots path above: truncating # would silently draft from a stale prefix, but killing the # forward would take every other in-flight request with it. logger.warning( f"DFlash: ctx overflow on slot {slot} " - f"({cur} + {slen} > {self._max_ctx}); skipping its context " + f"({cur} + {slen} > {cap}); skipping its context " "store, so this request drafts nothing." ) self._req_to_slot.pop(req_id, None) @@ -1065,17 +1100,9 @@ def _forward_impl( ) spec_metadata._dflash_worker = self # Before any store: prefill and decode both address pages through it. - refreshed = self._refresh_ctx_block_tables(attn_metadata, batch_size) - if self._ctx_block_tables is not None and not refreshed: - # False is normal only while no paged pool is bound (the private - # arena addresses slots directly). Once one IS bound, the counts - # drive the advertised context length, so a missed refresh leaves - # them at zero and every drafter attends an empty context -- no - # error, just acceptance quietly collapsing. - raise RuntimeError( - "DFlash: draft block tables are bound to a paged pool but this " - "iteration's attn_metadata carried no draft_kv_cache_block_offsets." - ) + # Returning False here means an empty batch -- the missing-offsets case + # raises inside, with the metadata type in the message. + self._refresh_ctx_block_tables(attn_metadata, batch_size) # Save context lengths so both warmup and a failed forward can roll # back the in-place _ctx_len updates made during drafting. @@ -1527,11 +1554,13 @@ def prepare_1st_drafter_inputs( col_idx = self._ctx_len[slots].unsqueeze(1) + offsets_kp1.unsqueeze(0) write_mask = offsets_kp1.unsqueeze(0) < gen_num_accepted_long.unsqueeze(1) if self._ctx_block_tables is not None: - # Bound by the table width. A per-request bound is not - # recoverable here: V2 writes 0 for BAD_PAGE_INDEX and V1 - # leaves the tail stale, so a placeholder reads as block 0. - ctx_capacity = self._ctx_block_tables.size(1) * self._ctx_page_size - col_idx = col_idx.clamp(max=ctx_capacity - 1) + # Bound by what the manager allocated for THIS request, not + # by the table width: the masked entries below are still + # written (fixed-size, for graph safety), and every page + # past the allocation is a placeholder that _refresh clamped + # to 0 -- i.e. another request's block. + gen_capacity = self._ctx_block_counts[gen_rows_out] * self._ctx_page_size + col_idx = torch.minimum(col_idx, (gen_capacity - 1).unsqueeze(1)).clamp_(min=0) else: col_idx = col_idx.clamp(max=self._max_ctx - 1) @@ -1585,16 +1614,8 @@ def prepare_1st_drafter_inputs( # block table left clamped to 0 -- another request's data. # Truncating the advertised length keeps the two consistent: # the drafter attends a short context instead of a wrong one. - # - # Leave the block room too. The MLA path writes its own latents - # at ctx_len..ctx_len+block_size (_build_mla_block_fixup), so - # stopping at `allocated` puts the first block position on the - # first UNallocated page -- whose block-table entry _refresh - # clamped from its negative placeholder to 0, i.e. another - # request's block. That is a silent cross-request write, not an - # out-of-range fault. - allocated = (self._ctx_block_counts * self._ctx_page_size - block_size).clamp_( - min=0 + allocated = dflash_allocated_ctx_limit( + self._ctx_block_counts, self._ctx_page_size, block_size ) num_ctx_per_req_t = torch.minimum(num_ctx_per_req_t, allocated[gen_rows_out]) noise_embedding = noise_embed_2d diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py index 01f4f60f40ac..dd92926d7d26 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py @@ -1173,27 +1173,35 @@ def test_mla_dspark_rope_conventions_agree_on_scores(): def test_mla_block_fixup_stays_inside_the_allocation(): """A context that fills its allocation must still leave the block room. - dflash.py truncates ctx_len to what the draft KV manager allocated, but the - MLA path then writes the block's own latents at ctx_len..ctx_len+block_size. - Stopping at exactly `allocated` puts the first block position on the first - unallocated page, whose block-table entry _refresh_ctx_block_tables clamped - from a negative placeholder to 0 -- another request's block. Silent - cross-request corruption, not a fault, so only an explicit bound catches it. + The bound comes from dflash.py, so it is called here rather than restated: + a test that computed `allocated - block_size` itself would still pass + against a production path truncating to `allocated`. The MLA writes the + block's own latents at ctx_len..ctx_len+block_size, so that regression puts + the first of them on the first unallocated page, whose block-table entry + _refresh_ctx_block_tables clamped from a negative placeholder to 0 -- + another request's block. Silent cross-request corruption, not a fault. """ from tensorrt_llm._torch.models.modeling_dspark import _build_mla_block_fixup + from tensorrt_llm._torch.speculative.dflash import dflash_allocated_ctx_limit page_size, block_size, allocated_pages = 8, 3, 2 allocated = allocated_pages * page_size # Entries past the allocation are the clamped placeholders: physical 0. page_tables = torch.tensor([[41, 42, 0, 0]]) - ctx_len = torch.tensor([max(allocated - block_size, 0)]) + ctx_len = dflash_allocated_ctx_limit(torch.tensor([allocated_pages]), page_size, block_size) fixup = _build_mla_block_fixup(ctx_len, page_tables, block_size, page_size) - allocated_ids = page_tables[0, :allocated_pages].tolist() - assert set(fixup.pages.flatten().tolist()) <= set(allocated_ids) + # Exactly the last allocated page, not merely a subset of the allocation: + # `<= {41, 42}` also passes for a block that landed a page early. + assert set(fixup.pages.flatten().tolist()) == {42} assert int(fixup.seq_lens_i32[0]) <= allocated + # Negative control: the value the production path yields if the block room + # is dropped. It must reach the unallocated page, or this bound is untested. + overrun = _build_mla_block_fixup(torch.tensor([allocated]), page_tables, block_size, page_size) + assert 0 in overrun.pages.flatten().tolist() + def test_mla_rope_table_follows_runtime_max_seq_len(): """The position table is sized by what is served, not what is advertised. From c1c19a09cccee458062a5536ced510e6a4c745a4 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Thu, 10 Sep 2026 21:25:41 -0700 Subject: [PATCH 05/25] [None][chore] Address review on the MLA drafter port Narrow the draft-KV dtype carve-out to DFlash/DSpark, declare _runtime_position_ceiling, key the MLA RoPE table on its cap, drop a dead try/except, and cover both aux-capture conventions. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/models/modeling_dflash.py | 7 + tensorrt_llm/_torch/models/modeling_dspark.py | 27 ++-- tensorrt_llm/_torch/pyexecutor/_util.py | 26 ++-- .../kv_cache/kv_cache_manager_v2.py | 14 +- .../test_kimi_k3_dflash_scaffold.py | 133 ++++++++++++++++++ 5 files changed, 180 insertions(+), 27 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index 5bd49cc08074..c25e1a050182 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -335,6 +335,13 @@ class DFlashForCausalLM(nn.Module): # is about which kernel reads it. _paged_ctx_cache = False + # Positions the worker will actually serve, published once by + # DFlashDrafter._lazy_init_ctx_buffers before the first forward. A drafter + # that sizes an absolute-position table reads it in place of its + # config-derived cap; None means no worker has run (direct construction in + # tests), so the config cap stands. + _runtime_position_ceiling = None + def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): """Build the draft model, resolving its architecture from the draft config (falling back to a model_type-derived name when the checkpoint uses a diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 0fbdb8553bff..253b077d3a7d 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -2676,6 +2676,7 @@ def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): self._mla_rope_params = rope self._mla_workspace = None self._mla_freqs = None + self._mla_freqs_cap = None self._mla_noop_weight = None # HF DeepSeek folds YaRN's attention scaling into the softmax scale: # softmax_scale = mscale^2 / sqrt(qk_nope + qk_rope). NOT 1/sqrt(576) -- @@ -2835,16 +2836,21 @@ def _mla_freqs_cis(self, device): on. The config-derived cap is the fallback for direct construction in tests, where no worker runs. """ - if self._mla_freqs is None: - rope = dict(self._mla_rope_params) - runtime_cap = getattr(self, "_runtime_position_ceiling", None) - if runtime_cap is not None: - # Outright, not max(): the published ceiling is already - # min(runtime, max_position_embeddings) + lookahead, so it is - # both sufficient and the only thing that keeps the table small - # when max_seq_len is unset and the config cap is the - # checkpoint's advertised 1,048,576. - rope["max_positions"] = int(runtime_cap) + rope = dict(self._mla_rope_params) + runtime_cap = self._runtime_position_ceiling + if runtime_cap is not None: + # Outright, not max(): the published ceiling is already + # min(runtime, max_position_embeddings) + lookahead, so it is + # both sufficient and the only thing that keeps the table small + # when max_seq_len is unset and the config cap is the + # checkpoint's advertised 1,048,576. + rope["max_positions"] = int(runtime_cap) + # Keyed on the cap, not just built once: _lazy_init_ctx_buffers + # re-publishes the ceiling when the KV-estimation probe manager is + # swapped for the real one, and drafter forwards do run during + # estimation. A table built short for a later, larger cap would index + # out of range on freqs[positions]. + if self._mla_freqs is None or self._mla_freqs_cap != rope["max_positions"]: self._mla_freqs = build_dspark_mla_yarn_freqs_cis( dim=self.qk_rope_head_dim, base=rope["theta"], @@ -2857,6 +2863,7 @@ def _mla_freqs_cis(self, device): mscale_all_dim=rope["mscale_all_dim"], device=device, ) + self._mla_freqs_cap = rope["max_positions"] return self._mla_freqs def _build_fused_kv_buffers(self) -> None: diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index cf9b07464a5e..a2c4b488c203 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1679,7 +1679,7 @@ def _get_draft_kv_model_config(self) -> ModelConfig: """The draft ModelConfig describing the KV pool as it is ALLOCATED. The args-level ``kv_cache_config.dtype`` sync stamps the TARGET's fp8 - KV algo onto every loaded model, including an external drafter. The + KV algo onto every loaded model, including a standalone drafter. The drafter stores and reads its pool in its weights dtype (DFlash validates a bf16 pool and otherwise falls back to the max_seq_len-dense private arena, which OOMs at long context), so the pool dtype must @@ -1696,7 +1696,16 @@ def _get_draft_kv_model_config(self) -> ModelConfig: rank, as soon as resident context passes ~50% of target utilization. """ effective_draft_config = self._get_effective_draft_config() - if not self._speculative_config.spec_dec_mode.is_external_drafter(): + # Narrower than is_external_drafter(), matching + # _should_create_separate_draft_kv_cache: PARD and + # DRAFT_TARGET_ONE_MODEL reach here too, via + # should_use_separate_draft_kv_cache, and their draft checkpoints can + # carry a genuine fp8 KV algo of their own -- with dtype="auto", + # validate_and_set_kv_cache_quant returns early and keeps it. Dropping + # that would allocate bf16 while their attention modules, built from + # the un-neutralized config, still read and write fp8. + spec_dec_mode = self._speculative_config.spec_dec_mode + if not (spec_dec_mode.is_dflash() or spec_dec_mode.is_dspark()): return effective_draft_config quant_config = getattr(effective_draft_config, "quant_config", None) if quant_config is None or not quant_config.quant_mode.has_fp8_kv_cache( @@ -1709,15 +1718,16 @@ def _get_draft_kv_model_config(self) -> ModelConfig: neutral_quant.kv_cache_quant_algo = None # QuantConfig.quant_mode and .layer_quant_mode are both cached_property # and the copy carries the already-computed caches, so BOTH must be - # dropped for the mutation to take. Leaving layer_quant_mode stale - # silently no-ops the `draft_kv_config.dtype -> "auto"` guard in - # _create_one_model_draft_kv_cache_manager, which reads it. + # dropped for the mutation to take: _create_kv_cache_manager reads + # quant_mode off this copy, and layer_quant_mode is the pair's other + # half, stale in the same way. neutral_quant.__dict__.pop("quant_mode", None) neutral_quant.__dict__.pop("layer_quant_mode", None) + # No _frozen dance: ModelConfig.__setattr__ exempts quant_config by + # name, and restoring _frozen to True would freeze a copy whose source + # may not have been frozen. effective_draft_config = copy.copy(effective_draft_config) - effective_draft_config._frozen = False effective_draft_config.quant_config = neutral_quant - effective_draft_config._frozen = True return effective_draft_config def _get_num_draft_layers(self) -> int: @@ -1777,7 +1787,7 @@ def _create_one_model_draft_kv_cache_manager( # Get the effective draft config (explicit draft_config if available, # otherwise fall back to target model config for MTP), with the - # target's inherited fp8 KV algo dropped for an external drafter. The + # target's inherited fp8 KV algo dropped for a standalone drafter. The # budget split in _get_kv_size_per_token resolves it through the SAME # helper, so the bytes/token it charges match the pool allocated here. effective_draft_config = self._get_draft_kv_model_config() diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 5ee10cb09c99..89ec6ae2c605 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -3700,16 +3700,12 @@ def _draft_pool_diagnostic(self) -> str: and how full was it", so answer it in the message rather than leaving it to post-hoc arithmetic over the budget-split log line. - Best-effort: never let a diagnostic mask the failure it describes. """ - try: - live = sum(c.capacity for c in self.kv_cache_map.values()) - return ( - f" [draft pool: {len(self.kv_cache_map)} live caches holding " - f"{live} tokens, gpu_max_tokens={self._gpu_max_tokens}]" - ) - except Exception: # noqa: BLE001 - diagnostic only - return "" + live = sum(c.capacity for c in self.kv_cache_map.values()) + return ( + f" [draft pool: {len(self.kv_cache_map)} live caches holding " + f"{live} tokens, gpu_max_tokens={self._gpu_max_tokens}]" + ) def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): """Create/resize KV caches in the draft V2 manager for scheduled requests. diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dflash_scaffold.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dflash_scaffold.py index fdd023bae7f8..71bd32bff6b0 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dflash_scaffold.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dflash_scaffold.py @@ -462,3 +462,136 @@ def __call__(self, hidden_states, block_residual, num_snapshots, attn_metadata, seen[last_idx], _apply_attn_res(h, br[:num_layers], proj_last, norm_last) ), "tail used the last layer's weights instead of the output-side ones" assert not torch.allclose(seen[last_idx], h), "tail captured the raw prefix sum" + + +@pytest.mark.parametrize("aggregated", [True, False]) +def test_aux_capture_taps_the_selected_stream(monkeypatch, aggregated): + """Both KIMI_K3_AUX_ATTN_RES_STREAM conventions must be executed. + + The switch picks which residual value the drafter sees, and a mismatch only + lowers acceptance -- nothing raises. The test above stubs the layer and so + never runs either branch; this one calls the real layer forward. + """ + pytest.importorskip("fla") + from types import SimpleNamespace + + from tensorrt_llm._torch.models import modeling_kimi_linear as mkl + + monkeypatch.setattr(mkl, "_AUX_ATTN_RES_STREAM_ENABLED", aggregated) + + torch.manual_seed(0) + num_tokens, hidden, k_max, num_snapshots = 3, 8, 4, 2 + proj = torch.nn.Linear(hidden, 1, bias=False, dtype=torch.float32) + norm = mkl.KimiK3RMSNorm(hidden, eps=1e-6, dtype=torch.float32) + with torch.no_grad(): + proj.weight.copy_(torch.randn(1, hidden)) + norm.weight.copy_(torch.randn(hidden) + 1.0) + + prefix_sum = torch.randn(num_tokens, hidden) + block_residual = torch.randn(k_max, num_tokens, hidden) + + seen = {} + layer = SimpleNamespace( + self_attention_res_proj=proj, + self_attention_res_norm=norm, + # Not a block boundary, so prefix_sum survives to the end of forward. + layer_idx=1, + attn_res_block_size=4, + input_layernorm=lambda h: h, + is_kda=True, + linear_attn=lambda h, md: torch.zeros_like(h), + mlp_res_proj=proj, + mlp_res_norm=norm, + post_attention_layernorm=lambda h: h, + is_moe=False, + mlp=lambda h: torch.zeros_like(h), + ) + spec_md = SimpleNamespace( + maybe_capture_hidden_states=lambda lid, h, r: seen.__setitem__(lid, h.clone()) + ) + + mkl.KimiLinearDecoderLayer.forward( + layer, + prefix_sum.clone(), + block_residual.clone(), + num_snapshots, + attn_metadata=SimpleNamespace(), + capture=(spec_md, 0), + ) + + mixture = mkl._apply_attn_res(prefix_sum, block_residual[:num_snapshots], proj, norm) + expected, other = (mixture, prefix_sum) if aggregated else (prefix_sum, mixture) + torch.testing.assert_close(seen[0], expected) + assert not torch.allclose(seen[0], other), ( + "the tapped tensor did not follow _AUX_ATTN_RES_STREAM_ENABLED" + ) + + +@pytest.mark.parametrize("aggregated", [True, False]) +def test_aux_capture_tail_follows_the_same_switch(monkeypatch, aggregated): + """The final layer has no successor, so the tail recompute must match. + + A tail that ignored the switch would hand the last captured layer a tensor + from the other convention -- the one inconsistency the in-loop test cannot + see, because the loop never reaches the final layer's tap. + """ + pytest.importorskip("fla") + from types import SimpleNamespace + + from tensorrt_llm._torch.models import modeling_kimi_linear as mkl + + monkeypatch.setattr(mkl, "_AUX_ATTN_RES_STREAM_ENABLED", aggregated) + + torch.manual_seed(1) + num_layers, num_tokens, hidden = 3, 2, 8 + out_proj = torch.nn.Linear(hidden, 1, bias=False, dtype=torch.float32) + out_norm = mkl.KimiK3RMSNorm(hidden, eps=1e-6, dtype=torch.float32) + with torch.no_grad(): + out_proj.weight.copy_(torch.randn(1, hidden) * 3.0) + out_norm.weight.copy_(torch.randn(hidden) + 1.0) + + class _Layer: + def __init__(self, idx): + self.layer_idx = idx + + def __call__(self, hidden_states, block_residual, num_snapshots, attn_metadata, capture): + block_residual[num_snapshots] = hidden_states * (self.layer_idx + 1) + return hidden_states + (self.layer_idx + 1), num_snapshots + 1 + + layers = [_Layer(i) for i in range(num_layers)] + last_idx = num_layers - 1 + seen = {} + # Only the final layer: the in-loop tap fires for layers[i-1], so this set + # leaves the tail as the single capture site. + spec_md = SimpleNamespace( + _capture_layer_set=frozenset({last_idx}), + maybe_capture_hidden_states=lambda lid, h, r: seen.__setitem__(lid, h.clone()), + ) + embeds = torch.randn(num_tokens, hidden, dtype=torch.float32) + fake = SimpleNamespace( + embed_tokens=lambda ids: embeds, + layers=layers, + norm=lambda h: h, + output_attn_res_proj=out_proj, + output_attn_res_norm=out_norm, + num_attn_res_snapshots=num_layers, + ) + + mkl.KimiLinearModel.forward( + fake, + attn_metadata=SimpleNamespace(num_tokens=num_tokens), + input_ids=torch.zeros(num_tokens, dtype=torch.int32), + spec_metadata=spec_md, + ) + + br = torch.empty(num_layers, num_tokens, hidden) + h = embeds + for i in range(num_layers): + br[i] = h * (i + 1) + h = h + (i + 1) + mixture = mkl._apply_attn_res(h, br[:num_layers], out_proj, out_norm) + expected, other = (mixture, h) if aggregated else (h, mixture) + torch.testing.assert_close(seen[last_idx], expected) + assert not torch.allclose(seen[last_idx], other), ( + "the tail tensor did not follow _AUX_ATTN_RES_STREAM_ENABLED" + ) From cdca3487ea1bcf39a829d5d996c72f6d8cd8fdeb Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Thu, 10 Sep 2026 21:46:27 -0700 Subject: [PATCH 06/25] [None][doc] State which V2 backend makes the draft page count per-request copyBatchBlockOffsetsToDeviceKernel maps BAD_PAGE_INDEX to 0, so on the default cpp backend the count saturates at the table width and the bound degenerates to upstream's. It is per-request only on the python backend, which keeps the -1. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/speculative/dflash.py | 29 ++++++++++++++--------- 1 file changed, 18 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index a1a56471d0e9..801e4abe0a4b 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -101,6 +101,16 @@ def dflash_allocated_ctx_limit( negative placeholder to 0, i.e. another request's block. That is a silent cross-request write, not an out-of-range fault, so the room has to be left here rather than caught downstream. + + How tight this is depends on the V2 backend, and only one of the two keeps + the information. ``copyBatchBlockOffsetsToDeviceKernel`` + (kvCacheManagerV2Utils.cu) writes ``0`` for ``BAD_PAGE_INDEX``, so under the + default ``cpp`` backend every row counts full and this degenerates to the + block-table-width bound upstream already applied -- correct, just not + per-request. The ``python`` backend propagates ``BAD_PAGE_INDEX`` through + ``_copy_swa_block_offsets_with_scratch``, and there the bound is real. + Sized from the counts either way so the tight case needs no second code + path; do not read a per-request guarantee into it on ``cpp``. """ return (block_counts * page_size - block_size).clamp(min=0) @@ -306,8 +316,8 @@ def __init__( self._ctx_page_table = None self._ctx_block_tables = None # [max_batch+1, max_blocks_per_seq] int32 self._ctx_block_indptr = None - # Pages the manager has actually handed this request, per batch row. - # Recovered from the negative placeholders before _refresh clamps them. + # Pages the manager handed this request, per batch row. Saturates at + # the table width on the default backend -- dflash_allocated_ctx_limit. self._ctx_block_counts = None self._ctx_pool_idx = 0 self._ctx_block_divisor = 1 # encoded offset -> raw pool block index @@ -539,9 +549,8 @@ def _refresh_ctx_block_tables(self, attn_metadata, num_seqs: int) -> bool: # clone(): .to() is a no-op for an int64 source, and the in-place ops # below would then edit attn_metadata's own tensor. encoded = src[self._ctx_pool_idx, :num_seqs, 0].to(torch.int64).clone() - # Count the placeholders BEFORE the clamp below erases them. Unlike the - # private arena, whose slots each owned a fixed page range, a write past - # a request's allocation lands in whatever block 0 belongs to. + # Before the clamp erases the placeholders. Informative only on the + # python V2 backend; the cpp one already mapped BAD_PAGE_INDEX to 0. self._ctx_block_counts[:num_seqs].copy_((encoded >= 0).sum(dim=1)) decoded = encoded.clamp_(min=0).div_(self._ctx_block_divisor, rounding_mode="floor") self._ctx_block_tables[:num_seqs].copy_(decoded.to(torch.int32)) @@ -956,7 +965,7 @@ def _store_prefill_context( offset = 0 num_contexts = attn_metadata.num_contexts # A context write addresses pages through row i of the block table, so - # it has to fit that row's allocation as well as the arena's length. + # it must fit that row's allocation as well as the arena's length. # One tolist() for the whole loop: the lengths themselves come from # _ctx_len_host, so the loop body below stays sync-free. ctx_alloc = ( @@ -1554,11 +1563,9 @@ def prepare_1st_drafter_inputs( col_idx = self._ctx_len[slots].unsqueeze(1) + offsets_kp1.unsqueeze(0) write_mask = offsets_kp1.unsqueeze(0) < gen_num_accepted_long.unsqueeze(1) if self._ctx_block_tables is not None: - # Bound by what the manager allocated for THIS request, not - # by the table width: the masked entries below are still - # written (fixed-size, for graph safety), and every page - # past the allocation is a placeholder that _refresh clamped - # to 0 -- i.e. another request's block. + # Bound by what the manager handed THIS request: the masked + # entries below are still written (fixed-size, for graph + # safety) and a placeholder page resolves to another's. gen_capacity = self._ctx_block_counts[gen_rows_out] * self._ctx_page_size col_idx = torch.minimum(col_idx, (gen_capacity - 1).unsqueeze(1)).clamp_(min=0) else: From 01c4e32bbf3aa174249f7e91fb4d839352d9db3d Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Fri, 11 Sep 2026 00:58:27 -0700 Subject: [PATCH 07/25] [None][chore] Build the MLA drafter's YaRN table with RopeEmbeddingUtils create_sinusoidal_positions_yarn is the same HF DeepSeek-V2 transcription the hand-rolled helpers restated. duplicate_data=True then sliced: the util's two modes disagree by 1 ulp on 6 entries, and the duplicated one is what ran. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/models/modeling_dspark.py | 78 ++++++++----------- 1 file changed, 34 insertions(+), 44 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 253b077d3a7d..bfffd6fdadc9 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -90,6 +90,7 @@ class or the edge would become a cycle. from torch import nn from transformers import PretrainedConfig +from tensorrt_llm.functional import RopeEmbeddingUtils from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping from tensorrt_llm.quantization.mode import QuantAlgo @@ -2030,29 +2031,6 @@ def _yarn_get_mscale(scale: float, mscale: float = 1.0) -> float: return 0.1 * mscale * math.log(scale) + 1.0 -def _yarn_find_correction_dim( - num_rotations: float, dim: int, base: float, max_position_embeddings: int -) -> float: - return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / ( - 2 * math.log(base) - ) - - -def _yarn_find_correction_range( - low_rot: float, high_rot: float, dim: int, base: float, max_position_embeddings: int -) -> tuple[int, int]: - low = math.floor(_yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings)) - high = math.ceil(_yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings)) - return max(low, 0), min(high, dim - 1) - - -def _yarn_linear_ramp_mask(minimum: float, maximum: float, dim: int) -> torch.Tensor: - if minimum == maximum: - maximum += 0.001 # Prevent singularity. - linear_func = (torch.arange(dim, dtype=torch.float32) - minimum) / (maximum - minimum) - return torch.clamp(linear_func, 0, 1) - - def build_dspark_mla_yarn_rope( *, dim: int, @@ -2068,33 +2046,40 @@ def build_dspark_mla_yarn_rope( ) -> tuple[torch.Tensor, torch.Tensor]: """YaRN cos/sin tables for an MLA drafter's ``qk_rope_head_dim`` slice. - Transcribed from HF ``DeepseekV3YarnRotaryEmbedding``: the drafters are - distilled under HF/vLLM numerics, so the table is built here rather than - routed through ``RopeParams`` (whose MLA convention carries the fused - kernel's ``duplicate_data`` / GPT-J packing). + Repacks ``RopeEmbeddingUtils.create_sinusoidal_positions_yarn``, which is + the same HF DeepSeek-V2 transcription the drafters are distilled under -- + verified bit-identical to the hand-rolled table this replaced, on the K3 + Inferact parameters. What is deliberately NOT reused is + ``RopeParams.create_rope_const_params()``: it pins the fused MLA kernel's + GPT-J packing, which rotates differently and costs acceptance silently. + + ``duplicate_data=True`` even though only the first half is read as a + complex pair. The util's two modes disagree by 1 ulp on 6 of 8224 entries + (torch's ``cos`` vectorises a ``[n, dim, 1]`` input differently than a + ``[n, dim/2, 1]`` one), and the duplicated one is what every measured + drafter run used. Returns ``(cos, sin)``, both ``[max_position_embeddings, dim]`` fp32, with the frequency half duplicated so ``rotate_half`` applies. """ - freq_extra = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) - freq_inter = 1.0 / ( - scaling_factor * base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim) + # The util returns a flat numpy [pos][slot](cos, sin); split the pair. + _, packed = RopeEmbeddingUtils.create_sinusoidal_positions_yarn( + num_pos=max_position_embeddings, + dim=dim, + base=base, + scaling_factor=scaling_factor, + original_max_position_embeddings=original_max_position_embeddings, + beta_fast=beta_fast, + beta_slow=beta_slow, + mscale=mscale, + mscale_all_dim=mscale_all_dim, + duplicate_data=True, ) - low, high = _yarn_find_correction_range( - beta_fast, beta_slow, dim, base, original_max_position_embeddings + table = torch.from_numpy(packed).reshape(max_position_embeddings, dim, 2) + return ( + table[..., 0].contiguous().to(device), + table[..., 1].contiguous().to(device), ) - inv_freq_mask = 1.0 - _yarn_linear_ramp_mask(low, high, dim // 2) - inv_freq = freq_inter * (1 - inv_freq_mask) + freq_extra * inv_freq_mask - - t = torch.arange(max_position_embeddings, dtype=torch.float32) - freqs = torch.outer(t, inv_freq) - # HF divides the two mscales; both published drafters set them equal, so - # this is 1.0 there. Kept general because the checkpoint declares both. - scale = _yarn_get_mscale(scaling_factor, mscale) / _yarn_get_mscale( - scaling_factor, mscale_all_dim - ) - emb = torch.cat((freqs, freqs), dim=-1) - return (emb.cos() * scale).to(device), (emb.sin() * scale).to(device) def build_dspark_mla_yarn_freqs_cis(*, device: str = "cuda", **kwargs) -> torch.Tensor: @@ -2127,6 +2112,11 @@ def apply_dspark_mla_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) gathered at the query positions. HF interleaves the rope slice, so it first reshapes ``(d//2, 2) -> transpose -> (d,)`` and only then applies the half-split ``rotate_half``; the net rotation is GPT-J style. + + Equal to ``RotaryEmbedding.apply_rotary_pos_emb(..., is_neox=False)`` up to + an even/odd lane split of its output (checked exactly, max diff 0.0). Kept + separate because that helper unsqueezes cos/sin itself, so routing through + it would push a rank contract onto every caller for no numerical gain. """ d = x.shape[-1] x = x.reshape(*x.shape[:-1], d // 2, 2).transpose(-1, -2).reshape(*x.shape[:-1], d) From dd00d571b18a11f864a72df9c616697861ef2bdb Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Fri, 11 Sep 2026 01:51:36 -0700 Subject: [PATCH 08/25] [None][chore] Trim the MLA drafter's comments to what cannot be re-derived Drops change history, design justification and explanation already carried by another docstring, from the ten longest blocks: -50 lines. Every measured number, error string and file:line pointer kept. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/models/modeling_dflash.py | 21 ++-- tensorrt_llm/_torch/models/modeling_dspark.py | 99 +++++++------------ tensorrt_llm/_torch/pyexecutor/_util.py | 26 ++--- tensorrt_llm/_torch/speculative/dflash.py | 50 ++++------ 4 files changed, 73 insertions(+), 123 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index c25e1a050182..fb8037eb53c1 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -869,18 +869,15 @@ def _assert_backbone_complete(self, weights: Dict, weight_mapper=None) -> None: The truth is the constructed module tree, not a hand-kept list that rots as the backbone changes. - Checked at module granularity. For a PLAIN module that is the hole the - flag cannot close: the loader skips one whose subtree filters to - nothing (modeling_utils.py `if module_weights:`) whatever the flag - says. For a FUSED module `allow_partial_loading=False` would catch it - (linear.py asserts all three shards) -- but the flag has to stay True - for the target-shared modules, so the check covers that case here - instead, and requires every component rather than any. - - Missing parameters INSIDE a present component stay tolerated: a - checkpoint with all three weights but only `q_proj.bias` takes the same - per-shard copy and leaves the rest at `torch.empty`. Module-granular - checking cannot see that and does not pretend to. + Module granularity. A PLAIN module is the hole `allow_partial_loading` + cannot close -- the loader skips one whose subtree filters to nothing + (modeling_utils.py `if module_weights:`) whatever the flag says. A FUSED + module `allow_partial_loading=False` would catch (linear.py asserts all + three shards), but the flag must stay True for the target-shared + modules, so this requires every component rather than any. + + Missing parameters INSIDE a present component stay tolerated: all three + weights but only `q_proj.bias` leaves the rest at `torch.empty`. """ provided = set(weights) diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index bfffd6fdadc9..8ed54b704b75 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -2013,15 +2013,11 @@ def load_weights_from_target_model(self, target_model): # ---------------------------------------------------------------------------- # MLA-shaped standalone drafter backbone (Inferact/Kimi-K3-DSpark). # -# A weight container, not a servable model: ``DFlashForCausalLM`` builds this -# through the registry only to own the drafter's modules, then runs its own -# hand-written block decode over them (see ``MLADSparkForCausalLM``). Nothing -# here registers with the attention backend or the KV cache manager, so the -# layers deliberately have no working ``forward``. -# -# Not a ``DeepseekV3`` reuse: that model carries MoE branches, a fused -# ``qkv_a_proj`` switch, quantization and KV-manager registration, none of -# which this bf16 five-layer drafter has. +# A weight container, not a servable model: nothing here registers with the +# attention backend or the KV cache manager, so the layers have no working +# ``forward``. ``MLADSparkForCausalLM`` runs its own block decode over them. +# Not a ``DeepseekV3`` subclass -- that carries MoE, a fused ``qkv_a_proj`` +# switch, quantization and KV-manager registration, none of which apply. # ---------------------------------------------------------------------------- @@ -2046,18 +2042,14 @@ def build_dspark_mla_yarn_rope( ) -> tuple[torch.Tensor, torch.Tensor]: """YaRN cos/sin tables for an MLA drafter's ``qk_rope_head_dim`` slice. - Repacks ``RopeEmbeddingUtils.create_sinusoidal_positions_yarn``, which is - the same HF DeepSeek-V2 transcription the drafters are distilled under -- - verified bit-identical to the hand-rolled table this replaced, on the K3 - Inferact parameters. What is deliberately NOT reused is - ``RopeParams.create_rope_const_params()``: it pins the fused MLA kernel's - GPT-J packing, which rotates differently and costs acceptance silently. + Repacks ``RopeEmbeddingUtils.create_sinusoidal_positions_yarn``, the same HF + DeepSeek-V2 transcription the drafters are distilled under. NOT reused: + ``RopeParams.create_rope_const_params()`` pins the fused MLA kernel's GPT-J + packing and hard-codes ``device='cuda'``. - ``duplicate_data=True`` even though only the first half is read as a - complex pair. The util's two modes disagree by 1 ulp on 6 of 8224 entries - (torch's ``cos`` vectorises a ``[n, dim, 1]`` input differently than a - ``[n, dim/2, 1]`` one), and the duplicated one is what every measured - drafter run used. + ``duplicate_data=True`` though only the first half is read as a complex + pair: the util's two modes disagree by 1 ulp on 6 of 8224 entries, and the + duplicated one is what every measured drafter run used. Returns ``(cos, sin)``, both ``[max_position_embeddings, dim]`` fp32, with the frequency half duplicated so ``rotate_half`` applies. @@ -2089,16 +2081,11 @@ def build_dspark_mla_yarn_freqs_cis(*, device: str = "cuda", **kwargs) -> torch. fused ``cute_dsl_dspark_rmsnorm_rope`` kernel consumes, where the last dim of the rotated slice is read as (re, im) pairs. - HF DeepSeek instead de-interleaves and then rotates halves. The two produce - the *same values in a different order*: with even/odd the pair members, - HF writes ``even*cos - odd*sin`` to lane i and ``odd*cos + even*sin`` to lane - i + dim/2, while this one writes them to lanes 2i and 2i+1. That permutation - of the rope slice cancels inside ``q_rope . k_rope``, so the drafter is free - to use either -- but every producer of a rope slice must use the SAME one. - Mixing them changes the scores silently and only shows up as lower AL. - - ``mscale`` rides in the modulus rather than the phase, exactly as the - cos/sin builder folds it into the table. + HF DeepSeek de-interleaves and rotates halves instead, writing the same + values to lanes i and i + dim/2 rather than 2i and 2i+1. That permutation + cancels inside ``q_rope . k_rope``, so either convention works -- but every + producer of a rope slice must use the SAME one, and mixing them only shows + up as lower AL. ``mscale`` rides in the modulus, not the phase. """ cos, sin = build_dspark_mla_yarn_rope(device=device, **kwargs) half = cos.shape[-1] // 2 @@ -2352,22 +2339,16 @@ def __init__(self, model_config): class _MLABlockFixup(NamedTuple): """Layer-invariant state for the MLA block-decode fixup. - The block's own latents are written into the pool at ``ctx_len + j`` first. - Those slots belong to this request -- the draft KV manager adds - ``max(draft_len, max_total_draft_tokens)`` tokens to every generation - request each step (resource_manager.py:1060) -- and the accepted tokens - overwrite them next step, so the write is transient rather than a claim on - the cache. It is what the GQA drafter's TRTLLM backend already does - (modeling_dflash.py append_paged_kv_cache). - - With ``seq_lens = ctx_len + block_size`` the kernel's in-block causality - then leaves query j seeing block keys 0..j, so all that is left to fix up is - the block's strict upper triangle -- ``blk_kv`` itself, already in hand. The - earlier form kept the pool read-only and paid for it twice: a gather of the - hidden context tail, and a 15-key rather than 8-key fixup. - - Every field depends only on ctx_len, the page table and the block geometry, - so it is built once per forward rather than once per layer. + The block's own latents are written into the pool at ``ctx_len + j``. Those + slots belong to this request -- the draft KV manager adds + ``max(draft_len, max_total_draft_tokens)`` tokens to every generation request + each step (resource_manager.py:1060) -- and the accepted tokens overwrite + them next step, so the write is transient. Same thing the GQA drafter's + TRTLLM backend does (modeling_dflash.py append_paged_kv_cache). + + In-block causality then leaves query j seeing block keys 0..j, so only the + strict upper triangle needs fixing up. Every field depends on ctx_len, the + page table and the block geometry alone, so it is built once per forward. """ pages: torch.Tensor # [B, block] page holding each block position @@ -2602,12 +2583,10 @@ def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): class MLADSparkForCausalLM(_DSparkHeadMixin, DFlashForCausalLM): """DSpark drafter on an MLA-shaped backbone, from a standalone checkpoint. - The sibling ``modeling_dspark``'s GQA class could not be: DFlash's block - decode splits a fused ``qkv_proj`` by (num_heads, num_kv_heads, head_dim) - and fuses per-head K/V across layers, and MLA has neither. What it keeps is - the DFlash *contract* -- dual-source KV (context from the projected target - hidden, block from the draft hidden), non-causal block attention, one - precomputed context cache -- and the DSpark head set on top. + Separate from the GQA class because DFlash's block decode splits a fused + ``qkv_proj`` by (num_heads, num_kv_heads, head_dim) and fuses per-head K/V + across layers, neither of which MLA has. It keeps the DFlash contract -- + dual-source KV, non-causal block attention, one precomputed context cache. The cache stores one ``kv_lora_rank + qk_rope_head_dim`` latent per token per layer and no V, which is the point: 5 x 576 x 2B = 5760 B/token/rank @@ -2745,18 +2724,12 @@ def _mla_paged_attention(self, query, blk_kv, layer_cache, fixup, block_size, ma full visibility). Its LSE is base 2 (also measured: exact against ``log2(sum exp)``, 1.84 off against the natural log). - The block's own latents go into the pool first, at ``ctx_len + j`` -- - slots the draft KV manager already reserves for this request each step + The block's own latents go into the pool first, at ``ctx_len + j`` (see _MLABlockFixup). With ``seq_lens = ctx_len + block_size`` the kernel - then covers context AND block, and in-block causality leaves only the - block's strict upper triangle: 8 keys, no context-tail gather. Those are - attended eagerly against ``blk_kv``, already in hand, and merged by - log-sum-exp. - - Writing into the pool is safe because dflash.py bounds the advertised - context at ``allocated - block_size``; without that the first block - position lands on the first unallocated page, whose table entry was - clamped to physical block 0 -- another request's. + covers context AND block, and in-block causality leaves only the block's + strict upper triangle -- 8 keys, attended eagerly against ``blk_kv`` and + merged by log-sum-exp. The room those writes need is reserved by + ``dflash_allocated_ctx_limit``. Returns the latent-space output ``[B, block, heads, kv_lora_rank]``. """ diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index a2c4b488c203..032804ae3d43 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1685,25 +1685,19 @@ def _get_draft_kv_model_config(self) -> ModelConfig: private arena, which OOMs at long context), so the pool dtype must follow the drafter, not the target. - Every consumer that reasons about draft KV bytes must go through here. - Construction and the target/draft budget split reading different dtypes - is not cosmetic: the split then charges the draft fp8 bytes for a pool - allocated bf16, so the draft manager receives exactly HALF the tokens - the target gets. Because the capacity scheduler admits on the target - pool alone, the draft pool cannot backpressure -- it can only raise - "Draft KV cache context resize failed" out of - ``KVCacheManagerV2._prepare_draft_resources``, which is fatal to every - rank, as soon as resident context passes ~50% of target utilization. + Every consumer of draft KV bytes must go through here. If the budget + split and the allocation read different dtypes, the split charges fp8 + bytes for a bf16 pool and the draft manager gets HALF the target's + tokens. The capacity scheduler admits on the target pool alone, so the + draft pool cannot backpressure -- past ~50% target utilization it raises + "Draft KV cache context resize failed", fatal to every rank. """ effective_draft_config = self._get_effective_draft_config() # Narrower than is_external_drafter(), matching - # _should_create_separate_draft_kv_cache: PARD and - # DRAFT_TARGET_ONE_MODEL reach here too, via - # should_use_separate_draft_kv_cache, and their draft checkpoints can - # carry a genuine fp8 KV algo of their own -- with dtype="auto", - # validate_and_set_kv_cache_quant returns early and keeps it. Dropping - # that would allocate bf16 while their attention modules, built from - # the un-neutralized config, still read and write fp8. + # _should_create_separate_draft_kv_cache. PARD and DRAFT_TARGET_ONE_MODEL + # reach here too and can carry a genuine fp8 KV algo of their own, which + # dtype="auto" keeps; dropping it would allocate bf16 under attention + # modules that still read and write fp8. spec_dec_mode = self._speculative_config.spec_dec_mode if not (spec_dec_mode.is_dflash() or spec_dec_mode.is_dspark()): return effective_draft_config diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 801e4abe0a4b..8f0936e0307b 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -94,23 +94,17 @@ def dflash_allocated_ctx_limit( ) -> torch.Tensor: """Context length a request may advertise, given the pages it holds. - Leaves ``block_size`` positions of room on purpose. The MLA path writes the - block's own latents at ``ctx_len .. ctx_len + block_size``, so stopping at - exactly the allocation puts the first of them on the first UNallocated page - -- whose block-table entry ``_refresh_ctx_block_tables`` clamped from a - negative placeholder to 0, i.e. another request's block. That is a silent - cross-request write, not an out-of-range fault, so the room has to be left - here rather than caught downstream. - - How tight this is depends on the V2 backend, and only one of the two keeps - the information. ``copyBatchBlockOffsetsToDeviceKernel`` - (kvCacheManagerV2Utils.cu) writes ``0`` for ``BAD_PAGE_INDEX``, so under the - default ``cpp`` backend every row counts full and this degenerates to the - block-table-width bound upstream already applied -- correct, just not - per-request. The ``python`` backend propagates ``BAD_PAGE_INDEX`` through - ``_copy_swa_block_offsets_with_scratch``, and there the bound is real. - Sized from the counts either way so the tight case needs no second code - path; do not read a per-request guarantee into it on ``cpp``. + Leaves ``block_size`` positions of room: the MLA path writes the block's own + latents at ``ctx_len .. ctx_len + block_size``, so stopping at exactly the + allocation puts the first of them on an unallocated page, whose table entry + ``_refresh_ctx_block_tables`` clamped to 0 -- another request's block. A + silent cross-request write, not a fault. + + Per-request only on the ``python`` V2 backend, which propagates + ``BAD_PAGE_INDEX``. The default ``cpp`` one maps it to 0 + (kvCacheManagerV2Utils.cu), every row counts full, and this degenerates to + the block-table-width bound upstream already applied. Do not read a + per-request guarantee into it there. """ return (block_counts * page_size - block_size).clamp(min=0) @@ -659,14 +653,10 @@ def _lazy_init_ctx_buffers( hd = draft_model._head_dim kv_factor = getattr(draft_model, "_kv_factor", 2) capacity = self._max_ctx + self._compute_block_size - # The one number a drafter that builds an absolute-position table needs, - # published from the only place that knows it: _max_ctx comes from - # attn_metadata.max_seq_len, the engine's value, which + # Published before any drafter table is built, so a drafter may read it + # instead of its config cap. _max_ctx is attn_metadata.max_seq_len, which # py_executor_creator raises past model_config.max_seq_len and never - # writes back. Tables are built lazily on first forward and this runs - # before any of them, so a drafter may read it in place of its - # config-derived cap. _compute_block_size, not _resolved_block_size: - # the block decode's j runs over the slots the forward computes. + # writes back. _compute_block_size: the block decode's j runs over it. draft_model._runtime_position_ceiling = dflash_position_ceiling( self._max_ctx, self._compute_block_size, self.max_draft_len ) @@ -1613,14 +1603,10 @@ def prepare_1st_drafter_inputs( num_ctx_per_req_t = self._ctx_len[slots] if self._ctx_block_tables is not None: - # ctx_len tracks the target sequence, but the write above clamps - # columns to what the manager allocated for this request, so a - # context that outruns its allocation has its tail written on - # top of the last valid slot. Reading past the allocation would - # then attend either that clobbered value or -- for pages the - # block table left clamped to 0 -- another request's data. - # Truncating the advertised length keeps the two consistent: - # the drafter attends a short context instead of a wrong one. + # The write above clamps columns to the request's allocation, so + # a context that outruns it has its tail written over the last + # valid slot. Truncate what is advertised to match, or the read + # attends that clobbered value -- or another request's page. allocated = dflash_allocated_ctx_limit( self._ctx_block_counts, self._ctx_page_size, block_size ) From 50039e24a11e9ff2e036387927d24ea400cc3aea Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Wed, 16 Sep 2026 02:58:19 -0700 Subject: [PATCH 09/25] [None][feat] Make the standalone drafter's block-decode backend selectable Ports AUTO/VANILLA/TRTLLM/CUTEDSL resolution from the development branch, with AUTO degrading per drafter family. CUTEDSL is gated off here: it needs a cute_dsl_mla_decode_fp16_blackwell taking per-token kv_bounds, absent upstream. Signed-off-by: Zhenhuan Chen --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 18 +- .../dspark_rmsnorm_rope_custom_op.py | 200 +++++++++- .../blackwell/dspark_rmsnorm_rope.py | 51 ++- tensorrt_llm/_torch/models/modeling_dflash.py | 47 ++- tensorrt_llm/_torch/models/modeling_dspark.py | 377 +++++++++++++++--- .../_torch/models/modeling_speculative.py | 2 +- tensorrt_llm/_torch/speculative/dflash.py | 6 + .../_torch/speculative/dflash_attention.py | 11 +- tensorrt_llm/llmapi/llm_args.py | 30 +- .../test_kimi_k3_dspark_semantics.py | 359 +++++++++++++---- .../test_dspark_cute_dsl_rmsnorm_rope.py | 37 +- 11 files changed, 962 insertions(+), 176 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 0e02a8759ab4..91a1053e525a 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -11559,12 +11559,28 @@ def forward( compiled_mla = CuteDSLNVMlaDecodeBlackwellRunner.kernel_cache[ cache_key] + page_table_arg = page_table + if page_table.shape[0] == 1 and page_table.shape[1] == 1: + # A raw torch.Tensor is re-adapted at call time by TensorAdapter, + # which runs a bare `from_dlpack(arg).mark_layout_dynamic()` + # (cute/runtime.py:915, registered for torch.Tensor at :941), so + # the leading_dim=0 given to cute.compile above does not carry + # over. A (1, 1) page table -- one page, batch 1 -- has stride + # (1, 1), no dimension with size > 1, and the deduction then + # raises "Can't deduce the leading dimension from layout" even + # though with both extents 1 the choice cannot change an address. + # An already-marked tensor has no registered adapter + # (jit_executor.py:650), so the bare call is skipped. Only the + # degenerate shape pays the wrapper. + page_table_arg = cute.runtime.from_dlpack( + page_table, + assumed_align=16).mark_layout_dynamic(leading_dim=0) runtime_args = [ q_latent, q_rope, c_latent, c_rope, - page_table, + page_table_arg, o, lse, ] diff --git a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py index 98ed0d4e3b75..56bc9f62d707 100644 --- a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py @@ -35,18 +35,44 @@ def _get_dspark_arch_str(sm_version: int | None = None) -> str | None: return _DSV4_DSPARK_ARCH_BY_SM.get(sm_version) +def _has_regular_row_stride(x: torch.Tensor) -> bool: + """Return whether leading dimensions flatten to non-overlapping strided rows.""" + if x.is_contiguous(): + return True + if x.stride(-1) != 1: + return False + + outer_dims = [dim for dim in range(x.ndim - 1) if x.shape[dim] > 1] + if not outer_dims: + return True + + row_dim = outer_dims[-1] + if x.stride(row_dim) < x.shape[-1]: + return False + return all( + x.stride(dim) == x.stride(next_dim) * x.shape[next_dim] + for dim, next_dim in zip(outer_dims, outer_dims[1:]) + ) + + def is_fused_dspark_rmsnorm_rope_supported( x: torch.Tensor, - weight: torch.Tensor, + weight: torch.Tensor | None, freqs: torch.Tensor, num_heads: int, rope_dim: int, norm_dim: int | None = None, ) -> bool: - """Return whether tensors satisfy the production fused-op contract.""" - if _get_dspark_arch_str() is None or not all(t.is_cuda for t in (x, weight, freqs)): + """Return whether tensors satisfy the production fused-op contract. + + ``weight`` is None for a kernel built with ``apply_weight=False``; the + weight checks are then vacuous rather than a reason to reject, and the + caller has no tensor to offer in the first place. + """ + operands = (x, freqs) if weight is None else (x, weight, freqs) + if _get_dspark_arch_str() is None or not all(t.is_cuda for t in operands): return False - if x.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16: + if x.dtype != torch.bfloat16 or (weight is not None and weight.dtype != torch.bfloat16): return False if freqs.dtype != torch.float32: return False @@ -59,7 +85,9 @@ def is_fused_dspark_rmsnorm_rope_supported( effective_norm_dim = x.shape[-1] if norm_dim is None else norm_dim if effective_norm_dim not in (x.shape[-1], x.shape[-1] - rope_dim): return False - if effective_norm_dim % 32 != 0 or weight.shape != (effective_norm_dim,): + if effective_norm_dim % 32 != 0: + return False + if weight is not None and weight.shape != (effective_norm_dim,): return False if (x.shape[-1] - rope_dim) % 32 != 0 or (rope_dim // 2) % 32 != 0: return False @@ -71,8 +99,8 @@ def is_fused_dspark_rmsnorm_rope_supported( and freqs.shape[0] == rows // num_heads and freqs.shape[1] >= max(1, rope_dim // 2) and freqs.shape[2] == 2 - and x.is_contiguous() - and weight.is_contiguous() + and _has_regular_row_stride(x) + and (weight is None or weight.is_contiguous()) and freqs.is_contiguous() ) @@ -157,8 +185,13 @@ def _compile_fused_dspark_rmsnorm_rope( x_fake = cute.runtime.make_fake_compact_tensor( cutlass.BFloat16, (rows, hidden_dim), stride_order=(1, 0) ) - weight_fake = cute.runtime.make_fake_compact_tensor( - cutlass.BFloat16, (norm_dim,), stride_order=(0,) + # None, not a fake tensor, when the kernel does not scale by a weight: the + # operand then does not exist in the compiled signature, so the call site + # has nothing to pass and nothing to allocate. + weight_fake = ( + cute.runtime.make_fake_compact_tensor(cutlass.BFloat16, (norm_dim,), stride_order=(0,)) + if apply_weight + else None ) freqs_fake = cute.runtime.make_fake_compact_tensor( cutlass.Float32, @@ -190,6 +223,64 @@ def _compile_fused_dspark_rmsnorm_rope( ) +@functools.cache +def _compile_fused_dspark_rope_into( + hidden_dim: int, + rope_dim: int, + num_heads: int, + eps: float, + out_dim: int, + out_rope_offset: int, +): + """Rope-only variant that writes into a wider destination row. + + Every argument is part of the @functools.cache key on purpose: out_dim and + out_rope_offset change the generated addressing, so sharing a compiled + kernel across two offsets would write the rotated pairs to the wrong + columns -- wrong numbers, no error, and for a speculative drafter that only + shows up as a lower acceptance length. + """ + rows = cute.sym_int() + freq_rows = cute.sym_int() + x_fake = cute.runtime.make_fake_tensor( + cutlass.BFloat16, + (rows, hidden_dim), + stride=(cute.sym_int64(), 1), + ) + freqs_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (freq_rows, cute.sym_int(), 2), + stride_order=(2, 1, 0), + ) + output_fake = cute.runtime.make_fake_tensor( + cutlass.BFloat16, + (rows, out_dim), + stride=(cute.sym_int64(), 1), + ) + stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = DSparkRMSNormRoPEKernel( + hidden_dim, + rope_dim, + num_heads, + eps, + False, + False, + False, + norm_dim=hidden_dim - rope_dim, + out_rope_offset=out_rope_offset, + write_nope=False, + ) + return cute.compile( + kernel, + x_fake, + None, + freqs_fake, + output_fake, + stream_fake, + options="--opt-level 2 --enable-tvm-ffi", + ) + + @functools.cache def _compile_dspark_rmsnorm_rope_cache_write(eps: float): if torch.cuda.is_current_stream_capturing(): @@ -296,7 +387,7 @@ def _compile_dspark_rmsnorm_rope_draft_block(block_size: int, eps: float): ) def cute_dsl_dspark_rmsnorm_rope( x: torch.Tensor, - weight: torch.Tensor, + weight: torch.Tensor | None, freqs: torch.Tensor, num_heads: int, rope_dim: int, @@ -306,11 +397,19 @@ def cute_dsl_dspark_rmsnorm_rope( inverse_rope: bool, norm_dim: int | None = None, ) -> torch.Tensor: - """Apply fused RMSNorm and adjacent-pair RoPE to contiguous BF16 rows.""" + """Apply fused RMSNorm and adjacent-pair RoPE to regular BF16 rows. + + ``weight`` may be None when ``apply_weight`` is False. It is dropped on the + way to the kernel either way, so a caller that has a weight lying around + can keep passing it, and one that does not need not invent one. + """ + if apply_weight and weight is None: + raise ValueError("cute_dsl_dspark_rmsnorm_rope needs a weight when apply_weight is set") + weight = weight if apply_weight else None if not is_fused_dspark_rmsnorm_rope_supported(x, weight, freqs, num_heads, rope_dim, norm_dim): raise ValueError( - "cute_dsl_dspark_rmsnorm_rope requires contiguous BF16 tensors on " - "an SM100 or SM103 GPU with a valid FP32 frequency view; " + "cute_dsl_dspark_rmsnorm_rope requires regular row-strided BF16 tensors on " + "an SM100, SM103, or SM107 GPU with a valid FP32 frequency view; " f"got SM {get_sm_version()}" ) @@ -334,7 +433,7 @@ def cute_dsl_dspark_rmsnorm_rope( @torch.library.register_fake("trtllm::cute_dsl_dspark_rmsnorm_rope") def _( x: torch.Tensor, - weight: torch.Tensor, + weight: torch.Tensor | None, freqs: torch.Tensor, num_heads: int, rope_dim: int, @@ -347,6 +446,79 @@ def _( return torch.empty_like(x) +@torch.library.custom_op( + "trtllm::cute_dsl_dspark_rope_into", + mutates_args=("out",), + device_types="cuda", +) +def cute_dsl_dspark_rope_into( + x: torch.Tensor, + freqs: torch.Tensor, + out: torch.Tensor, + num_heads: int, + rope_dim: int, + out_rope_offset: int, +) -> None: + """Rotate x's trailing rope_dim and store it at out[..., out_rope_offset:]. + + x keeps its own (narrower) row width; nothing is written outside the rope + columns of `out`, so the caller owns the rest of the destination row. + """ + # No weight operand: _compile_fused_dspark_rope_into builds the kernel with + # apply_weight=False, so there is nothing for the call to scale by and + # nothing to allocate per decode step. + if not is_fused_dspark_rmsnorm_rope_supported( + x, None, freqs, num_heads, rope_dim, x.shape[-1] - rope_dim + ): + raise ValueError( + "cute_dsl_dspark_rope_into requires regular row-strided BF16 tensors on " + "an SM100, SM103, or SM107 GPU with a valid FP32 frequency view; " + f"got SM {get_sm_version()}" + ) + # The destination is compiled with x's dtype and a (dyn, 1) stride, so it + # needs the same checks as the source; view() below would otherwise fail + # with an opaque RuntimeError, and a mismatched dtype would reach a kernel + # compiled for the other one. + if out.dtype != x.dtype or out.device != x.device or not _has_regular_row_stride(out): + raise ValueError( + "cute_dsl_dspark_rope_into needs a row-strided out on x's device with " + f"x's dtype; got dtype={out.dtype}, device={out.device}, " + f"stride={tuple(out.stride())}" + ) + out_flat = out.view(-1, out.shape[-1]) + x_flat = x.view(-1, x.shape[-1]) + if out_flat.shape[0] != x_flat.shape[0]: + raise ValueError( + f"out must have one row per x row; got {out_flat.shape[0]} vs {x_flat.shape[0]}" + ) + if out_rope_offset + rope_dim > out.shape[-1]: + raise ValueError( + f"rope slice [{out_rope_offset}, {out_rope_offset + rope_dim}) does not fit " + f"a row of width {out.shape[-1]}" + ) + compiled = _compile_fused_dspark_rope_into( + x.shape[-1], + rope_dim, + num_heads, + 0.0, + out.shape[-1], + out_rope_offset, + ) + compiled(x_flat, None, freqs, out_flat) + + +@torch.library.register_fake("trtllm::cute_dsl_dspark_rope_into") +def _( + x: torch.Tensor, + freqs: torch.Tensor, + out: torch.Tensor, + num_heads: int, + rope_dim: int, + out_rope_offset: int, +) -> None: + return None + + @torch.library.custom_op( "trtllm::cute_dsl_dspark_rmsnorm_rope_cache_write", mutates_args=("kv_cache",), diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark_rmsnorm_rope.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark_rmsnorm_rope.py index f14c51aa3650..bfdaa933119a 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark_rmsnorm_rope.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark_rmsnorm_rope.py @@ -3,6 +3,8 @@ """Fused DSpark RMSNorm and adjacent-pair RoPE for Blackwell.""" +from typing import Optional + import cutlass import cutlass.cute as cute @@ -27,6 +29,8 @@ def __init__( apply_rmsnorm: bool, inverse_rope: bool, norm_dim: int | None = None, + out_rope_offset: int | None = None, + write_nope: bool = True, ): if hidden_dim % self.num_threads != 0: raise ValueError( @@ -56,6 +60,22 @@ def __init__( f"({self.nope_dim}); got {self.norm_dim}" ) self.norm_covers_rope = self.norm_dim == hidden_dim + # Where the rotated pairs land in `output`. None means "same column as in + # the input", which is every caller that rewrites a row in place. The MLA + # drafter instead reads a 192-wide q row and writes into the rope half of + # a 576-wide fused query, so its input and output offsets differ and it + # wants the nope half left to bmm_out. + self.out_rope_offset = self.nope_dim if out_rope_offset is None else out_rope_offset + if self.out_rope_offset < 0 or self.out_rope_offset % 2 != 0: + raise ValueError( + f"out_rope_offset must be even and non-negative; got {self.out_rope_offset}" + ) + self.write_nope = write_nope + if not self.write_nope and self.norm_covers_rope: + # The rope lanes would need the RMS scale computed over a row whose + # nope half is never written -- readable, but no caller wants it and + # it would silently pair a raw-passthrough offset with a scaled rope. + raise ValueError("write_nope=False requires norm_dim == nope_dim") if self.norm_dim % self.num_threads != 0: raise ValueError( f"norm_dim must be divisible by {self.num_threads}; got {self.norm_dim}" @@ -77,11 +97,19 @@ def __init__( def __call__( self, x: cute.Tensor, - weight: cute.Tensor, + weight: Optional[cute.Tensor], freqs: cute.Tensor, output: cute.Tensor, stream: cuda.CUstream, ): + """weight is None exactly when apply_weight is False. + + Both reads sit behind const_expr(self.apply_weight), so with the flag + off the operand never reaches the generated signature -- which is the + point: a rope-only caller has no weight and should not have to + materialize one. The compile wrappers derive both from the same flag, + so they cannot disagree. + """ self.kernel(x, weight, freqs, output).launch( grid=[x.shape[0], 1, 1], block=[self.num_threads, 1, 1], @@ -92,7 +120,7 @@ def __call__( def kernel( self, x: cute.Tensor, - weight: cute.Tensor, + weight: Optional[cute.Tensor], freqs: cute.Tensor, output: cute.Tensor, ): @@ -109,12 +137,13 @@ def kernel( sum_sq = cute.arch.warp_reduction_sum(sum_sq) inverse_rms = cute.math.rsqrt(sum_sq / self.norm_dim + self.eps) - for item in cutlass.range_constexpr(self.nope_elements_per_thread): - dim = tidx + item * self.num_threads - value = cutlass.Float32(x[row, dim]) * inverse_rms - if cutlass.const_expr(self.apply_weight): - value *= cutlass.Float32(weight[dim]) - output[row, dim] = value.to(output.element_type) + if cutlass.const_expr(self.write_nope): + for item in cutlass.range_constexpr(self.nope_elements_per_thread): + dim = tidx + item * self.num_threads + value = cutlass.Float32(x[row, dim]) * inverse_rms + if cutlass.const_expr(self.apply_weight): + value *= cutlass.Float32(weight[dim]) + output[row, dim] = value.to(output.element_type) if cutlass.const_expr(self.rope_pairs > 0): freq_row = row // self.num_heads @@ -122,6 +151,8 @@ def kernel( pair = tidx + item * self.num_threads real_dim = self.nope_dim + pair * 2 imag_dim = real_dim + 1 + out_real = self.out_rope_offset + pair * 2 + out_imag = out_real + 1 real = cutlass.Float32(x[row, real_dim]) imag = cutlass.Float32(x[row, imag_dim]) # Outside norm_dim the rope lanes are passed through raw: no RMS @@ -136,8 +167,8 @@ def kernel( sin = cutlass.Float32(freqs[freq_row, pair, 1]) if cutlass.const_expr(self.inverse_rope): sin = -sin - output[row, real_dim] = (real * cos - imag * sin).to(output.element_type) - output[row, imag_dim] = (imag * cos + real * sin).to(output.element_type) + output[row, out_real] = (real * cos - imag * sin).to(output.element_type) + output[row, out_imag] = (imag * cos + real * sin).to(output.element_type) class DSparkRMSNormRoPECacheWriteKernel(DSparkRMSNormRoPEKernel): diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index fb8037eb53c1..c49919541696 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -20,6 +20,7 @@ _flashinfer_rope = None from ..pyexecutor.config_utils import _is_sliding_attention_layer, get_layer_attention_window from ..speculative.dflash_attention import ( + dflash_trtllm_gen_unavailability_reason, get_dflash_fa4_fwd, get_dflash_flash_attention, get_dflash_paged_append, @@ -329,6 +330,14 @@ class DFlashForCausalLM(nn.Module): # dependencies, and the worker's per-backend shape checks do not apply. _uses_worker_attention_backend = True + # What ``attention_backend="AUTO"`` resolves to, and what the field may say + # at all. Per drafter family, because the fastest kernel that can express + # the shape differs: GQA DFlash cross-attention is what the two worker op + # sets were built for, while an MLA drafter runs its own block decode and + # has a third implementation they cannot express. + _default_attention_backend = "VANILLA" + _supported_attention_backends = ("VANILLA", "TRTLLM", "FA4") + # Whether the drafter's context KV lives in the draft KV cache manager's # paged pool rather than a private arena dense in max_seq_len. Orthogonal # to the attention backend: paging is about where the KV lives, the backend @@ -342,7 +351,30 @@ class DFlashForCausalLM(nn.Module): # tests), so the config cap stands. _runtime_position_ceiling = None - def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): + @classmethod + def _resolve_auto_attention_backend(cls) -> str: + """Turn ``AUTO`` into a concrete backend for this drafter family. + + Only TRTLLM degrades: its ops are an optional dependency and the two + drafter families that prefer it differ in whether a build can load + them. CUTEDSL does not -- a build that cannot run it raises in + ``MLADSparkForCausalLM.__init__`` with the reason, which is the + intended behaviour for the only hardware that lacks it. + """ + want = cls._default_attention_backend + if want != "TRTLLM": + return want + reason = dflash_trtllm_gen_unavailability_reason() + if reason is None: + return "TRTLLM" + logger.info_once( + f"{cls.__name__} prefers the TRTLLM attention backend but it is " + f"unavailable ({reason}); falling back to VANILLA.", + key=f"dflash_auto_backend_fallback_{cls.__name__}", + ) + return "VANILLA" + + def __init__(self, draft_config, *, dflash_attention_backend: str = "AUTO"): """Build the draft model, resolving its architecture from the draft config (falling back to a model_type-derived name when the checkpoint uses a custom DFlash architecture label).""" @@ -389,13 +421,17 @@ def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): ) self.target_layer_ids = dflash_config.get("target_layer_ids", None) + # Upstream keeps the dflash_config override; rubin-advance dropped it. self.block_size = dflash_config.get( "block_size", getattr(pretrained_config, "block_size", None) ) + if dflash_attention_backend == "AUTO": + dflash_attention_backend = self._resolve_auto_attention_backend() self.dflash_attention_backend = dflash_attention_backend - if self.dflash_attention_backend not in ("VANILLA", "TRTLLM", "FA4"): + if self.dflash_attention_backend not in self._supported_attention_backends: raise ValueError( - "DFlash attention backend must be VANILLA, TRTLLM or FA4, got " + f"{type(self).__name__} attention backend must be one of " + f"{list(self._supported_attention_backends)}, got " f"{self.dflash_attention_backend!r}." ) # Each backend loads only its own ops; the rest stay None so the @@ -409,7 +445,8 @@ def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): # but no op set is loaded: this drafter calls none of them. logger.info_once( f"{type(self).__name__} brings its own block decode; " - f"attention_backend={self.dflash_attention_backend!r} is not used.", + f"attention_backend={self.dflash_attention_backend!r} selects " + f"among its implementations, not the shared DFlash op sets.", key=f"dflash_own_attention_{type(self).__name__}", ) elif self.dflash_attention_backend == "VANILLA": @@ -1826,7 +1863,7 @@ def _normalize_config(config: PretrainedConfig) -> None: if isinstance(dflash_config, dict): config.block_size = dflash_config.get("block_size", None) - def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): + def __init__(self, draft_config, *, dflash_attention_backend: str = "AUTO"): """Pin the Laguna draft-layer class and enable Laguna-specific behaviors (context input_layernorm, causal sliding blocks); reject non-per-head gating.""" diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 8ed54b704b75..42f414618a83 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -95,7 +95,7 @@ class or the edge would become a cycle. from tensorrt_llm.mapping import Mapping from tensorrt_llm.quantization.mode import QuantAlgo -from ..._utils import is_sm_100f +from ..._utils import get_sm_version, is_sm_100f from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from ..distributed import AllReduceParams from ..model_config import ModelConfig @@ -384,9 +384,35 @@ def _rope_last_dims_batched( return torch.cat([nope, rope], dim=-1) +def _mla_rope_into( + q: torch.Tensor, + freqs_cis: torch.Tensor, + out: torch.Tensor, + num_heads: int, + rope_head_dim: int, + out_rope_offset: int, +) -> bool: + """Rotate q's rope tail into ``out[..., out_rope_offset:]``; False if unavailable. + + Returns a bool rather than raising so the caller keeps the eager path for + builds without the fused kernel, the same way _rmsnorm_rope_batched does. + """ + if not (IS_CUTLASS_DSL_AVAILABLE and is_sm_100f()): + return False + freqs_real = torch.view_as_real(freqs_cis).reshape(-1, freqs_cis.shape[-1], 2) + if not is_fused_dspark_rmsnorm_rope_supported( + q, None, freqs_real, num_heads, rope_head_dim, q.shape[-1] - rope_head_dim + ): + return False + torch.ops.trtllm.cute_dsl_dspark_rope_into( + q, freqs_real, out, num_heads, rope_head_dim, out_rope_offset + ) + return True + + def _rmsnorm_rope_batched( t: torch.Tensor, - weight: torch.Tensor, + weight: Optional[torch.Tensor], eps: float, rope_head_dim: int, freqs_cis: torch.Tensor, @@ -399,6 +425,9 @@ def _rmsnorm_rope_batched( ) -> torch.Tensor: """Fuse DSpark RMSNorm and last-dimension RoPE when supported. + ``weight`` may be None when ``apply_weight`` is False -- no path reads it + then, fused or eager. + ``norm_dim`` bounds the RMSNorm (and the weight) to a prefix of the row. None means the whole row, which is DSpark's own convention. The other supported value is ``t.shape[-1] - rope_head_dim``: DeepSeek-style MLA @@ -604,7 +633,7 @@ def dspark_attention_forward_batched( q = F.linear(q, wq_b).unflatten(-1, (n_heads, head_dim)) # [G, block, h, head_dim] q = _rmsnorm_rope_batched( q, - kv_norm_w, + None, eps, rd, blk_freqs, @@ -2356,6 +2385,44 @@ class _MLABlockFixup(NamedTuple): valid: torch.Tensor # [1, block, 1, block] strict upper triangle page_tables_i32: torch.Tensor # kernel operands, cast once seq_lens_i32: torch.Tensor # ctx_len + block_size + kv_bounds_i32: torch.Tensor # [B * block] seq_lens repeated per query token + + +@lru_cache(maxsize=1) +def cute_dsl_mla_decode_unavailability_reason() -> Optional[str]: + """Return why the cute-dsl MLA decode cannot run here, or None. + + Mirrors ``dflash_trtllm_gen_unavailability_reason``: the op itself only + reports by raising, and ``attention_backend="CUTEDSL"`` has to fail at + construction rather than five layers into the first draft step. + + The ``kv_bounds`` probe is not paranoia. The op exists upstream without + per-token KV bounds, and this backend's whole shape -- one pass over + context and block, no fixup -- is expressed through them. Without the probe + the mismatch surfaces as a torch dispatch TypeError at the first draft + step, naming an argument count rather than a missing kernel feature. + """ + if (sm := get_sm_version()) not in (100, 103, 107): + return f"requires SM100, SM103 or SM107, got SM{sm}" + try: + import cutlass # noqa: F401 + + from ..custom_ops.cute_dsl_custom_ops import CuteDSLNVMlaDecodeBlackwellRunner # noqa: F401 + except ImportError as error: + return f"the cute-dsl kernels are not importable: {error}" + op = getattr(torch.ops.trtllm, "cute_dsl_mla_decode_fp16_blackwell", None) + if op is None: + return "trtllm::cute_dsl_mla_decode_fp16_blackwell is not registered" + if not any( + arg.name == "kv_bounds" + for overload in op.overloads() + for arg in getattr(op, overload)._schema.arguments + ): + return ( + "the registered cute_dsl_mla_decode_fp16_blackwell takes no kv_bounds, " + "so per-token KV bounds are unavailable in this build" + ) + return None def _build_mla_block_fixup(ctx_len, page_tables, block_size, page_size) -> _MLABlockFixup: @@ -2363,6 +2430,7 @@ def _build_mla_block_fixup(ctx_len, page_tables, block_size, page_size) -> _MLAB device = ctx_len.device pos = ctx_len.view(-1, 1) + torch.arange(block_size, device=device) idx = torch.arange(block_size, device=device) + seq_lens = (ctx_len + block_size).to(torch.int32) return _MLABlockFixup( pages=page_tables.gather(1, pos // page_size), offsets=pos % page_size, @@ -2370,7 +2438,11 @@ def _build_mla_block_fixup(ctx_len, page_tables, block_size, page_size) -> _MLAB 1, block_size, 1, block_size ), page_tables_i32=page_tables.to(torch.int32), - seq_lens_i32=(ctx_len + block_size).to(torch.int32), + seq_lens_i32=seq_lens, + # repeat_interleave, not repeat: the kernel indexes kv_bounds at + # b * seq_len_q + q (mla_decode_fp16.py:400-407), so the batch must be + # the slow axis. B == 1 cannot tell the two apart. + kv_bounds_i32=seq_lens.repeat_interleave(block_size), ) @@ -2570,7 +2642,19 @@ class GQADSparkForCausalLM(_DSparkHeadMixin, DFlashForCausalLM): Reference: arXiv 2607.05147; deepseek-ai/DeepSpec. """ - def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): + # TRTLLM, matching what the shipped K3 configs already pin + # (examples/kimi_k3/eval_extra_llm_options_dspark{,_nvfp4}.yaml); VANILLA + # additionally needs the optional flash-attn package. AUTO degrades to + # VANILLA off SM100/SM103. + # + # This also moves the drafter's context KV into the manager's paged pool, + # and that is the point, not a side effect: dflash.py:592 keys `use_paged` + # off the backend name, and this class leaves _paged_ctx_cache False, so + # VANILLA is the only way to get the arena that is dense in max_seq_len and + # that free_gpu_memory_fraction never bounds. + _default_attention_backend = "TRTLLM" + + def __init__(self, draft_config, *, dflash_attention_backend: str = "AUTO"): super().__init__(draft_config, dflash_attention_backend=dflash_attention_backend) self._init_dspark_heads(draft_config.pretrained_config) @@ -2593,27 +2677,44 @@ class MLADSparkForCausalLM(_DSparkHeadMixin, DFlashForCausalLM): against the GQA drafter's 20480 under attention-DP, where KV heads are unsharded. - Block attention runs on flashinfer's trtllm-gen absorbed-MLA paged decode - (``_mla_paged_attention``); the eager torch path stays only as the - reference, for builds without flashinfer and for the unpaged arena. - ``spec_config.attention_backend`` does NOT select between those two -- it - picks the *GQA* drafter's backend in the shared DFlash worker, and this - class overrides the attention entirely. + ``spec_config.attention_backend`` selects among this class's own three + block-decode implementations, not among the shared DFlash worker op sets: + + ``CUTEDSL`` one cute-dsl pass with ``kv_bounds``, no fixup + (``_cute_dsl_block_decode``). Paged only. What ``AUTO`` picks + on a build that can run it. + ``TRTLLM`` flashinfer's trtllm-gen absorbed-MLA paged decode plus the + upper-triangle fixup (``_mla_paged_attention``). + ``VANILLA`` the eager torch reference. It is also what the other two + degrade to on the unpaged arena, where there are no page + tables to hand a kernel. Reference: arXiv 2607.05147; deepseek-ai/DeepSpec. """ - # Attention is _mla_paged_attention, so neither worker backend applies and + # Attention is _mla_paged_attention, so neither worker OP SET applies and # neither backend's shape checks bind -- the absorbed shape (64 heads : 1 KV - # head, head_dim 576) would fail both. + # head, head_dim 576) would fail both. The backend NAME still selects, via + # _mla_block_decode_variant below. _uses_worker_attention_backend = False + # CUTEDSL stays selectable but is NOT the default here: it needs a + # cute_dsl_mla_decode_fp16_blackwell that takes per-token kv_bounds, which + # this branch's kernel does not have. Flip this one line once it does -- + # measured 35.7% off the drafter region (947.50 -> 608.95 us/iteration, + # jobs 3006025/3006026) at no cost in acceptance or accuracy, on both + # aggregated and disaggregated gsm8k (jobs 3006856/3006857, 3006673/3006684). + # CUTEDSL never degrades: an explicit request raises in __init__ with the + # reason rather than hiding a slower path behind the name. + _default_attention_backend = "TRTLLM" + _supported_attention_backends = ("VANILLA", "TRTLLM", "CUTEDSL") + # The whole point of the MLA drafter: its context KV comes out of the # manager's pool (5760 B/token/rank) instead of an arena dense in # max_seq_len that free_gpu_memory_fraction never bounds. Independent of the # backend field, which is why either value of it works here. _paged_ctx_cache = True - def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): + def __init__(self, draft_config, *, dflash_attention_backend: str = "AUTO"): cfg = draft_config.pretrained_config # Pin the backbone instead of relying on the model_type-derived name # (the Laguna precedent in modeling_dflash.py): the checkpoint labels @@ -2646,7 +2747,7 @@ def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): self._mla_workspace = None self._mla_freqs = None self._mla_freqs_cap = None - self._mla_noop_weight = None + self._cute_dsl_checked_shape = None # HF DeepSeek folds YaRN's attention scaling into the softmax scale: # softmax_scale = mscale^2 / sqrt(qk_nope + qk_rope). NOT 1/sqrt(576) -- # the absorbed product reproduces the 192-dim un-absorbed score. @@ -2654,6 +2755,16 @@ def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): self.softmax_scale = (mscale * mscale) / math.sqrt( self.qk_nope_head_dim + self.qk_rope_head_dim ) + # Raise here rather than five layers into the first draft step, and name + # the reason. Reached for an explicit CUTEDSL and for the AUTO default + # alike -- neither degrades. + if self.dflash_attention_backend == "CUTEDSL": + reason = cute_dsl_mla_decode_unavailability_reason() + if reason is not None: + raise ValueError( + f"attention_backend='CUTEDSL' was requested but the cute-dsl " + f"MLA decode is unavailable: {reason}." + ) # -- shape / buffers --------------------------------------------------- @@ -2703,19 +2814,150 @@ def _mla_decode_op(): except (ImportError, AttributeError): return None - def _mla_rope_noop_weight(self, like: torch.Tensor) -> torch.Tensor: - """Placeholder weight for a rope-only call; never read by the kernel.""" - if self._mla_noop_weight is None: - self._mla_noop_weight = torch.ones(like.shape[-1], device=like.device, dtype=like.dtype) - return self._mla_noop_weight + def _mla_block_decode_variant(self, paged: bool) -> str: + """Which of the three block-decode implementations this step runs. - def _mla_decode_workspace(self, device): - if self._mla_workspace is None: - self._mla_workspace = torch.zeros( - _MLA_DECODE_WORKSPACE_BYTES, dtype=torch.uint8, device=device - ) + ``paged`` is False on the private arena, where there are no page tables + for either kernel, so both kernel backends degrade to eager there. + """ + if not paged or self.dflash_attention_backend == "VANILLA": + return "eager" + if self.dflash_attention_backend == "CUTEDSL": + return "cute_dsl" + return "trtllm_gen" if self._mla_decode_op() is not None else "eager" + + def _mla_decode_workspace(self, device, min_bytes: int = 0): + want = max(_MLA_DECODE_WORKSPACE_BYTES, int(min_bytes)) + if self._mla_workspace is None or self._mla_workspace.numel() < want: + self._mla_workspace = torch.zeros(want, dtype=torch.uint8, device=device) return self._mla_workspace + def _assert_cute_dsl_can_implement(self, batch, block_size, page_size, dtype) -> None: + """Reject a shape the kernel cannot serve, here, naming the reason. + + Without this the rejection is silent until it is not: ``get_valid_tactics`` + returns [] (cute_dsl_custom_ops.py:8508), the AutoTuner falls back to its + -1 sentinel, and ``forward`` compiles an unvalidated config -- so a draft + pool with ``tokens_per_block=256`` (``128 % 256 != 0``, + mla_decode_fp16.py:3798) surfaces as a CUTLASS error out of the fifth + drafter layer rather than as the configuration problem it is. + + Page size is the only operand here that is not fixed by the checkpoint, + and it comes from the draft KV pool, so this cannot move to __init__. + """ + key = (batch, block_size, page_size, dtype) + if self._cute_dsl_checked_shape == key: + return + import cutlass + + from ..custom_ops.cute_dsl_custom_ops import CuteDSLNVMlaDecodeBlackwellRunner + + in_dtype = cutlass.Float16 if dtype == torch.float16 else cutlass.BFloat16 + runner = CuteDSLNVMlaDecodeBlackwellRunner + # Mirror the runner's own probe (cute_dsl_custom_ops.py:8542) through its + # own constants, so this cannot drift from the tactic it will pick. Its + # candidate tiler list has one entry; split_kv=1 is always in range. + ok = runner._KERNEL_CLASS_BY_DTYPE[in_dtype].can_implement( + batch, + block_size, + page_size, + self._num_heads, + self.kv_lora_rank, + self.qk_rope_head_dim, + in_dtype, + in_dtype, + cutlass.Float32, + cutlass.Float32, + (128, 128), + (128, 256), + 1, + False, + runner._IS_VAR_SEQ, + runner._IS_VAR_SPLIT_KV, + page_size, + ) + if not ok: + raise ValueError( + f"attention_backend='CUTEDSL' cannot serve this shape: batch={batch}, " + f"block_size={block_size}, page_size={page_size}, " + f"num_heads={self._num_heads}, dtype={dtype}. Use " + f"attention_backend='TRTLLM', or set kv_cache_config.tokens_per_block " + f"to a divisor of 128 other than 1." + ) + self._cute_dsl_checked_shape = key + + def _cute_dsl_block_decode(self, query, layer_cache, fixup, block_size): + """One cute-dsl MLA decode over context AND block, no fixup. + + The trtllm-gen path has to split the work: its q_len semantics are + causal, so query i cannot see block positions j > i, and the strict + upper triangle is recovered eagerly afterwards and merged by LSE. The + cute-dsl kernel takes ``kv_bounds`` -- a per-query-token KV length that + REPLACES the implicit causal bound (mla_decode_fp16.py:400-407) -- so + filling it with ctx_len + block_size makes every block position visible + to every other in one pass and deletes the fixup outright. + + No fallback: what the kernel cannot serve raises. Build-level + availability is checked at construction through + ``cute_dsl_mla_decode_unavailability_reason``; the shape depends on the + draft pool's page size, which is not known until here, so it is checked + on the first call. + """ + import cutlass + + from ..custom_ops.cute_dsl_custom_ops import CuteDSLNVMlaDecodeBlackwellRunner + + batch, _, num_heads, _ = query.shape + device = query.device + pool = layer_cache[:, 0, 0] # [pages, page_size, head_dim] + page_size = pool.shape[1] + d_latent = self.kv_lora_rank + self._assert_cute_dsl_can_implement(batch, block_size, page_size, query.dtype) + + # stride[1] == 1 is the kernel's only layout demand (mla_decode_fp16.py:465-470); + # a 576-wide row sliced at 512 keeps that on both halves, so no copy. + # detach(): cute.runtime.from_dlpack refuses a tensor that requires grad + # ("Can't export tensors that require gradient"). Production runs under + # inference_mode so nothing here carries grad, but the in-place pool + # write above pulls the pool into the autograd graph whenever a caller + # does not -- a unit test, for one. Views, so no copy. + q_latent = query.detach()[..., :d_latent].permute(2, 3, 1, 0) + q_rope = query.detach()[..., d_latent:].permute(2, 3, 1, 0) + c_latent = pool.detach()[..., :d_latent].permute(1, 2, 0) + c_rope = pool.detach()[..., d_latent:].permute(1, 2, 0) + + out = torch.empty(batch, block_size, num_heads, d_latent, dtype=query.dtype, device=device) + workspace_numel = CuteDSLNVMlaDecodeBlackwellRunner.get_max_padded_workspace_size( + num_heads, block_size, d_latent, batch, cutlass.Float32 + ) + workspace = self._mla_decode_workspace(device, workspace_numel) + torch.ops.trtllm.cute_dsl_mla_decode_fp16_blackwell( + q_latent, + q_rope, + c_latent, + c_rope, + # .t() WITHOUT .contiguous(): the kernel wants (max_blocks, B) with + # dim 0 as the leading (stride-1) dimension. contiguous() would make + # it row-major, i.e. stride (B, 1), and the kernel rejects it with + # "Expected strides[leading_dim] == 1, but got B" for every B > 1 + # (probe job 3005553; B=1 passes by coincidence). The transposed + # view of a row-major (B, max_blocks) already has stride (1, max_blocks). + fixup.page_tables_i32.t(), + fixup.seq_lens_i32, + out.permute(2, 3, 1, 0), + workspace, + num_heads, + block_size, + page_size, + self.softmax_scale, + 1.0, + batch, + None, + # Every block token sees the whole block, not just its causal prefix. + fixup.kv_bounds_i32, + ) + return out + def _mla_paged_attention(self, query, blk_kv, layer_cache, fixup, block_size, max_kv_len): """Absorbed-MLA block attention: paged kernel + a block-local fixup. @@ -2738,10 +2980,15 @@ def _mla_paged_attention(self, query, blk_kv, layer_cache, fixup, block_size, ma device = query.device # --- the block's own latents into the pool, then one causal pass --- + # Must precede BOTH decode paths: each reads the block back out of the + # pool, so writing it afterwards would attend uninitialised slots. pool = layer_cache[:, 0, 0] # [pages, page_size, head_dim] pool[fixup.pages.reshape(-1), fixup.offsets.reshape(-1)] = blk_kv.reshape( -1, blk_kv.shape[-1] ) + + if self._mla_block_decode_variant(paged=True) == "cute_dsl": + return self._cute_dsl_block_decode(query, layer_cache, fixup, block_size) out_lower = torch.empty( batch, block_size, num_heads, self.kv_lora_rank, dtype=query.dtype, device=device ) @@ -2793,7 +3040,7 @@ def _mla_freqs_cis(self, device): """Adjacent-pair YaRN table, cached. See build_dspark_mla_yarn_freqs_cis. Sized from the ceiling the worker publishes once it knows the runtime - one (_runtime_position_ceiling, set in DFlashDrafter._lazy_init_ctx_buffers + one (_runtime_position_ceiling, set in DFlashWorker._lazy_init_ctx_buffers before any forward), because that is the value ctx_len is clamped to and it is strictly above model_config.max_seq_len whenever spec decoding is on. The config-derived cap is the fallback for direct construction in @@ -2945,9 +3192,10 @@ def dflash_forward( # is what keeps them out -- not the block table. page_size = None if page_tables is None else ctx_kv_cache[0].shape[-2] capacity = ctx_k_cache.shape[2] if page_tables is None else page_tables.shape[1] * page_size - # The paged kernel is the fast path; the eager gather stays as the - # reference and covers builds without flashinfer and the unpaged arena. - use_kernel = page_tables is not None and self._mla_decode_op() is not None + # attention_backend picks the implementation; both kernel backends need + # page tables, so the eager gather also covers the private arena. + variant = self._mla_block_decode_variant(page_tables is not None) + use_kernel = variant != "eager" # Kernel path: index and mask operands are layer-invariant, so build # them once here rather than five times inside the decode loop. fixup = ( @@ -2957,12 +3205,13 @@ def dflash_forward( ) logger.info_once( "MLA DSpark block decode: " - + ( - "trtllm-gen paged kernel, block in pool + upper-triangle fixup" - if use_kernel - else f"eager (paged={page_tables is not None}, " - f"flashinfer={self._mla_decode_op() is not None})" - ), + + { + "trtllm_gen": "trtllm-gen paged kernel, block in pool + upper-triangle fixup", + "cute_dsl": "one cute-dsl pass over context and block, no fixup", + "eager": f"eager (backend={self.dflash_attention_backend!r}, " + f"paged={page_tables is not None}, " + f"flashinfer={self._mla_decode_op() is not None})", + }[variant], key="mla_dspark_block_decode_variant", ) # Eager path only: dense key mask over the gathered context. @@ -2988,24 +3237,52 @@ def dflash_forward( q = attn.q_b_proj(attn.q_a_layernorm(attn.q_a_proj(hs_normed))) q = q.view(batch, block_size, num_heads, self.qk_nope_head_dim + self.qk_rope_head_dim) - # Rotate the trailing rope slice of every head in one kernel. No - # norm and no weight here -- q_a_layernorm already ran -- but the - # fused op still shape-checks the weight, hence the dummy. - q = _rmsnorm_rope_batched( - q, - self._mla_rope_noop_weight(q), - 0.0, - self.qk_rope_head_dim, - blk_freqs, - num_heads=num_heads, - apply_weight=False, - apply_rmsnorm=False, + # The 576-wide query is assembled in one buffer rather than + # cat([absorbed, rope]): q's rope half is a 64-slice of a 192-wide + # row, so that cat had a non-contiguous input and PyTorch dropped + # off CatArrayBatchedCopy_alignedK_contig onto the generic kernel -- + # 17.3 us vs 1.6 us per launch, five launches per iteration + # (measured job 2986420). Mirrors fused_q in modules/mla.py + # forward_absorption_generation. + query = torch.empty( + (batch, block_size, num_heads, self.kv_lora_rank + self.qk_rope_head_dim), + dtype=q.dtype, + device=q.device, ) - q_nope, q_rope = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + # Rotate the trailing rope slice of every head straight into the + # destination's rope columns. No norm and no weight here -- + # q_a_layernorm already ran. Writing the rotation into `query` + # instead of back into `q` removes the strided + # `query[..., kv_lora_rank:] = q_rope` copy that replaced the cat. + q_rope_into = _mla_rope_into( + q, blk_freqs, query, num_heads, self.qk_rope_head_dim, self.kv_lora_rank + ) + if not q_rope_into: + q = _rmsnorm_rope_batched( + q, + None, + 0.0, + self.qk_rope_head_dim, + blk_freqs, + num_heads=num_heads, + apply_weight=False, + apply_rmsnorm=False, + ) + query[..., self.kv_lora_rank :] = q[..., self.qk_nope_head_dim :] + q_nope = q[..., : self.qk_nope_head_dim] # Absorb the nope half into latent space so the block attends the # stored latent directly (MQA), instead of expanding it per head. - q_absorbed = torch.einsum("bshd,hdc->bshc", q_nope, attn._k_b_proj.to(q_nope.dtype)) - query = torch.cat([q_absorbed, q_rope], dim=-1) + # [h, b*s, nope] x [h, nope, kv_lora_rank] straight into the latent + # half. view(), never reshape(): both of these are strided views and + # reshape would silently stage a copy that bmm_out then writes into + # and throws away, whereas view() raises. + torch.ops.trtllm.bmm_out( + q_nope.view(batch * block_size, num_heads, self.qk_nope_head_dim).transpose(0, 1), + attn._k_b_proj.to(q.dtype), + query[..., : self.kv_lora_rank] + .view(batch * block_size, num_heads, self.kv_lora_rank) + .transpose(0, 1), + ) # The block's own latent comes from the draft hidden; the context's # from the projected target hidden (precompute_context_kv). That diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 4d80300d2567..559b9cb8a7cb 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1413,7 +1413,7 @@ def external_drafter_config_kwargs(model_config, spec_config) -> dict: # The user's value, NOT the engine's. py_executor_creator raises # model_engine_max_seq_len past this and never writes it back, so a # drafter that indexes absolute positions must read the raised value at - # runtime (DFlashDrafter publishes it as _runtime_position_ceiling) + # runtime (DFlashWorker publishes it as _runtime_position_ceiling) # rather than have this line predict it -- reproducing that arithmetic # here is what let the two drift apart in the first place. max_seq_len=model_config.max_seq_len, diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 8f0936e0307b..f03374914337 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -368,6 +368,12 @@ def _validate_draft_attention_backend(self, draft_model) -> None: ) if draft_model_dflash_attention_backend is None: raise ValueError("DFlash draft model is missing dflash_attention_backend.") + if self._dflash_attention_backend == "AUTO": + # The draft model resolved AUTO against its own family and this + # build; adopt that, because _lazy_init_ctx_buffers keys paging and + # the trtllm-gen shape checks off the concrete value. + self._dflash_attention_backend = draft_model_dflash_attention_backend + return if draft_model_dflash_attention_backend != self._dflash_attention_backend: raise ValueError( "DFlash worker and draft model attention backends must match; " diff --git a/tensorrt_llm/_torch/speculative/dflash_attention.py b/tensorrt_llm/_torch/speculative/dflash_attention.py index 17954dbd8db5..5759211742e1 100644 --- a/tensorrt_llm/_torch/speculative/dflash_attention.py +++ b/tensorrt_llm/_torch/speculative/dflash_attention.py @@ -103,8 +103,13 @@ def validate_dflash_fa4_runtime( raise RuntimeError(f"DFlash FA4 attention does not support head_dim={head_dim} on SM90.") -def _get_trtllm_gen_unavailability_reason() -> Optional[str]: - """Return why the DFlash TRTLLM backend cannot be initialized.""" +def dflash_trtllm_gen_unavailability_reason() -> Optional[str]: + """Return why the DFlash TRTLLM backend cannot be initialized, or None. + + Public because ``attention_backend="AUTO"`` resolves through it: a + drafter that prefers TRTLLM has to know whether to fall back before it + commits, and ``get_dflash_trtllm_gen_ops`` only reports by raising. + """ if not IS_FLASHINFER_AVAILABLE: return "flashinfer is not installed" @@ -124,7 +129,7 @@ def _get_trtllm_gen_unavailability_reason() -> Optional[str]: @lru_cache(maxsize=1) def get_dflash_trtllm_gen_ops() -> DFlashTrtllmGenOps: """Load TRTLLM-Gen operations after validating common prerequisites.""" - unavailable_reason = _get_trtllm_gen_unavailability_reason() + unavailable_reason = dflash_trtllm_gen_unavailability_reason() if unavailable_reason is not None: raise RuntimeError(f"DFlash TRTLLM attention backend is unavailable: {unavailable_reason}.") diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 27dfc35fcafc..1d07b7633ce3 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3133,18 +3133,24 @@ class DSparkDecodingConfig(DecodingBaseConfig): decoding_type: Literal["DSpark"] = Field(default="DSpark") - attention_backend: Literal["VANILLA", "TRTLLM"] = Field( - default="VANILLA", - description= - "Attention backend for the pooled-context cross-attention of a " - "standalone DSpark drafter (one shipped as its own checkpoint rather " - "than inside the target's mtp.* namespace). Ignored by the embedded " - "DeepSeek-V4-Pro draft, which uses its own captured-context attention. " - "This is independent of the backend used to construct the drafter's " - "standard attention modules. TRTLLM requires FlashInfer and an NVIDIA " - "Blackwell GPU with SM100 or SM103, and uses generated FMHA kernels " - "with a private paged context cache; VANILLA uses FlashAttention with " - "a contiguous cache.") + attention_backend: Literal["AUTO", "VANILLA", "TRTLLM", "CUTEDSL"] = Field( + default="AUTO", + description= + "Attention backend for the block decode of a standalone DSpark drafter " + "(one shipped as its own checkpoint rather than inside the target's " + "mtp.* namespace). Ignored by the embedded DeepSeek-V4-Pro draft, which " + "uses its own captured-context attention. This is independent of the " + "backend used to construct the drafter's standard attention modules. " + "AUTO picks per drafter family and degrades when a kernel is " + "unavailable: TRTLLM for a GQA-backboned drafter, CUTEDSL for an " + "MLA-backboned one. TRTLLM requires FlashInfer and an NVIDIA Blackwell " + "GPU with SM100 or SM103; for a GQA backbone it uses generated FMHA " + "kernels with a private paged context cache, and for an MLA backbone " + "the absorbed-MLA paged decode plus a block-local fixup. VANILLA uses " + "FlashAttention with a contiguous cache on a GQA backbone and the eager " + "torch reference on an MLA one. CUTEDSL is MLA-only: one cute-dsl pass " + "over context and block that replaces the fixup; it does not degrade, " + "and a build that cannot run it raises with the reason.") @model_validator(mode="after") def set_max_total_draft_tokens(self): diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py index dd92926d7d26..52ff3b95db0b 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py @@ -62,7 +62,7 @@ def _has_absorbed_mla_kernel() -> bool: ) # --------------------------------------------------------------------------- -# Reference oracle: line-for-line port of DeepSpec VanillaMarkov +# Reference: line-for-line port of DeepSpec VanillaMarkov # (deepspec/modeling/dspark/markov_head.py) at temperature 0. # --------------------------------------------------------------------------- @@ -203,7 +203,7 @@ def test_swa_window_conventions(): # --------------------------------------------------------------------------- # Tiny end-to-end drafter: config parsing, weight loading, block-decode -# parity vs an fp32 eager oracle (needs CUDA + flash_attn). +# parity vs an eager bf16 reference (needs CUDA + flash_attn). # --------------------------------------------------------------------------- TINY = dict( @@ -350,8 +350,11 @@ def _build_drafter( def _rms(x, w, eps=1e-6): + # fp32 accumulation, x's dtype out -- the same contract as trtllm's RMSNorm. + # Returning fp32 unconditionally would silently upcast a bf16 reference at + # the first norm and make every downstream matmul fp32. xf = x.float() - return (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps)) * w.float() + return ((xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps)) * w.float()).to(x.dtype) def _rope(x, positions, theta=10000.0): @@ -360,16 +363,21 @@ def _rope(x, positions, theta=10000.0): hd = x.shape[-1] inv = 1.0 / theta ** (torch.arange(0, hd, 2, dtype=torch.float64) / hd) ang = positions.double().unsqueeze(-1) * inv # [T, hd/2] - cos = ang.cos().float().unsqueeze(1) - sin = ang.sin().float().unsqueeze(1) + cos = ang.cos().to(x.dtype).unsqueeze(1) + sin = ang.sin().to(x.dtype).unsqueeze(1) x1, x2 = x[..., : hd // 2], x[..., hd // 2 :] return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1) -def _oracle_block_decode(weights, captured, noise_embed, use_swa): - """fp32 eager port of the DeepSpec dspark block decode - (Qwen3DSparkDecoderLayer stack over [context ; draft block]).""" - w = {k: v.float() for k, v in weights.items()} +def _reference_block_decode(weights, captured, noise_embed, use_swa): + """Eager bf16 port of the DeepSpec dspark block decode + (Qwen3DSparkDecoderLayer stack over [context ; draft block]). + + bf16 matches the drafter, so the comparison measures the algorithm rather + than the dtype; torch keeps fp32 accumulation inside its bf16 matmuls. + """ + bf16 = torch.bfloat16 + w = {k: v.to(bf16) for k, v in weights.items()} nh, nkv, hd = (TINY["num_attention_heads"], TINY["num_key_value_heads"], TINY["head_dim"]) ctx = captured.shape[0] blk = noise_embed.shape[0] @@ -379,9 +387,9 @@ def _oracle_block_decode(weights, captured, noise_embed, use_swa): # Target feature projection: hidden_norm(fc(captured)); constant across # layers, no input_layernorm on the context path (generic DFlash). - ctx_feat = _rms(captured.float() @ w["fc.weight"].T, w["hidden_norm.weight"]) + ctx_feat = _rms(captured.to(bf16) @ w["fc.weight"].T, w["hidden_norm.weight"]) - hs = noise_embed.float() + hs = noise_embed.to(bf16) for i in range(TINY["num_hidden_layers"]): p = f"layers.{i}." h = _rms(hs, w[p + "input_layernorm.weight"]) @@ -620,11 +628,11 @@ def _run_block_decode(drafter, weights, captured, noise_embed): @needs_gpu -def test_dspark_block_decode_matches_reference_oracle(): +def test_dspark_block_decode_matches_reference(): """The full drafter block decode (fc/hidden_norm projection, per-layer QKV + q/k-norm + RoPE, non-causal SWA flash attention over - [context ; block], MLP, final norm) matches the fp32 eager oracle; the - no-window oracle does NOT match (the window demonstrably binds).""" + [context ; block], MLP, final norm) matches the eager bf16 reference; the + no-window reference does NOT match (the window demonstrably binds).""" torch.manual_seed(0) weights = _tiny_weights() drafter = _build_drafter(True, weights) @@ -635,24 +643,24 @@ def test_dspark_block_decode_matches_reference_oracle(): out = _run_block_decode(drafter, weights, captured, noise_embed) - # Oracle consumes the same bf16-quantized inputs the drafter sees. + # The reference consumes the same bf16 inputs the drafter sees. captured_q = captured.to(torch.bfloat16) noise_q = noise_embed.to(torch.bfloat16) - oracle_swa = _oracle_block_decode(weights, captured_q, noise_q, True) - oracle_full = _oracle_block_decode(weights, captured_q, noise_q, False) + expected_swa = _reference_block_decode(weights, captured_q, noise_q, True) + expected_full = _reference_block_decode(weights, captured_q, noise_q, False) - diff_swa = (out - oracle_swa).abs().max().item() - diff_full = (out - oracle_full).abs().max().item() - # bf16 forward vs fp32 oracle: tolerance well below the SWA-vs-full gap. + diff_swa = (out - expected_swa).abs().max().item() + diff_full = (out - expected_full).abs().max().item() + # Tolerance well below the SWA-vs-full gap; see the measured values below. assert diff_swa < 0.02, f"SWA parity failed: max abs diff {diff_swa}" assert diff_full > 4 * max(diff_swa, 1e-4), ( - f"negative control failed: no-window oracle too close " + f"negative control failed: no-window reference too close " f"({diff_full} vs swa {diff_swa}) — window may not be applied" ) @needs_gpu -def test_plain_dflash_block_decode_matches_full_attention_oracle(): +def test_plain_dflash_block_decode_matches_full_attention_reference(): """No-regression numeric check: the plain-DFlash drafter (no dspark fields) still runs full non-causal attention over the whole context.""" weights = _tiny_weights() @@ -661,10 +669,10 @@ def test_plain_dflash_block_decode_matches_full_attention_oracle(): captured = torch.randn(CTX_LEN, TINY["hidden_size"] * NUM_CAPTURE, generator=g) * 0.5 noise_embed = torch.randn(TINY["block_size"], TINY["hidden_size"], generator=g) * 0.5 out = _run_block_decode(drafter, weights, captured, noise_embed) - oracle = _oracle_block_decode( + expected = _reference_block_decode( weights, captured.to(torch.bfloat16), noise_embed.to(torch.bfloat16), False ) - diff = (out - oracle).abs().max().item() + diff = (out - expected).abs().max().item() assert diff < 0.02, f"plain DFlash parity failed: max abs diff {diff}" @@ -784,7 +792,7 @@ def rnd(*shape): return w -def _build_mla_drafter(weights, *, dflash_attention_backend="VANILLA"): +def _build_mla_drafter(weights, *, dflash_attention_backend="TRTLLM"): from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_dspark import MLADSparkForCausalLM @@ -862,31 +870,37 @@ def test_mla_dspark_yarn_rope_matches_hf_reference(): ) -def _oracle_mla_block_decode(weights, captured, noise_embed, *, swap_kv_b=False): - """fp32 eager MLA block decode, written in the UN-ABSORBED form. +def _reference_mla_block_decode(weights, captured, noise_embed, *, swap_kv_b=False): + """Eager bf16 MLA block decode, written in the UN-ABSORBED form. The drafter attends the stored latent directly (``q_nope`` folded through ``kv_b_proj``'s K half); this expands the latent into per-head K/V and runs plain attention instead, so agreement is evidence about the absorption rather than a restatement of it. ``swap_kv_b`` is the negative control: exchanging the K and V halves must break parity. + + bf16, matching the drafter, so the comparison measures the ALGORITHM and + not the dtype. torch keeps fp32 accumulation inside its bf16 matmuls, which + is also what the drafter's GEMMs and its fp32 attention accumulator do. """ from tensorrt_llm._torch.models.modeling_dspark import apply_dspark_mla_rope - w = {k: v.float() for k, v in weights.items()} + bf16 = torch.bfloat16 + w = {k: v.to(bf16) for k, v in weights.items()} ctx, blk = captured.shape[0], noise_embed.shape[0] ctx_pos = torch.arange(ctx, dtype=torch.long) q_pos = torch.arange(ctx, ctx + blk, dtype=torch.long) cos, sin = _hf_yarn_reference( MLA_ROPE, MLA_THETA, MLA_FACTOR, MLA_ORIG_MAX, MLA_ORIG_MAX * int(MLA_FACTOR) ) + cos, sin = cos.to(bf16), sin.to(bf16) scale = _yarn_mscale_sq() / (MLA_NOPE + MLA_ROPE) ** 0.5 # Context features: context_norm(context_proj(captured)); constant across # layers and NOT passed through input_layernorm (the DFlash contract). - ctx_feat = _rms(captured.float() @ w["context_proj.weight"].T, w["context_norm.weight"]) + ctx_feat = _rms(captured.to(bf16) @ w["context_proj.weight"].T, w["context_norm.weight"]) - hs = noise_embed.float() + hs = noise_embed.to(bf16) for i in range(MLA_LAYERS): p = f"layers.{i}." h = _rms(hs, w[p + "input_layernorm.weight"]) @@ -944,9 +958,27 @@ def _yarn_mscale_sq(): # looser, so a smaller value silently tests one variant and not the other. MLA_PAGE = 32 +# Two pages so the batch crosses a page boundary -- a single page never +# exercises the addressing this path exists for. It is no longer a hard +# requirement: a (1, 1) page table used to be rejected by the CuTe DSL runtime +# with "Can't deduce the leading dimension from layout", and +# cute_dsl_custom_ops.py now hands that degenerate shape an already-marked +# tensor (see test_mla_dspark_cute_dsl_takes_a_single_page_batch_one). The +# shared CTX_LEN is 24 and is pinned by the SWA tests, so this one is local. +MLA_CTX_LEN = 40 # ceil((40 + MLA_BLOCK) / MLA_PAGE) = 2 pages +MLA_SHORT_CTX_LEN = 20 # one page, and the block stays inside it + -def _run_mla_block_decode(drafter, captured, noise_embed, paged=False): - """Drive one block decode, optionally through the paged cache. +def _run_mla_block_decode(drafter, requests, paged=False): + """Drive one block decode over a BATCH, optionally through the paged cache. + + ``requests`` is a list of ``(captured, noise_embed)`` pairs; each one's + context length is its own ``captured.shape[0]``, and they are deliberately + unequal. B >= 2 with unequal lengths is what makes the paged path testable + at all: at B == 1 the page table's stride is 1 either way, so + ``.t()`` and ``.t().contiguous()`` are indistinguishable, and + ``kv_bounds`` built with ``repeat`` instead of ``repeat_interleave`` gives + the same vector. Both are silent at B == 1 and wrong at B == 2. ``paged=False`` keeps the dense arena the eager reference path uses. ``paged=True`` hands ``dflash_forward`` a page table, which is what selects @@ -954,58 +986,87 @@ def _run_mla_block_decode(drafter, captured, noise_embed, paged=False): variant silently falls back to the same eager branch. """ dev = "cuda" - proj = drafter.project_target_hidden(captured.to(dev, torch.bfloat16)) - ctx_pos = torch.arange(CTX_LEN, device=dev) - latent, v = drafter.precompute_context_kv(proj, ctx_pos) - assert v is None, "an MLA drafter stores no V half" - assert latent.shape == (CTX_LEN, MLA_LAYERS, 1, MLA_LATENT) + batch = len(requests) + ctx_lens = [c.shape[0] for c, _ in requests] + max_ctx = max(ctx_lens) + + latents = [] + for captured, _ in requests: + proj = drafter.project_target_hidden(captured.to(dev, torch.bfloat16)) + latent, v = drafter.precompute_context_kv(proj, torch.arange(captured.shape[0], device=dev)) + assert v is None, "an MLA drafter stores no V half" + assert latent.shape == (captured.shape[0], MLA_LAYERS, 1, MLA_LATENT) + latents.append(latent) kwargs = dict( - noise_embedding=noise_embed.to(dev, torch.bfloat16).unsqueeze(0), - query_positions=torch.arange(CTX_LEN, CTX_LEN + MLA_BLOCK, device=dev).unsqueeze(0), - num_ctx_per_req=torch.tensor([CTX_LEN], device=dev), + noise_embedding=torch.stack([n.to(dev, torch.bfloat16) for _, n in requests]), + query_positions=torch.stack([torch.arange(L, L + MLA_BLOCK, device=dev) for L in ctx_lens]), + num_ctx_per_req=torch.tensor(ctx_lens, device=dev), ctx_v_cache=None, - ctx_cache_batch_idx=torch.tensor([0], device=dev), + ctx_cache_batch_idx=torch.arange(batch, device=dev), ) if not paged: pool = torch.zeros( - 1, MLA_LAYERS, CTX_LEN + MLA_BLOCK, 1, MLA_LATENT, dtype=torch.bfloat16, device=dev + batch, + MLA_LAYERS, + max_ctx + MLA_BLOCK, + 1, + MLA_LATENT, + dtype=torch.bfloat16, + device=dev, ) - pool[0, :, :CTX_LEN] = latent.permute(1, 0, 2, 3) + for b, latent in enumerate(latents): + pool[b, :, : ctx_lens[b]] = latent.permute(1, 0, 2, 3) out = drafter.dflash_forward(ctx_k_cache=pool, **kwargs) - return out.float().cpu() + # dflash_forward returns [B * block, hidden]; give the caller the batch + # axis back. At B == 1 the flat form happened to match the reference's + # [block, hidden], which is why this never surfaced before. + return out.reshape(batch, MLA_BLOCK, -1).float().cpu() # [pages, kv_factor=1, nkv=1, page_size, head_dim] per layer, plus a block's # worth of slack: the manager reserves it, and the non-causal path writes - # the draft block into it. - npages = -(-(CTX_LEN + MLA_BLOCK) // MLA_PAGE) + # the draft block into it. Each request owns a DISJOINT page range, so a + # transposed or mis-ordered page table reads another request's context + # rather than silently reading the same one back. + per_req = -(-(max_ctx + MLA_BLOCK) // MLA_PAGE) pool = [ - torch.zeros(npages, 1, 1, MLA_PAGE, MLA_LATENT, dtype=torch.bfloat16, device=dev) + torch.zeros(batch * per_req, 1, 1, MLA_PAGE, MLA_LATENT, dtype=torch.bfloat16, device=dev) for _ in range(MLA_LAYERS) ] - rows, cols = ctx_pos // MLA_PAGE, ctx_pos % MLA_PAGE - for layer_idx in range(MLA_LAYERS): - pool[layer_idx][rows, 0, 0, cols] = latent[:, layer_idx, 0] - page_table = torch.arange(npages, device=dev, dtype=torch.int32).unsqueeze(0) + for b, latent in enumerate(latents): + pos = torch.arange(ctx_lens[b], device=dev) + rows = b * per_req + pos // MLA_PAGE + cols = pos % MLA_PAGE + for layer_idx in range(MLA_LAYERS): + pool[layer_idx][rows, 0, 0, cols] = latent[:, layer_idx, 0] + page_table = torch.arange(batch * per_req, device=dev, dtype=torch.int32).view(batch, per_req) out = drafter.dflash_forward( ctx_k_cache=pool[0], ctx_kv_cache=pool, ctx_page_table=page_table, **kwargs ) - return out.float().cpu() + return out.reshape(batch, MLA_BLOCK, -1).float().cpu() @needs_cuda @pytest.mark.parametrize( - "paged", - [False, pytest.param(True, marks=needs_absorbed_mla)], - ids=["eager", "kernel_fixup"], + "paged,backend", + [ + (False, "VANILLA"), + pytest.param(True, "TRTLLM", marks=needs_absorbed_mla), + pytest.param(True, "CUTEDSL", marks=needs_absorbed_mla), + ], + ids=["eager", "kernel_fixup", "cute_dsl"], ) -def test_mla_dspark_block_decode_matches_unabsorbed_oracle(monkeypatch, paged): +def test_mla_dspark_block_decode_matches_unabsorbed_reference(monkeypatch, paged, backend): """Absorbed block decode == un-absorbed eager MLA, and the halves are not interchangeable (the negative control keeps this from being a tautology). - The paged variant writes the draft block into the pool and reads it back - through the kernel, so a bug in that write shows up here as a parity - failure rather than as lower acceptance length. + The paged variants write the draft block into the pool and read it back + through a kernel, so a bug in that write shows up here as a parity failure + rather than as lower acceptance length. Both kernels are compared against + the same reference: they differ in whether the block's strict upper triangle + comes from a separate fixup (TRTLLM) or from kv_bounds inside one pass + (CUTEDSL), and that is precisely the part a reference can catch and + acceptance cannot. """ import tensorrt_llm._torch.models.modeling_dspark as md @@ -1022,24 +1083,61 @@ def _counted(*a, **kw): return _original(*a, **kw) monkeypatch.setattr(md, "_build_mla_block_fixup", _counted) + if backend == "CUTEDSL": + reason = md.cute_dsl_mla_decode_unavailability_reason() + if reason is not None: + pytest.skip(f"cute-dsl MLA decode unavailable: {reason}") torch.manual_seed(0) weights = _tiny_mla_weights() - drafter = _build_mla_drafter(weights) + drafter = _build_mla_drafter(weights, dflash_attention_backend=backend) g = torch.Generator().manual_seed(42) - captured = torch.randn(CTX_LEN, MLA_HIDDEN * NUM_CAPTURE, generator=g) * 0.5 - noise_embed = torch.randn(MLA_BLOCK, MLA_HIDDEN, generator=g) * 0.5 + # Two requests with UNEQUAL context lengths. MLA_CTX_LEN spans two pages and + # puts the draft block in the second; MLA_SHORT_CTX_LEN sits inside the + # first, so the two requests differ in page count, in intra-page offset and + # in kv_bounds. At B == 1 none of those can disagree. + requests = [ + ( + torch.randn(L, MLA_HIDDEN * NUM_CAPTURE, generator=g) * 0.5, + torch.randn(MLA_BLOCK, MLA_HIDDEN, generator=g) * 0.5, + ) + for L in (MLA_CTX_LEN, MLA_SHORT_CTX_LEN) + ] - out = _run_mla_block_decode(drafter, captured, noise_embed, paged=paged) - captured_q, noise_q = captured.to(torch.bfloat16), noise_embed.to(torch.bfloat16) - oracle = _oracle_mla_block_decode(weights, captured_q, noise_q) - swapped = _oracle_mla_block_decode(weights, captured_q, noise_q, swap_kv_b=True) + out = _run_mla_block_decode(drafter, requests, paged=paged) + expected = torch.stack( + [ + _reference_mla_block_decode(weights, c.to(torch.bfloat16), n.to(torch.bfloat16)) + for c, n in requests + ] + ) + swapped = torch.stack( + [ + _reference_mla_block_decode( + weights, c.to(torch.bfloat16), n.to(torch.bfloat16), swap_kv_b=True + ) + for c, n in requests + ] + ) assert calls == {"fixup": 1 if paged else 0}, f"paged={paged} took the wrong branch: {calls}" + assert out.shape[0] == len(requests) - diff = (out - oracle).abs().max().item() + diff = (out - expected).abs().max().item() diff_swapped = (out - swapped).abs().max().item() - assert diff < 0.02, f"MLA parity failed: max abs diff {diff}" + # Measured per request against the bf16 reference, job 3049452: + # ctx 40 0.023438 eager / 0.015625 paged + # ctx 20 0.015625 in both + # Every value is an exact multiple of 2^-6, which is one bf16 ULP at this + # output magnitude -- the residual is the last bit, not an algorithm gap. + # 0.03 leaves ~1.5 ULP of headroom, and the negative control lands 89-201x + # away, so the bound is nowhere near admitting a degenerate output. + # + # Two facts say the batch is wired right rather than the tolerance being + # generous: ctx 20 is IDENTICAL across the eager and paged paths (a wrong + # page range or kv_bounds for the short request would not be), and ctx 40 + # is bit-identical to the same request run alone. + assert diff < 0.03, f"MLA parity failed: max abs diff {diff}" assert diff_swapped > 4 * max(diff, 1e-4), ( f"negative control failed: kv_b K/V halves swapped is too close " f"({diff_swapped} vs {diff}) -- the absorption may not be exercised" @@ -1080,37 +1178,140 @@ def test_mla_dspark_loads_the_torchspec_spelling_and_keeps_its_embedding(): @needs_cuda -@pytest.mark.parametrize("backend", ["VANILLA", "TRTLLM"]) -def test_mla_dspark_ignores_the_worker_attention_backend(backend): - """The MLA drafter runs _mla_paged_attention, so the field selects nothing. - - Both values must construct, and neither worker op set may be loaded -- - otherwise a drafter that never calls them drags in an optional dependency, - and the worker's per-backend shape checks (which the absorbed 64:1 / - head_dim 576 shape fails) would bind on a path that does not use them. +@pytest.mark.parametrize( + "backend,variant", + [("VANILLA", "eager"), ("TRTLLM", "trtllm_gen"), ("CUTEDSL", "cute_dsl")], +) +def test_mla_dspark_backend_selects_its_own_block_decode(backend, variant): + """The field picks among this class's implementations, not the worker's. + + Neither worker op set may be loaded -- otherwise a drafter that never calls + them drags in an optional dependency, and the worker's per-backend shape + checks (which the absorbed 64:1 / head_dim 576 shape fails) would bind on a + path that does not use them. What the field DOES drive is + _mla_block_decode_variant, and getting that wrong is silent: every variant + returns the same shape and only acceptance would move. """ from tensorrt_llm._torch.model_config import ModelConfig - from tensorrt_llm._torch.models.modeling_dspark import MLADSparkForCausalLM + from tensorrt_llm._torch.models.modeling_dspark import ( + MLADSparkForCausalLM, + cute_dsl_mla_decode_unavailability_reason, + ) - model_config = ModelConfig(pretrained_config=_tiny_mla_config(), attn_backend=backend) + if backend == "CUTEDSL": + reason = cute_dsl_mla_decode_unavailability_reason() + if reason is not None: + pytest.skip(f"cute-dsl MLA decode unavailable: {reason}") + + model_config = ModelConfig(pretrained_config=_tiny_mla_config(), attn_backend="VANILLA") drafter = MLADSparkForCausalLM(model_config, dflash_attention_backend=backend) assert drafter._uses_worker_attention_backend is False assert drafter.dflash_attention_backend == backend assert drafter._dflash_flash_attention is None assert drafter._dflash_trtllm_gen_ops is None + assert drafter._mla_block_decode_variant(paged=True) == variant + # No page tables, no kernel: the private arena is eager whatever was asked. + assert drafter._mla_block_decode_variant(paged=False) == "eager" + + +@needs_cuda +@needs_absorbed_mla +def test_mla_dspark_cute_dsl_takes_a_single_page_batch_one(): + """One page and one request -- the page table the runtime used to reject. + + (max_blocks, B) == (1, 1) has stride (1, 1) and no dimension with size > 1, + so CuTe DSL's layout deduction refused to pick a leading dimension even + though, with both extents 1, the choice cannot change an address. It is + reachable whenever the draft pool is one page deep and a single request is + resident, which is a first-draft-step shape, not an exotic one -- and since + AUTO resolves to CUTEDSL it would be the default path that died. + + The parity bound is the same one the batched test uses; a wrongly deduced + axis would show up as a magnitude error, not a last-bit one. + """ + import tensorrt_llm._torch.models.modeling_dspark as md + + reason = md.cute_dsl_mla_decode_unavailability_reason() + if reason is not None: + pytest.skip(f"cute-dsl MLA decode unavailable: {reason}") + + torch.manual_seed(0) + weights = _tiny_mla_weights() + drafter = _build_mla_drafter(weights, dflash_attention_backend="CUTEDSL") + g = torch.Generator().manual_seed(42) + captured = torch.randn(MLA_SHORT_CTX_LEN, MLA_HIDDEN * NUM_CAPTURE, generator=g) * 0.5 + noise = torch.randn(MLA_BLOCK, MLA_HIDDEN, generator=g) * 0.5 + + out = _run_mla_block_decode(drafter, [(captured, noise)], paged=True) + expected = _reference_mla_block_decode( + weights, captured.to(torch.bfloat16), noise.to(torch.bfloat16) + ) + diff = (out[0] - expected).abs().max().item() + assert diff < 0.03, f"single-page parity failed: max abs diff {diff}" + + +@needs_cuda +def test_mla_dspark_auto_backend_resolves_to_cutedsl(): + """AUTO is the one-pass kernel, and it must reach the dispatch as such. + + Asserting the dispatch and not just the field: every variant returns the + same shape, so a default that resolved but never reached + _mla_block_decode_variant would be silent. + """ + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models.modeling_dspark import ( + MLADSparkForCausalLM, + cute_dsl_mla_decode_unavailability_reason, + ) + + reason = cute_dsl_mla_decode_unavailability_reason() + if reason is not None: + pytest.skip(f"cute-dsl MLA decode unavailable: {reason}") + model_config = ModelConfig(pretrained_config=_tiny_mla_config(), attn_backend="VANILLA") + drafter = MLADSparkForCausalLM(model_config, dflash_attention_backend="AUTO") + assert drafter.dflash_attention_backend == "CUTEDSL" + assert drafter._mla_block_decode_variant(paged=True) == "cute_dsl" @needs_cuda +def test_gqa_dspark_auto_backend_prefers_trtllm_and_degrades(): + """AUTO on the GQA drafter follows the op set's own availability probe. + + VANILLA needs the optional flash-attn package and TRTLLM needs flashinfer + on SM100/SM103, so AUTO must consult the probe rather than assume either. + """ + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models.modeling_dspark import GQADSparkForCausalLM + from tensorrt_llm._torch.speculative.dflash_attention import ( + dflash_trtllm_gen_unavailability_reason, + ) + + model_config = ModelConfig(pretrained_config=_tiny_config(True), attn_backend="VANILLA") + drafter = GQADSparkForCausalLM(model_config, dflash_attention_backend="AUTO") + expected = "VANILLA" if dflash_trtllm_gen_unavailability_reason() else "TRTLLM" + assert drafter.dflash_attention_backend == expected + + def test_mla_dspark_still_rejects_an_unknown_backend(): """Dropping the VANILLA-only guard must not drop the typo check.""" from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_dspark import MLADSparkForCausalLM model_config = ModelConfig(pretrained_config=_tiny_mla_config(), attn_backend="VANILLA") - with pytest.raises(ValueError, match="must be VANILLA or TRTLLM"): + with pytest.raises(ValueError, match="attention backend must be one of"): MLADSparkForCausalLM(model_config, dflash_attention_backend="FLASHINFER") +def test_gqa_dspark_rejects_the_mla_only_backend(): + """CUTEDSL is MLA-only; the GQA drafter must not accept it silently.""" + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models.modeling_dspark import GQADSparkForCausalLM + + model_config = ModelConfig(pretrained_config=_tiny_config(True), attn_backend="VANILLA") + with pytest.raises(ValueError, match="attention backend must be one of"): + GQADSparkForCausalLM(model_config, dflash_attention_backend="CUTEDSL") + + def test_mla_dspark_rope_conventions_agree_on_scores(): """The drafter's adjacent-pair RoPE must be the HF one up to a lane swap. diff --git a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py index 7fa00613c361..494127a28ae3 100644 --- a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py +++ b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py @@ -110,7 +110,7 @@ def test_cute_dsl_dspark_rmsnorm_rope_rejects_invalid_inputs(): ) x, weight, freqs = _make_inputs(2, 5, 512, 64, 1) - with pytest.raises(ValueError, match="requires contiguous BF16"): + with pytest.raises(ValueError, match="requires regular row-strided BF16"): cute_dsl_dspark_rmsnorm_rope(x.float(), weight, freqs, 1, 64, 1e-6, True, True, False) @@ -246,3 +246,38 @@ def test_fused_dspark_rmsnorm_rope_norm_dim(split_norm): ref = torch.cat([ref[..., :nope], rot.to(ref.dtype)], dim=-1) torch.testing.assert_close(got.float(), ref.float(), atol=2e-2, rtol=2e-2) + + +def test_fused_dspark_rmsnorm_rope_accepts_no_weight() -> None: + """A gain-less call may pass None, and must match passing an ignored weight. + + The kernel is compiled with apply_weight=False, so the weight operand is not + in its signature at all; None is the only thing a caller without a norm gain + should have to produce. Asserting equality against a real-but-ignored weight + is what catches a compile wrapper that silently starts applying it. + """ + from tensorrt_llm._torch.custom_ops.dspark_rmsnorm_rope_custom_op import ( + cute_dsl_dspark_rmsnorm_rope, + is_fused_dspark_rmsnorm_rope_supported, + ) + + x, weight, freqs = _make_inputs(2, 5, 512, 64, 1, seed=11) + args = (1, 64, 1e-6, False, True, False) + + assert is_fused_dspark_rmsnorm_rope_supported(x, None, freqs, 1, 64) + without = cute_dsl_dspark_rmsnorm_rope(x, None, freqs, *args) + with_ignored = cute_dsl_dspark_rmsnorm_rope(x, weight, freqs, *args) + expected = _reference(x, weight, freqs, *args) + + torch.testing.assert_close(without, with_ignored, rtol=0, atol=0) + torch.testing.assert_close(without, expected, rtol=2e-2, atol=2e-2) + + +def test_fused_dspark_rmsnorm_rope_rejects_missing_weight() -> None: + from tensorrt_llm._torch.custom_ops.dspark_rmsnorm_rope_custom_op import ( + cute_dsl_dspark_rmsnorm_rope, + ) + + x, _, freqs = _make_inputs(2, 5, 512, 64, 1, seed=12) + with pytest.raises(ValueError, match="needs a weight when apply_weight is set"): + cute_dsl_dspark_rmsnorm_rope(x, None, freqs, 1, 64, 1e-6, True, True, False) From 029e0a619ebe73fc860297e2cea758ba1cee46cf Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Wed, 16 Sep 2026 03:24:56 -0700 Subject: [PATCH 10/25] [None][fix] Keep degenerate rows out of the rejection sampling kernel Cherry-picked from ftp/tekit!10642. A non-greedy batch with rejection sampling hit a device-side assert in vectorized_gather_kernel; the degenerate rows are now masked before they reach it. Signed-off-by: Zhenhuan Chen --- .../_torch/models/modeling_speculative.py | 28 +++++++- tensorrt_llm/_torch/speculative/interface.py | 46 ++++++++++++ .../hw_agnostic/test_dspark_heads.py | 32 +++++++++ .../test_rejection_buffers_guard.py | 71 +++++++++++++++++++ 4 files changed, 174 insertions(+), 3 deletions(-) mode change 100755 => 100644 tensorrt_llm/_torch/models/modeling_speculative.py diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py old mode 100755 new mode 100644 index 559b9cb8a7cb..8780e187bfc5 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -21,7 +21,7 @@ from ..attention.mla import MLA from ..model_config import ModelConfig, TConfig from ..modules.decoder_layer import DecoderLayer -from ..modules.embedding import Embedding +from ..modules.embedding import Embedding, get_masked_input_and_mask from ..modules.gated_mlp import GatedMLP from ..modules.linear import (Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig) @@ -92,6 +92,28 @@ def greedy_or_sample(logits: torch.Tensor, temperature: float) -> torch.Tensor: return sampled.view(probs.shape[:-1]) +def markov_prev_embeddings(prev_tokens: torch.Tensor, + markov_w1: torch.Tensor) -> torch.Tensor: + """``markov_w1[prev_tokens]`` with out-of-vocab anchors masked to zero. + + The anchor is the last accepted token, which on the one-model rejection + path comes from flashinfer's ``chain_speculative_sampling``: it pads + non-accepted positions with ``-1`` and returns an out-of-range id for a row + whose ``relu(target - draft)`` residual has no mass. Mask like + ``modules/embedding.py`` does for the target embedding, so such a row + contributes no bias instead of tripping a device-side assert. + + Args: + prev_tokens: previous token ids (draft vocab), any shape. + markov_w1: [vocab, rank]. + Returns: + ``prev_tokens.shape + (rank,)`` in ``markov_w1``'s dtype. + """ + prev_tokens, invalid = get_masked_input_and_mask(prev_tokens.long(), 0, + markov_w1.shape[0]) + return F.embedding(prev_tokens, markov_w1).masked_fill(invalid, 0) + + def dspark_markov_step_bias(prev_tokens: torch.Tensor, markov_w1: torch.Tensor, markov_w2: torch.Tensor) -> torch.Tensor: """Vanilla Markov head logit bias for one intra-block draft step. @@ -109,7 +131,7 @@ def dspark_markov_step_bias(prev_tokens: torch.Tensor, markov_w1: torch.Tensor, Returns: [B, vocab_or_shard] bias in the markov weights' dtype. """ - return F.linear(F.embedding(prev_tokens, markov_w1), markov_w2) + return F.linear(markov_prev_embeddings(prev_tokens, markov_w1), markov_w2) def dspark_markov_chain( @@ -226,7 +248,7 @@ def __init__(self, *, vocab_size: int, markov_rank: int): bias=False) def get_prev_embeddings(self, token_ids: torch.Tensor) -> torch.Tensor: - return F.embedding(token_ids.long(), self.markov_w1.weight) + return markov_prev_embeddings(token_ids, self.markov_w1.weight) def project_bias(self, latent_states: torch.Tensor, diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index bbd785ddd618..b6bab90e3e5a 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -2068,6 +2068,43 @@ def maybe_gather_sharded_draft_logits(self, from ..distributed.ops import allgather return allgather(logits, self.mapping, dim=-1) + @staticmethod + def _zero_padding_rows(logits: torch.Tensor, spec_metadata, + num_contexts: int, batch_size: int, + rows_per_request: int) -> torch.Tensor: + """Zero the logits rows belonging to CUDA-graph padding requests. + + A padding request decodes from an uninitialized KV/hidden state, so its + logits can be non-finite and its probs degenerate. flashinfer's + ``chain_speculative_sampling`` then rejects at a position whose + ``relu(target - draft)`` residual has no mass, where it reads an + uninitialized shared-memory slot and emits an out-of-range token id. + Zeroed logits give a uniform -- and therefore legal -- distribution. + + Padding requests are the ones ``populate_sampling_params_for_one_model`` + routed to ``dummy_slot_row`` (``py_seq_slot is None``), so this needs no + extra host-side state. Shapes are static: CUDA-graph safe. No-op while + ``dummy_slot_row`` is still 0, which is a live request's row rather than + a padding marker. + + Args: + logits: ``[(batch_size - num_contexts) * rows_per_request, vocab]``; + returned unchanged if the row count does not match. + rows_per_request: logits rows each request contributes. + Returns: + ``logits`` with the padding requests' rows zeroed. + """ + slot_ids = getattr(spec_metadata, "batch_slot_ids", None) + dummy_slot_row = getattr(spec_metadata, "dummy_slot_row", 0) + num_gens = batch_size - num_contexts + if (slot_ids is None or dummy_slot_row <= 0 or num_gens <= 0 + or logits.shape[0] != num_gens * rows_per_request): + return logits + is_padding = slot_ids[num_contexts:batch_size] == dummy_slot_row + if rows_per_request > 1: + is_padding = is_padding.repeat_interleave(rows_per_request) + return logits.masked_fill(is_padding.unsqueeze(-1), 0.0) + def advanced_sample_draft(self, logits: torch.Tensor, spec_metadata: "SpecMetadata", @@ -2098,6 +2135,10 @@ def advanced_sample_draft(self, step_offset=1 + (draft_step or 0)) if spec_metadata.use_rejection_sampling and draft_step is not None: + # The proposal stored below is read back by next iteration's + # rejection kernel, so a padding row must not poison it either. + logits = self._zero_padding_rows(logits, spec_metadata, 0, + batch_size, 1) draft_tokens, probs = ( sampling_batch_spec_dec_one_model_for_rejection(logits, temperatures, @@ -2272,6 +2313,9 @@ def _sample_and_accept_draft_tokens_rejection( spec_metadata.top_ks[gen_start:gen_end], spec_metadata.top_ps[gen_start:gen_end]) + gen_logits = self._zero_padding_rows(gen_logits, spec_metadata, + num_contexts, batch_size, + runtime_draft_len + 1) target_probs_flat = compute_probs_from_logits( gen_logits, temperatures, top_ks, top_ps) target_probs = target_probs_flat.reshape(num_gens, @@ -2541,6 +2585,8 @@ def advanced_sample_draft_block(self, gen_logits: torch.Tensor, step_offset=1) if getattr(spec_metadata, "use_rejection_sampling", False): + flat_logits = self._zero_padding_rows(flat_logits, spec_metadata, + num_contexts, batch_size, K) flat_tokens, flat_probs = ( sampling_batch_spec_dec_one_model_for_rejection(flat_logits, temps, diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py index 5546c1b03f85..036152f00b08 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py @@ -23,6 +23,8 @@ VanillaMarkov, build_markov_head, confident_prefix_length, + dspark_markov_step_bias, + markov_prev_embeddings, ) VOCAB, RANK, HID, B, BLK = 257, 16, 32, 3, 5 @@ -115,3 +117,33 @@ def test_confidence_head_with_markov_concat_dim(): assert out.shape == (B, BLK) with pytest.raises(AssertionError): head(hid) # with_markov requires prev_embeddings + + +# The anchor is a rejection-sampled token id, which can be -1 or out of vocab; +# an unmasked lookup raises IndexError on CPU and asserts device-side on CUDA. +@pytest.mark.parametrize("bad_token", [-1, VOCAB]) +def test_markov_bias_is_zero_for_out_of_vocab_anchor(bad_token): + torch.manual_seed(3) + w1, w2 = torch.randn(VOCAB, RANK), torch.randn(VOCAB, RANK) + prev = torch.tensor([0, bad_token], dtype=torch.long) + + emb = markov_prev_embeddings(prev, w1) + assert torch.equal(emb[0], w1[0]) + assert torch.all(emb[1] == 0.0) + assert torch.all(dspark_markov_step_bias(prev, w1, w2)[1] == 0.0) + + +def test_markov_head_tolerates_out_of_vocab_anchor(): + # Same guard reached through the module's get_prev_embeddings, which the + # gated/RNN heads also use. + torch.manual_seed(4) + head = VanillaMarkov(vocab_size=VOCAB, markov_rank=RANK).eval() + base = torch.randn(2, BLK, VOCAB) + first = torch.tensor([-1, 0], dtype=torch.long) + with torch.no_grad(): + _, corrected = head.sample_block_tokens( + base, first_prev_token_ids=first, hidden_states=None, temperature=0.0 + ) + # Invalid anchor -> no step-0 bias; the valid neighbour still gets one. + assert torch.allclose(corrected[0, 0], base[0, 0], atol=1e-5) + assert not torch.allclose(corrected[1, 0], base[1, 0], atol=1e-5) diff --git a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py index a71b9f1e6640..03151e690212 100644 --- a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py +++ b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py @@ -12,6 +12,7 @@ ``max_draft_len``, ``draft_probs``, ``batch_slot_ids``, ``full_draft_probs``). - ``SpecWorkerBase._rejection_buffers_valid`` only reads its arguments and ``spec_metadata`` attributes (no ``self`` use), so it can be called unbound. +- ``SpecWorkerBase._zero_padding_rows`` is a staticmethod over plain tensors. Requires CUDA because the buffers are allocated on ``device='cuda'``. """ @@ -269,6 +270,76 @@ def test_accept_dispatch_base_when_all_greedy(): assert calls["base"] == 1 and calls["rejection"] == 0 +# -------------------------------------------------------------------------- +# Padding-row sanitization: SpecWorkerBase._zero_padding_rows keeps the +# CUDA-graph padding requests' (possibly non-finite) logits out of the +# rejection kernel. Another staticmethod over plain tensors, same stand-ins. +# -------------------------------------------------------------------------- + +_PAD_ROW = 16 # spec_metadata.dummy_slot_row +_zero = SpecWorkerBase._zero_padding_rows + + +def _pad_meta(slot_ids, dummy_slot_row=_PAD_ROW): + ids = None if slot_ids is None else torch.tensor(slot_ids, dtype=torch.long, device="cuda") + return types.SimpleNamespace(batch_slot_ids=ids, dummy_slot_row=dummy_slot_row) + + +def _nan_logits(num_rows): + # A padding request decodes from uninitialized state; softmax of such a row + # is NaN, which is what has to be kept out of the rejection kernel. + return torch.full((num_rows, V), float("nan"), device="cuda") + + +@pytest.mark.parametrize("rows_per_request", [1, 3]) +def test_only_padding_rows_are_zeroed(rows_per_request): + # Two real requests, then two CUDA-graph padding requests on the shared + # dummy row -- the layout pad_batch() produces (padding appended last). + n = 2 * rows_per_request + real = torch.randn(n, V, device="cuda") + logits = torch.cat([real, _nan_logits(n)]) + + out = _zero(logits, _pad_meta([0, 1, _PAD_ROW, _PAD_ROW]), 0, 4, rows_per_request) + + torch.testing.assert_close(out[:n], real) + assert torch.all(out[n:] == 0.0) + assert torch.isnan(logits[n:]).all(), "must not mutate the caller's logits" + + +def test_padding_mask_excludes_context_rows(): + # Mixed batch: the first num_contexts requests own one logits row each and + # are not part of the gen slice this helper is handed. + out = _zero(_nan_logits(4), _pad_meta([7, _PAD_ROW, 3]), 1, 3, 2) + + assert torch.all(out[:2] == 0.0) # the _PAD_ROW request + assert torch.isnan(out[2:]).all() # the real request, untouched + + +# Builders, not built values: the metas hold CUDA tensors, and parametrize +# arguments are evaluated at import time, before the CUDA skip applies. +def _pad_case_no_slot_ids(): + return _pad_meta(None), 2 # rejection buffers never allocated + + +def _pad_case_unpublished_scratch_row(): + return _pad_meta([0, 1], dummy_slot_row=0), 2 # slot 0 is a live request + + +def _pad_case_unexpected_row_count(): + return _pad_meta([0, _PAD_ROW]), 3 # does not match the batch + + +@pytest.mark.parametrize( + "build_case", + [_pad_case_no_slot_ids, _pad_case_unpublished_scratch_row, _pad_case_unexpected_row_count], + ids=lambda f: f.__name__.removeprefix("_pad_case_"), +) +def test_padding_mask_returns_input_when_layout_is_unrecognized(build_case): + meta, num_rows = build_case() + logits = torch.randn(num_rows, V, device="cuda") + assert _zero(logits, meta, 0, 2, 1) is logits + + if __name__ == "__main__": import sys From 9a76ec03261042cf13bab91c297f3ff413f4c60d Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 16 Sep 2026 03:25:40 -0700 Subject: [PATCH 11/25] [None][fix] DSV4 DSpark CUDA graph RoPE bounds Squashed from ftp/tekit!10658 (13 commits). Publishes the runtime position ceiling to both the wrapper and its inner dspark_model, and defaults unknown batch rows to the scratch slot instead of slot 0. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/models/modeling_dspark.py | 39 +++-- tensorrt_llm/_torch/speculative/dspark.py | 68 +++++++- .../hw_agnostic/test_dspark_worker.py | 162 ++++++++++++++++++ 3 files changed, 252 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 42f414618a83..2eacadfe36cb 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -1095,7 +1095,7 @@ def has_heads(self) -> bool: def _runtime_position_cap(model_config, pretrained_config, slack: int = 0) -> int: - """Positions a RoPE table must span: the served length, not the advertised one. + """Fallback RoPE table length: the served length, not the advertised one. Both drafter families build a position table, and both were sized from the checkpoint's max_position_embeddings. K3 advertises 1,048,576, which costs @@ -1103,16 +1103,23 @@ def _runtime_position_cap(model_config, pretrained_config, slack: int = 0) -> in ~768 MiB transient -- while the context cache is bounded by the runtime max_seq_len. Only the table length is affected; YaRN's correction range is computed from original_max_position_embeddings and does not move. + + Shrinking the table this way is also what made its bound matter: sized from + max_position_embeddings it could not be indexed past, sized from + model_config.max_seq_len it can. Hence _runtime_position_ceiling, which + every drafter prefers over this value whenever a worker has published it. """ runtime = getattr(model_config, "max_seq_len", None) return int(runtime or getattr(pretrained_config, "max_position_embeddings", 163840)) + slack -# NOTE on slack: this is a construction-time cap only. A drafter driven by -# DFlashWorker gets its real bound at runtime from dflash_position_ceiling -# (published as _runtime_position_ceiling), because the engine raises -# max_seq_len past model_config's value and never writes it back. The DSv4 site -# keeps its own slack because it is not driven by that worker. +# NOTE on slack: this is a construction-time FALLBACK only, for direct +# construction in tests. Every drafter's real bound comes at runtime from +# dflash_position_ceiling (published as _runtime_position_ceiling), because the +# engine raises max_seq_len past model_config's value and never writes it back. +# DSv4 needs it too: not being driven by DFlashWorker keeps its ctx_len from +# being clamped to the engine's value, but its positions still come from the +# target's position_ids, which are bounded by exactly that value. class DSv4DSparkDraftModel(nn.Module): @@ -1232,6 +1239,7 @@ def __init__( # batched paths. It is built once per device and gathered/sliced by the # runtime decode positions, so the cache does not grow with sequence # length and the batched consuming op's shape remains static. + # Fallback: _dspark_freqs_table prefers _runtime_position_ceiling. self._freqs_cap = _runtime_position_cap(model_config, config, self.block_size + 2) self._freqs_table_cache: Dict = {} @@ -1336,13 +1344,23 @@ def cache_attn_weights_from_state_dict(self, weights: Dict) -> None: self._cache_attn_weights(weights) def _dspark_freqs_table(self, device: torch.device) -> torch.Tensor: - """Return the fixed-size plain-RoPE table cached for ``device``.""" - key = str(device) + """Return the plain-RoPE table cached for ``(device, cap)``. + + Sized from the ceiling the worker publishes once it knows the runtime + one, because that is what the decode positions actually reach: they come + from the target's position_ids, and py_executor_creator raises the + engine's max_seq_len past model_config's value without writing it back. + ``_freqs_cap`` is the fallback for direct construction in tests, where no + worker runs. MLADSparkForCausalLM reads the same attribute in + ``_mla_freqs_cis``. + """ + cap = int(getattr(self, "_runtime_position_ceiling", None) or self._freqs_cap) + key = (str(device), cap) cached = self._freqs_table_cache.get(key) if cached is None: cached = precompute_dspark_freqs_cis( self._attn_params["rope_head_dim"], - self._freqs_cap, + cap, rope_theta=self._rope_theta, device=device, ) @@ -1564,7 +1582,8 @@ def write_context_windows_batched( eps = float(self._attn_params["eps"]) positions = positions.long() slots = slots.long() - freqs = self._dspark_freqs_table(main_hidden.device)[positions] # [G, M, rd//2] + safe_positions = torch.where(mask, positions, torch.zeros_like(positions)) + freqs = self._dspark_freqs_table(main_hidden.device)[safe_positions] # [G, M, rd//2] cols = positions % win # [G, M] rows = slots[:, None].expand(-1, M) # [G, M] mask3 = mask.unsqueeze(-1) # [G, M, 1] diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 808ea8840e2e..a5eb338aa090 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -37,6 +37,18 @@ from ...llmapi.llm_args import DSparkDecodingConfig +def _dspark_position_ceiling(max_ctx: int, block_size: int, max_draft_len: int) -> int: + """Return the number of RoPE entries needed by the DSv4 block drafter. + + A target verification can accept the target token plus every draft token, + so ``start_pos`` advances by ``max_draft_len + 1``. The block drafter then + indexes ``block_size`` positions beginning at that start. With ``old`` + bounded by ``max_ctx``, the largest index is + ``max_ctx + max_draft_len + 1 + block_size``; the length is one greater. + """ + return int(max_ctx) + int(max_draft_len) + int(block_size) + 2 + + @dataclass class DSparkSpecMetadata(SpecMetadata): """Metadata for DSpark speculative decoding. @@ -124,21 +136,32 @@ def prepare(self): # without this, all concurrent gen requests fall through to the shared # scratch row below and corrupt each other's draft window at batch # size > 1 (GitHub #16767). Context-prefix entries are left to - # ``_seed_context_windows``; the ADP-idle (id 0) and CUDA-graph - # padding dummies are kept on the scratch row. + # ``_seed_context_windows``. CUDA-graph metadata is an explicit + # engine contract: its synthetic generation rows get resettable + # per-request slots, independent of their numeric request IDs. + # Outside that path, ADP-idle and padding dummies use scratch. num_contexts = max(0, len(self.request_ids) - self.num_generations) + is_graph_warmup = self.is_cuda_graph and num_seqs > 0 for rid in self.request_ids[num_contexts:]: - if ( + if is_graph_warmup or ( rid != ATTENTION_DP_DUMMY_REQUEST_ID and rid < worker._graph_dummy_id_floor and rid not in worker._req_to_slot ): - worker._assign_slot(rid, reset=False) + worker._assign_slot(rid, reset=is_graph_warmup) # Unknown request IDs (synthetic warmup / CUDA-graph padding, ADP idle # requests, or disagg seed forwards without a real id) map to the # dedicated throwaway scratch row so they cannot overwrite a live # request's rolling window (they previously aliased to slot 0). scratch = worker._scratch_slot + # Reset every iteration. The scratch row has no prefill to seed it + # and no completion to free it, so _advance_generation_state's + # unconditional advance would leave its ABSOLUTE position climbing + # across graph shapes and iterations until it indexes past the RoPE + # table. A real slot is bounded instead by its request's lifetime. + worker._ctx_len[scratch] = 0 + worker._valid_len[scratch] = 0 + worker._position_initialized[scratch] = False mapping = torch.tensor( [worker._req_to_slot.get(rid, scratch) for rid in self.request_ids], dtype=torch.long, @@ -270,7 +293,7 @@ def __init__( def max_draft_len(self) -> int: return self.spec_config.max_draft_len - def _lazy_init(self, draft_model, spec_metadata) -> None: + def _lazy_init(self, draft_model, spec_metadata, attn_metadata=None) -> None: block_size = int(draft_model.block_size) if block_size != self.max_draft_len: raise ValueError( @@ -279,6 +302,24 @@ def _lazy_init(self, draft_model, spec_metadata) -> None: ) if not self._win_inited: + # Published before the RoPE table is built, which happens lazily on + # the first forward. The engine's max_seq_len is what positions + # actually reach and is strictly above model_config's whenever spec + # decoding is on, so the config-derived cap undersizes the table. + # Unlike DFlash, the DSv4 block path can advance by a fully accepted + # block and then index another full block, so it needs the + # DSpark-specific bound. + max_ctx = getattr(attn_metadata, "max_seq_len", None) + if max_ctx is not None: + ceiling = _dspark_position_ceiling(max_ctx, block_size, self.max_draft_len) + draft_model._runtime_position_ceiling = ceiling + # The worker owns the DSv4 wrapper, while the RoPE-table cache + # lives on its inner ``dspark_model``. Publish to both so the + # value that sizes the table is the runtime bound rather than + # the construction fallback. + inner_model = getattr(draft_model, "dspark_model", None) + if inner_model is not None: + inner_model._runtime_position_ceiling = ceiling max_batch = spec_metadata.max_num_requests num_stages = draft_model.num_stages self._win = int(draft_model._attn_params["window_size"]) @@ -313,7 +354,9 @@ def _lazy_init(self, draft_model, spec_metadata) -> None: self._ctx_len = torch.zeros(num_rows, dtype=torch.long, device="cuda") self._valid_len = torch.zeros(num_rows, dtype=torch.long, device="cuda") self._position_initialized = torch.zeros(num_rows, dtype=torch.bool, device="cuda") - self._batch_to_slot = torch.zeros(max_batch, dtype=torch.long, device="cuda") + self._batch_to_slot = torch.full( + (max_batch,), self._scratch_slot, dtype=torch.long, device="cuda" + ) self._free_slots = deque(range(max_batch)) self._req_to_slot = {} logger.info( @@ -424,6 +467,17 @@ def _advance_generation_state( input_positions: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """Bootstrap and advance per-slot decode state without host synchronization.""" + # ``prepare()`` runs before graph creation, not before every replay. A + # CUDA-graph padding / ADP-idle request therefore reuses this one + # scratch row across all capture shapes. It has no persistent request + # state: reset it in the captured forward so its absolute position + # always bootstraps from ``input_positions`` rather than accumulating + # until a RoPE-table gather goes out of range. The row is never assigned + # to a live request, so this cannot reset user-visible state. + scratch = self._scratch_slot + self._ctx_len[scratch].zero_() + self._valid_len[scratch].zero_() + self._position_initialized[scratch].zero_() old = torch.where(self._position_initialized[slots], self._ctx_len[slots], input_positions) start_pos = old + num_accepted_tokens self._ctx_len[slots] = start_pos @@ -579,7 +633,7 @@ def _forward_impl( raw_logits = logits K = self.max_draft_len - self._lazy_init(draft_model, spec_metadata) + self._lazy_init(draft_model, spec_metadata, attn_metadata) # Backref so DSparkSpecMetadata.prepare() can maintain the host slot map # and mirror it into _batch_to_slot for the CUDA-graph-safe gen path. spec_metadata._dspark_worker = self diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py index ace16cba8ed0..b27751925e94 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -140,6 +140,7 @@ def test_worker_lazy_init_window_buffers(): assert worker._batch_to_slot is not None assert worker._batch_to_slot.shape == (8,) assert worker._batch_to_slot.device.type == "cuda" + assert worker._batch_to_slot.tolist() == [worker._scratch_slot] * 8 # idempotent buf_id = id(worker._kv_windows) worker._lazy_init(dm, meta) @@ -384,6 +385,52 @@ def test_prepare_keeps_dummy_generation_requests_on_scratch_row(): assert list(worker._free_slots) == [1, 2, 3] +def test_prepare_keeps_small_real_request_slots_across_steps(): + """Small real IDs are not mistaken for CUDA-graph warmup dummies.""" + worker = _make_worker() + meta = _make_metadata(max_num_requests=4) + worker._lazy_init(_fake_draft_model(), meta) + meta._dspark_worker = worker + meta.request_ids = [1, 2] + meta.num_generations = 2 + meta.prepare() + + slots = worker._batch_to_slot[:2] + worker._ctx_len[slots] = torch.tensor([7, 8], device="cuda") + worker._valid_len[slots] = torch.tensor([1, 2], device="cuda") + worker._position_initialized[slots] = True + + meta.prepare() + assert worker._ctx_len[slots].tolist() == [7, 8] + assert worker._valid_len[slots].tolist() == [1, 2] + assert worker._position_initialized[slots].tolist() == [True, True] + + +def test_prepare_resets_small_cuda_graph_warmup_slots(): + """CUDA-graph metadata gives warmup rows distinct, resettable slots.""" + worker = _make_worker() + meta = _make_metadata(max_num_requests=4) + worker._lazy_init(_fake_draft_model(), meta) + meta._dspark_worker = worker + meta.request_ids = [0, 1, 2] + meta.num_generations = 3 + meta.is_cuda_graph = True + meta.prepare() + + slots = worker._batch_to_slot[:3] + assert len(set(slots.tolist())) == 3 + assert all(slot != worker._scratch_slot for slot in slots.tolist()) + worker._ctx_len[slots] = torch.tensor([7, 8, 9], device="cuda") + worker._valid_len[slots] = torch.tensor([1, 2, 3], device="cuda") + worker._position_initialized[slots] = True + + meta.prepare() + slots = worker._batch_to_slot[:3] + assert worker._ctx_len[slots].tolist() == [0, 0, 0] + assert worker._valid_len[slots].tolist() == [0, 0, 0] + assert worker._position_initialized[slots].tolist() == [False, False, False] + + class _RecordingDraftModel: num_stages = 1 block_size = 5 @@ -440,6 +487,34 @@ def test_generation_state_cuda_graph_bootstrap_and_replay(): assert worker._valid_len[slots].tolist() == [3, 5] +def test_generation_state_resets_scratch_row_on_every_cuda_graph_replay(): + """Graph replays must not advance the shared dummy slot indefinitely.""" + worker = _make_worker() + worker._lazy_init(_fake_draft_model(window_size=8), _make_metadata(max_num_requests=2)) + scratch = worker._scratch_slot + slots = torch.tensor([scratch], device="cuda", dtype=torch.long) + num_accepted = torch.tensor([1], device="cuda", dtype=torch.long) + input_positions = torch.tensor([4016], device="cuda", dtype=torch.long) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + worker._advance_generation_state(slots, num_accepted, input_positions) + + graph.replay() + assert worker._ctx_len[scratch].item() == 4017 + assert worker._valid_len[scratch].item() == 1 + + graph.replay() + assert worker._ctx_len[scratch].item() == 4017 + assert worker._valid_len[scratch].item() == 1 + + num_accepted.fill_(2) + graph.replay() + # The new replay bootstraps from 4016, rather than adding to 4017. + assert worker._ctx_len[scratch].item() == 4018 + assert worker._valid_len[scratch].item() == 2 + + def test_disagg_position_bootstrap_uses_actual_positions_and_target_width(): """A gen-only worker bootstraps absolute positions once and indexes packed target rows by their runtime width rather than the configured K+1 width.""" @@ -824,3 +899,90 @@ def test_dspark_worker_policies_come_from_the_drafter(): # No Markov head -> the backbone logits pass through untouched. logits = torch.randn(2, 3, 8, device="cuda") assert worker._refine_block_logits(legacy, logits, {}, None) is logits + + +def test_lazy_init_publishes_the_runtime_position_ceiling(): + """The RoPE table must be sized from the engine's bound, not model_config's.""" + worker = _make_worker() + dm = _fake_draft_model() + dm.dspark_model = types.SimpleNamespace() + attn_metadata = types.SimpleNamespace(max_seq_len=4096) + + worker._lazy_init(dm, _make_metadata(max_num_requests=2), attn_metadata) + + # A full target verification accepts K+1 tokens, then DSpark indexes K + # block positions from that new start. The table length is max index + 1. + expected = 4096 + 5 + 1 + 5 + 1 + assert dm._runtime_position_ceiling == expected + assert dm.dspark_model._runtime_position_ceiling == expected + + +def test_freqs_table_prefers_the_runtime_ceiling_over_the_config_cap(mocker): + """The cache key has to carry the cap, not just the device. + + Moving where the cap comes from without moving the key would let any table + built before the worker publishes the ceiling serve every later call -- the + undersized table, silently, for the life of the process. + """ + from tensorrt_llm._torch.models.modeling_dspark import DSv4DSparkDraftModel + + spy = mocker.patch( + "tensorrt_llm._torch.models.modeling_dspark.precompute_dspark_freqs_cis", + return_value=torch.zeros(1), + ) + model = types.SimpleNamespace( + _freqs_cap=1000, + _runtime_position_ceiling=5000, + _attn_params={"rope_head_dim": 64}, + _rope_theta=10000.0, + _freqs_table_cache={}, + ) + device = torch.device("cpu") + + DSv4DSparkDraftModel._dspark_freqs_table(model, device) + assert spy.call_args.args[1] == 5000 + + # The cache key carries the cap, so a table built for one bound can never be + # served for the other. + del model._runtime_position_ceiling + DSv4DSparkDraftModel._dspark_freqs_table(model, device) + assert spy.call_args.args[1] == 1000 + assert len(model._freqs_table_cache) == 2 + + +def test_prepare_resets_the_scratch_row_every_iteration(): + """The throwaway row must not carry an absolute position between forwards. + + It has no prefill to seed it and no completion to free it, so without this + reset ``_advance_generation_state``'s unconditional advance leaves its + position climbing across CUDA-graph shapes and iterations until it indexes + past the RoPE table. + """ + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID + + worker = _make_worker() + meta = _make_metadata(max_num_requests=4) + worker._lazy_init(_fake_draft_model(), meta) + meta._dspark_worker = worker + scratch = worker._scratch_slot + + # State a previous padding / ADP-idle forward left behind. + worker._ctx_len[scratch] = 4096 + worker._valid_len[scratch] = 8 + worker._position_initialized[scratch] = True + + live = worker._assign_slot(100, reset=True) + worker._ctx_len[live] = 17 + worker._valid_len[live] = 5 + worker._position_initialized[live] = True + + meta.request_ids = [100, CUDA_GRAPH_DUMMY_REQUEST_ID] + meta.prepare() + + assert int(worker._ctx_len[scratch]) == 0 + assert int(worker._valid_len[scratch]) == 0 + assert not bool(worker._position_initialized[scratch]) + # A live request's own state is bounded by its lifetime and must survive. + assert int(worker._ctx_len[live]) == 17 + assert int(worker._valid_len[live]) == 5 + assert bool(worker._position_initialized[live]) From d4ac584d2bee4dbe19994a3eedfdb91112945d87 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Wed, 16 Sep 2026 05:58:07 -0700 Subject: [PATCH 12/25] [None][fix] Let each drafter family judge its own attention backend The same backend NAME resolves to different kernels per family, so one probe cannot answer for both: MLA TRTLLM was consulting the GQA trtllm-gen op set. An explicit backend the build cannot serve now raises instead of degrading. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/models/modeling_dflash.py | 73 +++++++++++++------ tensorrt_llm/_torch/models/modeling_dspark.py | 31 +++++--- tensorrt_llm/llmapi/llm_args.py | 15 ++-- .../test_kimi_k3_dspark_semantics.py | 48 ++++++++++++ 4 files changed, 131 insertions(+), 36 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index c49919541696..a53463670349 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -337,6 +337,10 @@ class DFlashForCausalLM(nn.Module): # has a third implementation they cannot express. _default_attention_backend = "VANILLA" _supported_attention_backends = ("VANILLA", "TRTLLM", "FA4") + # Where AUTO lands when the preferred backend cannot run. None means AUTO + # propagates the reason instead: a family whose only deployable target can + # run the fast kernel gains nothing from a silently slower path. + _auto_fallback_attention_backend = "VANILLA" # Whether the drafter's context KV lives in the draft KV cache manager's # paged pool rather than a private arena dense in max_seq_len. Orthogonal @@ -353,26 +357,58 @@ class DFlashForCausalLM(nn.Module): @classmethod def _resolve_auto_attention_backend(cls) -> str: - """Turn ``AUTO`` into a concrete backend for this drafter family. - - Only TRTLLM degrades: its ops are an optional dependency and the two - drafter families that prefer it differ in whether a build can load - them. CUTEDSL does not -- a build that cannot run it raises in - ``MLADSparkForCausalLM.__init__`` with the reason, which is the - intended behaviour for the only hardware that lacks it. - """ + """Turn ``AUTO`` into a concrete backend for this drafter family.""" want = cls._default_attention_backend - if want != "TRTLLM": - return want - reason = dflash_trtllm_gen_unavailability_reason() + reason = cls._attention_backend_unavailability_reason(want) if reason is None: - return "TRTLLM" + return want + fallback = cls._auto_fallback_attention_backend + if fallback is None: + raise ValueError( + f"{cls.__name__} has no usable attention backend: its default " + f"{want!r} is unavailable ({reason}) and this family does not " + f"degrade." + ) logger.info_once( - f"{cls.__name__} prefers the TRTLLM attention backend but it is " - f"unavailable ({reason}); falling back to VANILLA.", + f"{cls.__name__} prefers the {want} attention backend but it is " + f"unavailable ({reason}); falling back to {fallback}.", key=f"dflash_auto_backend_fallback_{cls.__name__}", ) - return "VANILLA" + return fallback + + @classmethod + def _attention_backend_unavailability_reason(cls, backend: str) -> Optional[str]: + """Why this build cannot run ``backend``, or None. Overridden per family. + + The same backend NAME resolves to different kernels per drafter family + -- GQA TRTLLM is the trtllm-gen FMHA op set, MLA TRTLLM is flashinfer's + absorbed-MLA paged decode -- so the probe belongs to the class rather + than to the resolver, which only knows the name. + """ + if backend == "TRTLLM": + return dflash_trtllm_gen_unavailability_reason() + return None + + @classmethod + def check_valid_attention_backend(cls, backend: str) -> None: + """Raise unless this build can actually run ``backend``. + + Both halves are fatal on an EXPLICIT request: a typo, and a backend the + build cannot serve. Silently running something else is what this + replaces -- the MLA drafter used to fall through to its eager reference, + correct but orders slower, with nothing raised. + """ + if backend not in cls._supported_attention_backends: + raise ValueError( + f"{cls.__name__} attention backend must be one of " + f"{list(cls._supported_attention_backends)}, got {backend!r}." + ) + reason = cls._attention_backend_unavailability_reason(backend) + if reason is not None: + raise ValueError( + f"attention_backend={backend!r} was requested but it is " + f"unavailable for {cls.__name__}: {reason}." + ) def __init__(self, draft_config, *, dflash_attention_backend: str = "AUTO"): """Build the draft model, resolving its architecture from the draft config @@ -428,12 +464,7 @@ def __init__(self, draft_config, *, dflash_attention_backend: str = "AUTO"): if dflash_attention_backend == "AUTO": dflash_attention_backend = self._resolve_auto_attention_backend() self.dflash_attention_backend = dflash_attention_backend - if self.dflash_attention_backend not in self._supported_attention_backends: - raise ValueError( - f"{type(self).__name__} attention backend must be one of " - f"{list(self._supported_attention_backends)}, got " - f"{self.dflash_attention_backend!r}." - ) + self.check_valid_attention_backend(self.dflash_attention_backend) # Each backend loads only its own ops; the rest stay None so the # shared paged prologue can read them unconditionally. self._dflash_flash_attention = None diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 2eacadfe36cb..a840784a866c 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -2726,6 +2726,11 @@ class MLADSparkForCausalLM(_DSparkHeadMixin, DFlashForCausalLM): # reason rather than hiding a slower path behind the name. _default_attention_backend = "TRTLLM" _supported_attention_backends = ("VANILLA", "TRTLLM", "CUTEDSL") + # No AUTO degrade. This drafter needs a 288 GB Blackwell part to hold the + # target at all, so a build that cannot run its kernels cannot run the + # model either; falling back to the eager reference would hide that behind + # a path orders slower at identical output. + _auto_fallback_attention_backend = None # The whole point of the MLA drafter: its context KV comes out of the # manager's pool (5760 B/token/rank) instead of an arena dense in @@ -2774,16 +2779,6 @@ def __init__(self, draft_config, *, dflash_attention_backend: str = "AUTO"): self.softmax_scale = (mscale * mscale) / math.sqrt( self.qk_nope_head_dim + self.qk_rope_head_dim ) - # Raise here rather than five layers into the first draft step, and name - # the reason. Reached for an explicit CUTEDSL and for the AUTO default - # alike -- neither degrades. - if self.dflash_attention_backend == "CUTEDSL": - reason = cute_dsl_mla_decode_unavailability_reason() - if reason is not None: - raise ValueError( - f"attention_backend='CUTEDSL' was requested but the cute-dsl " - f"MLA decode is unavailable: {reason}." - ) # -- shape / buffers --------------------------------------------------- @@ -2833,6 +2828,22 @@ def _mla_decode_op(): except (ImportError, AttributeError): return None + @classmethod + def _attention_backend_unavailability_reason(cls, backend: str) -> Optional[str]: + """Probe THIS family's kernels, not the base's GQA op set. + + ``_uses_worker_attention_backend`` is False here, so neither worker op + set is ever loaded and the base's trtllm-gen probe would answer about + kernels this drafter does not call. + """ + if backend == "TRTLLM": + if cls._mla_decode_op() is None: + return "flashinfer.mla.trtllm_batch_decode_with_kv_cache_mla is not importable" + return None + if backend == "CUTEDSL": + return cute_dsl_mla_decode_unavailability_reason() + return None + def _mla_block_decode_variant(self, paged: bool) -> str: """Which of the three block-decode implementations this step runs. diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 1d07b7633ce3..f579e377c9e2 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3141,16 +3141,21 @@ class DSparkDecodingConfig(DecodingBaseConfig): "mtp.* namespace). Ignored by the embedded DeepSeek-V4-Pro draft, which " "uses its own captured-context attention. This is independent of the " "backend used to construct the drafter's standard attention modules. " - "AUTO picks per drafter family and degrades when a kernel is " - "unavailable: TRTLLM for a GQA-backboned drafter, CUTEDSL for an " - "MLA-backboned one. TRTLLM requires FlashInfer and an NVIDIA Blackwell " + "AUTO picks per drafter family: TRTLLM for both a GQA- and an " + "MLA-backboned drafter, though the two resolve that name to different " + "kernels. A GQA drafter degrades to VANILLA when its kernel is " + "unavailable; an MLA drafter raises instead, since a build that cannot " + "run its kernel cannot hold the target either. " + "TRTLLM requires FlashInfer and an NVIDIA Blackwell " "GPU with SM100 or SM103; for a GQA backbone it uses generated FMHA " "kernels with a private paged context cache, and for an MLA backbone " "the absorbed-MLA paged decode plus a block-local fixup. VANILLA uses " "FlashAttention with a contiguous cache on a GQA backbone and the eager " "torch reference on an MLA one. CUTEDSL is MLA-only: one cute-dsl pass " - "over context and block that replaces the fixup; it does not degrade, " - "and a build that cannot run it raises with the reason.") + "over context and block that replaces the fixup. It is selectable but " + "not a default, because it needs a cute-dsl MLA decode taking per-token " + "kv_bounds that is not upstream yet; requesting it on a build without " + "that kernel raises with the reason.") @model_validator(mode="after") def set_max_total_draft_tokens(self): diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py index 52ff3b95db0b..aa00232f1202 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py @@ -1565,3 +1565,51 @@ def test_mla_drafter_rejects_a_partial_fused_component_set(): with pytest.raises(ValueError, match="WEIGHTS_SHARED_WITH_TARGET"): drafter.load_weights(truncated) + + +def test_mla_dspark_raises_when_the_requested_backend_cannot_run(monkeypatch): + """An explicit backend the build cannot serve must not degrade silently. + + The old behaviour fell through ``_mla_block_decode_variant`` to the eager + reference: identical output, orders slower, nothing raised. Only a log line + distinguished it, and that line named the GQA op set this drafter never + loads. + """ + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models.modeling_dspark import MLADSparkForCausalLM + + monkeypatch.setattr(MLADSparkForCausalLM, "_mla_decode_op", staticmethod(lambda: None)) + model_config = ModelConfig(pretrained_config=_tiny_mla_config(), attn_backend="VANILLA") + with pytest.raises(ValueError, match="trtllm_batch_decode_with_kv_cache_mla"): + MLADSparkForCausalLM(model_config, dflash_attention_backend="TRTLLM") + + +def test_mla_dspark_auto_does_not_degrade_to_the_eager_reference(monkeypatch): + """AUTO on this family propagates the reason instead of falling back. + + The GQA drafter degrades because VANILLA/FlashAttention is a real path on a + smaller part. This drafter needs a 288 GB Blackwell part to hold the target + at all, so the eager reference is never the answer -- see + ``_auto_fallback_attention_backend = None``. + """ + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models.modeling_dspark import MLADSparkForCausalLM + + monkeypatch.setattr(MLADSparkForCausalLM, "_mla_decode_op", staticmethod(lambda: None)) + model_config = ModelConfig(pretrained_config=_tiny_mla_config(), attn_backend="VANILLA") + with pytest.raises(ValueError, match="does not degrade"): + MLADSparkForCausalLM(model_config, dflash_attention_backend="AUTO") + + +def test_gqa_dspark_auto_still_degrades(monkeypatch): + """The base policy is unchanged: a GQA drafter falls back rather than raise.""" + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models import modeling_dflash + from tensorrt_llm._torch.models.modeling_dspark import GQADSparkForCausalLM + + monkeypatch.setattr( + modeling_dflash, "dflash_trtllm_gen_unavailability_reason", lambda: "probe says no" + ) + model_config = ModelConfig(pretrained_config=_tiny_config(True), attn_backend="VANILLA") + drafter = GQADSparkForCausalLM(model_config, dflash_attention_backend="AUTO") + assert drafter.dflash_attention_backend == "VANILLA" From 77a3230ffcb4404cb934cdca82a74c745a3df35d Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Wed, 16 Sep 2026 06:21:26 -0700 Subject: [PATCH 13/25] [None][fix] Never reset a live DSpark slot from prepare() prepare() runs on the host before every graph replay, not just capture, so is_cuda_graph could not gate a reset: it memset every active request's rolling window each decode step. Synced from the ftp/tekit!10658 follow-up. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/speculative/dspark.py | 13 ++++++------- .../speculative/hw_agnostic/test_dspark_worker.py | 10 +++++----- 2 files changed, 11 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index a5eb338aa090..93febece0bbd 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -136,19 +136,18 @@ def prepare(self): # without this, all concurrent gen requests fall through to the shared # scratch row below and corrupt each other's draft window at batch # size > 1 (GitHub #16767). Context-prefix entries are left to - # ``_seed_context_windows``. CUDA-graph metadata is an explicit - # engine contract: its synthetic generation rows get resettable - # per-request slots, independent of their numeric request IDs. - # Outside that path, ADP-idle and padding dummies use scratch. + # ``_seed_context_windows``. CUDA-graph capture dummies and real + # generation requests acquire a persistent slot only once; prepare() + # runs before every replay, so it must never reset a live slot. + # ADP-idle (id 0) and high-ID CUDA-graph padding dummies use scratch. num_contexts = max(0, len(self.request_ids) - self.num_generations) - is_graph_warmup = self.is_cuda_graph and num_seqs > 0 for rid in self.request_ids[num_contexts:]: - if is_graph_warmup or ( + if ( rid != ATTENTION_DP_DUMMY_REQUEST_ID and rid < worker._graph_dummy_id_floor and rid not in worker._req_to_slot ): - worker._assign_slot(rid, reset=is_graph_warmup) + worker._assign_slot(rid, reset=False) # Unknown request IDs (synthetic warmup / CUDA-graph padding, ADP idle # requests, or disagg seed forwards without a real id) map to the # dedicated throwaway scratch row so they cannot overwrite a live diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py index b27751925e94..01fad4fab945 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -406,8 +406,8 @@ def test_prepare_keeps_small_real_request_slots_across_steps(): assert worker._position_initialized[slots].tolist() == [True, True] -def test_prepare_resets_small_cuda_graph_warmup_slots(): - """CUDA-graph metadata gives warmup rows distinct, resettable slots.""" +def test_prepare_keeps_cuda_graph_slots_across_replays(): + """Graph metadata must not reset rolling state before each replay.""" worker = _make_worker() meta = _make_metadata(max_num_requests=4) worker._lazy_init(_fake_draft_model(), meta) @@ -426,9 +426,9 @@ def test_prepare_resets_small_cuda_graph_warmup_slots(): meta.prepare() slots = worker._batch_to_slot[:3] - assert worker._ctx_len[slots].tolist() == [0, 0, 0] - assert worker._valid_len[slots].tolist() == [0, 0, 0] - assert worker._position_initialized[slots].tolist() == [False, False, False] + assert worker._ctx_len[slots].tolist() == [7, 8, 9] + assert worker._valid_len[slots].tolist() == [1, 2, 3] + assert worker._position_initialized[slots].tolist() == [True, True, True] class _RecordingDraftModel: From 4c0d7aef592577430e144528175ab5178090a049 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Wed, 16 Sep 2026 22:06:09 -0700 Subject: [PATCH 14/25] [None][fix] Fix drafter test fixtures against the rebased main Build the spec config from DSparkDecodingConfig / MTPDecodingConfig rather than a SimpleNamespace, so the cost path reading one more field (max_total_draft_tokens, tokens_per_gen_step) stops breaking the test. Request id 0 is ATTENTION_DP_DUMMY_REQUEST_ID, which prepare() routes to the scratch row by design, so it can never hold a persistent slot. Signed-off-by: Zhenhuan Chen --- .../kv_cache/test_kv_cache_budget_split.py | 37 ++++++++++++------- .../hw_agnostic/test_dspark_worker.py | 4 +- 2 files changed, 27 insertions(+), 14 deletions(-) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py index ecde9f7376cd..96c06555a956 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py @@ -25,7 +25,7 @@ from tensorrt_llm._torch.pyexecutor.config_utils import uses_vswa_kv_cache_layout from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode -from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig, KvCacheConfig, MTPDecodingConfig pytestmark = pytest.mark.cpu_only @@ -1065,7 +1065,13 @@ class DraftModelConfig: def get_num_attention_layers(self): return TestExternalDrafterKvDtype.DRAFT_LAYERS - target_model_config = SimpleNamespace(is_encoder_decoder=False) + # pretrained_config is read on the target (non-draft) cost path by + # _derive_v2_layer_type_attention_windows. Empty on purpose: with no + # layer_types it returns None, so the single-window default stands and + # this test keeps measuring only the draft dtype. + target_model_config = SimpleNamespace( + is_encoder_decoder=False, pretrained_config=SimpleNamespace() + ) draft_model_config = DraftModelConfig() seen_draft_model_configs = [] @@ -1089,14 +1095,15 @@ def get_cache_size_per_token(model_config, *args, **kwargs): c._mapping.has_cp_helix.return_value = False c._mapping.pp_layers.return_value = list(range(self.DRAFT_LAYERS)) c._mapping.is_last_pp_rank.return_value = True - # The real enum, not Mock(): a bare Mock answers True to EVERY - # predicate, so use_one_engine() and is_mtp_vanilla() both fire and the - # code walks branches an external drafter never takes. DSPARK satisfies - # is_external_drafter() via is_parallel_draft() and nothing else. - c._speculative_config = SimpleNamespace( - spec_dec_mode=SpeculativeDecodingMode.DSPARK, - max_draft_len=4, - ) + # The real config, not Mock() and not a SimpleNamespace: a bare Mock + # answers True to EVERY predicate, so use_one_engine() and + # is_mtp_vanilla() both fire and the code walks branches an external + # drafter never takes; a SimpleNamespace instead needs a new attribute + # stubbed every time the cost path reads one more field of the spec + # config (max_total_draft_tokens, tokens_per_gen_step, ...). DSparkDecodingConfig + # derives all of them and satisfies is_external_drafter() via + # is_parallel_draft() and nothing else. + c._speculative_config = DSparkDecodingConfig(max_draft_len=4) c._model_engine = SimpleNamespace(model=SimpleNamespace(model_config=target_model_config)) c._draft_model_engine = None c._draft_config = draft_model_config @@ -1158,8 +1165,12 @@ def test_non_external_drafter_is_untouched(self, mocker): inherited = QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) c, draft_model_config, _ = self._creator(mocker, inherited) - # A real non-external mode, not a patched predicate: MTP_EAGLE shares the - # target's KV layout, which is exactly the case this asserts is untouched. - c._speculative_config.spec_dec_mode = SpeculativeDecodingMode.MTP_EAGLE + # A real non-external config, not a patched predicate: + # MTP_EAGLE_ONE_MODEL shares the target's KV layout, exactly the case + # this asserts is untouched. Swapped whole rather than by assigning + # spec_dec_mode, which MTPDecodingConfig derives from + # num_nextn_predict_layers and does not accept being set. + c._speculative_config = MTPDecodingConfig(num_nextn_predict_layers=1, max_draft_len=4) + assert c._speculative_config.spec_dec_mode == SpeculativeDecodingMode.MTP_EAGLE_ONE_MODEL assert c._get_draft_kv_model_config() is draft_model_config diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py index 01fad4fab945..e32a24dc20f9 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -412,7 +412,9 @@ def test_prepare_keeps_cuda_graph_slots_across_replays(): meta = _make_metadata(max_num_requests=4) worker._lazy_init(_fake_draft_model(), meta) meta._dspark_worker = worker - meta.request_ids = [0, 1, 2] + # Not 0: that is ATTENTION_DP_DUMMY_REQUEST_ID, which prepare() routes to + # the scratch row by design, so it would never take a persistent slot. + meta.request_ids = [11, 12, 13] meta.num_generations = 3 meta.is_cuda_graph = True meta.prepare() From 373624f7899c7fd3ab53e524004e9141a669fb91 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Wed, 16 Sep 2026 22:12:50 -0700 Subject: [PATCH 15/25] [None][chore] Regenerate the LLM-args golden manifest DSparkDecodingConfig.attention_backend gained AUTO and CUTEDSL, and the manifest aggregates allowed_values by field name across configs. Produced by scripts/generate_llm_args_golden_manifest.py, not by hand. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/usage/llm_args_golden_manifest.json | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 943cbfb3d0f6..5fbae1a4bca4 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1501,7 +1501,9 @@ "allowed_values": [ "VANILLA", "TRTLLM", - "FA4" + "FA4", + "AUTO", + "CUTEDSL" ], "capture_policy": "literal", "kind": "categorical", From dfd03089143bb7210144ed724835ab395389fe39 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Wed, 16 Sep 2026 22:30:01 -0700 Subject: [PATCH 16/25] [None][fix] Raise on a declared backend that loads no op set The op dispatch ended in a bare else, so a subclass widening _supported_attention_backends without adding its loader would silently get FA4's ops. Restores the exhaustive form upstream had before the port. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/models/modeling_dflash.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index a53463670349..04f06423d7bb 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -484,9 +484,19 @@ def __init__(self, draft_config, *, dflash_attention_backend: str = "AUTO"): self._dflash_flash_attention = get_dflash_flash_attention() elif self.dflash_attention_backend == "TRTLLM": self._dflash_trtllm_gen_ops = get_dflash_trtllm_gen_ops() - else: + elif self.dflash_attention_backend == "FA4": self._dflash_fa4_fwd = get_dflash_fa4_fwd() self._dflash_paged_append = get_dflash_paged_append() + else: + # Not the user's typo -- check_valid_attention_backend rejected + # those above. This is a subclass that widened + # _supported_attention_backends without adding the branch that + # loads the ops, which a bare else would answer by silently + # handing it FA4's. + raise ValueError( + f"{type(self).__name__} allows attention_backend=" + f"{self.dflash_attention_backend!r} but loads no op set for it." + ) self._dflash_trtllm_gen_workspace = None self._dflash_trtllm_gen_counters = None self.register_buffer("_dflash_batch_indices", None, persistent=False) @@ -921,6 +931,7 @@ def take(key: str) -> torch.Tensor: self.mlp_convs = mlp_convs self.candidate_selector = selector return {k: v for k, v in weights.items() if k not in consumed} + #: Tensors the wrapper itself owns: not in draft_model_full, built from the #: checkpoint in load_weights, and required. WRAPPER_OWNED_WEIGHTS = ("fc.weight", "hidden_norm.weight") From d576296f7f380098607f5197c9e11f70d243aec6 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Wed, 16 Sep 2026 22:50:55 -0700 Subject: [PATCH 17/25] [None][chore] Cut the long inline comments to the facts Every inline block over five lines, compressed to the measured numbers, error strings and file:line pointers a reader cannot re-derive; the design rationale that earned the length moved into the docstring. Corrects a dflash.py:592 pointer the rebase had made stale. Signed-off-by: Zhenhuan Chen --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 14 +--- tensorrt_llm/_torch/models/modeling_dflash.py | 19 ++--- tensorrt_llm/_torch/models/modeling_dspark.py | 80 ++++++------------- .../_torch/models/modeling_kimi_linear.py | 10 +-- .../_torch/models/modeling_speculative.py | 14 +--- .../kv_cache/kv_cache_manager_v2.py | 12 +-- tensorrt_llm/_torch/speculative/dspark.py | 21 ++--- .../kv_cache/test_kv_cache_budget_split.py | 11 +-- .../test_kimi_k3_dspark_semantics.py | 37 +++------ 9 files changed, 66 insertions(+), 152 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 91a1053e525a..0b60225cb65f 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -11561,17 +11561,9 @@ def forward( cache_key] page_table_arg = page_table if page_table.shape[0] == 1 and page_table.shape[1] == 1: - # A raw torch.Tensor is re-adapted at call time by TensorAdapter, - # which runs a bare `from_dlpack(arg).mark_layout_dynamic()` - # (cute/runtime.py:915, registered for torch.Tensor at :941), so - # the leading_dim=0 given to cute.compile above does not carry - # over. A (1, 1) page table -- one page, batch 1 -- has stride - # (1, 1), no dimension with size > 1, and the deduction then - # raises "Can't deduce the leading dimension from layout" even - # though with both extents 1 the choice cannot change an address. - # An already-marked tensor has no registered adapter - # (jit_executor.py:650), so the bare call is skipped. Only the - # degenerate shape pays the wrapper. + # leading_dim=0 does not survive TensorAdapter's call-time re-adapt + # (cute/runtime.py:915), and a (1, 1) table has no extent > 1, so + # deduction raises "Can't deduce the leading dimension from layout". page_table_arg = cute.runtime.from_dlpack( page_table, assumed_align=16).mark_layout_dynamic(leading_dim=0) diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index 04f06423d7bb..338ae4be8f29 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -839,13 +839,9 @@ def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): remapped[key] = value # Wrapper-owned and built FROM the checkpoint, so they are never in - # draft_model_full's module tree and _assert_backbone_complete cannot - # see them: a module conjured from data cannot be reported missing by - # walking modules. One tuple drives both the extraction below and this - # check, so the two cannot drift. Without fc the drafter has no capture - # projection, has_target_features stays False, _ctx_len never advances - # and it drafts from an empty context forever (the `hasattr` guards on - # that path are degradation, not a supported mode). + # draft_model_full's module tree for _assert_backbone_complete to walk. + # Without fc the drafter has no capture projection and drafts from an + # empty context forever. wrapper_missing = [k for k in self.WRAPPER_OWNED_WEIGHTS if k not in remapped] if wrapper_missing: raise ValueError( @@ -960,12 +956,9 @@ def _assert_backbone_complete(self, weights: Dict, weight_mapper=None) -> None: """ provided = set(weights) - # Whichever fusion table the load below will actually use, not a third - # copy: with a mapper modeling_utils dispatches to _load_weights_impl_v2 - # and the mapper's own table applies; without one it falls back to - # _load_weights_impl, whose table is FUSED_MODULE_COMPONENTS. An empty - # `mapping` means init_model_and_config has not run, so the mapper has - # no table to offer yet and the constant is still the right answer. + # Whichever fusion table the load below uses, not a third copy: + # _load_weights_impl_v2 takes the mapper's, _load_weights_impl takes + # FUSED_MODULE_COMPONENTS. An empty mapping means the mapper has none yet. fusion = dict(getattr(weight_mapper, "mapping", None) or FUSED_MODULE_COMPONENTS) def _has(prefix: str) -> bool: diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index a840784a866c..7db8d1c97e16 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -1113,13 +1113,10 @@ def _runtime_position_cap(model_config, pretrained_config, slack: int = 0) -> in return int(runtime or getattr(pretrained_config, "max_position_embeddings", 163840)) + slack -# NOTE on slack: this is a construction-time FALLBACK only, for direct -# construction in tests. Every drafter's real bound comes at runtime from -# dflash_position_ceiling (published as _runtime_position_ceiling), because the -# engine raises max_seq_len past model_config's value and never writes it back. -# DSv4 needs it too: not being driven by DFlashWorker keeps its ctx_len from -# being clamped to the engine's value, but its positions still come from the -# target's position_ids, which are bounded by exactly that value. +# Construction-time FALLBACK only, for direct construction in tests. The +# real bound is _runtime_position_ceiling: the engine raises max_seq_len +# past model_config's and never writes back. DSv4 reaches it via the +# target's position_ids. class DSv4DSparkDraftModel(nn.Module): @@ -2661,16 +2658,9 @@ class GQADSparkForCausalLM(_DSparkHeadMixin, DFlashForCausalLM): Reference: arXiv 2607.05147; deepseek-ai/DeepSpec. """ - # TRTLLM, matching what the shipped K3 configs already pin - # (examples/kimi_k3/eval_extra_llm_options_dspark{,_nvfp4}.yaml); VANILLA - # additionally needs the optional flash-attn package. AUTO degrades to - # VANILLA off SM100/SM103. - # - # This also moves the drafter's context KV into the manager's paged pool, - # and that is the point, not a side effect: dflash.py:592 keys `use_paged` - # off the backend name, and this class leaves _paged_ctx_cache False, so - # VANILLA is the only way to get the arena that is dense in max_seq_len and - # that free_gpu_memory_fraction never bounds. + # TRTLLM matches the shipped K3 configs (examples/kimi_k3/ + # eval_extra_llm_options_dspark{,_nvfp4}.yaml) and pages the context KV; + # _paged_ctx_cache is False here, so VANILLA is the only route to the arena. _default_attention_backend = "TRTLLM" def __init__(self, draft_config, *, dflash_attention_backend: str = "AUTO"): @@ -2716,14 +2706,9 @@ class MLADSparkForCausalLM(_DSparkHeadMixin, DFlashForCausalLM): # head, head_dim 576) would fail both. The backend NAME still selects, via # _mla_block_decode_variant below. _uses_worker_attention_backend = False - # CUTEDSL stays selectable but is NOT the default here: it needs a - # cute_dsl_mla_decode_fp16_blackwell that takes per-token kv_bounds, which - # this branch's kernel does not have. Flip this one line once it does -- - # measured 35.7% off the drafter region (947.50 -> 608.95 us/iteration, - # jobs 3006025/3006026) at no cost in acceptance or accuracy, on both - # aggregated and disaggregated gsm8k (jobs 3006856/3006857, 3006673/3006684). - # CUTEDSL never degrades: an explicit request raises in __init__ with the - # reason rather than hiding a slower path behind the name. + # Selectable but not the default: needs a cute_dsl_mla_decode_fp16_blackwell + # taking per-token kv_bounds, absent on this branch. Worth 35.7% off the + # drafter region (947.50 -> 608.95 us/iter, jobs 3006025/3006026). _default_attention_backend = "TRTLLM" _supported_attention_backends = ("VANILLA", "TRTLLM", "CUTEDSL") # No AUTO degrade. This drafter needs a 288 GB Blackwell part to hold the @@ -2944,13 +2929,9 @@ def _cute_dsl_block_decode(self, query, layer_cache, fixup, block_size): d_latent = self.kv_lora_rank self._assert_cute_dsl_can_implement(batch, block_size, page_size, query.dtype) - # stride[1] == 1 is the kernel's only layout demand (mla_decode_fp16.py:465-470); - # a 576-wide row sliced at 512 keeps that on both halves, so no copy. - # detach(): cute.runtime.from_dlpack refuses a tensor that requires grad - # ("Can't export tensors that require gradient"). Production runs under - # inference_mode so nothing here carries grad, but the in-place pool - # write above pulls the pool into the autograd graph whenever a caller - # does not -- a unit test, for one. Views, so no copy. + # stride[1] == 1 is the kernel's only layout demand (mla_decode_fp16.py:465-470), + # kept by a 576-wide row sliced at 512. detach() because from_dlpack refuses + # grad-requiring tensors, which the pool write adds outside inference_mode. q_latent = query.detach()[..., :d_latent].permute(2, 3, 1, 0) q_rope = query.detach()[..., d_latent:].permute(2, 3, 1, 0) c_latent = pool.detach()[..., :d_latent].permute(1, 2, 0) @@ -2966,12 +2947,9 @@ def _cute_dsl_block_decode(self, query, layer_cache, fixup, block_size): q_rope, c_latent, c_rope, - # .t() WITHOUT .contiguous(): the kernel wants (max_blocks, B) with - # dim 0 as the leading (stride-1) dimension. contiguous() would make - # it row-major, i.e. stride (B, 1), and the kernel rejects it with - # "Expected strides[leading_dim] == 1, but got B" for every B > 1 - # (probe job 3005553; B=1 passes by coincidence). The transposed - # view of a row-major (B, max_blocks) already has stride (1, max_blocks). + # .t() WITHOUT .contiguous(): the kernel wants (max_blocks, B) with dim 0 + # stride-1. contiguous() gives stride (B, 1) and it raises "Expected + # strides[leading_dim] == 1, but got B" for B > 1 (probe job 3005553). fixup.page_tables_i32.t(), fixup.seq_lens_i32, out.permute(2, 3, 1, 0), @@ -3056,11 +3034,8 @@ def _mla_paged_attention(self, query, blk_kv, layer_cache, fixup, block_size, ma out_upper = torch.einsum("bsht,btc->bshc", probs, blk_kv[..., : self.kv_lora_rank]) # --- log-sum-exp merge (flash-decoding combine), kernel LSE is base 2 --- - # The weight pair collapses: w_u / (w_l + w_u) == sigmoid(lse_u - lse_l), - # so the merge is one sigmoid and one lerp instead of a peak subtraction, - # two exps and a division -- and it needs no explicit max for stability. - # The last query has an empty triangle, hence lse_upper -inf, sigmoid 0, - # pure kernel output. + # 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 weight_upper = torch.sigmoid(lse_upper - lse_l).unsqueeze(-1) merged = torch.lerp(out_lower.float(), out_upper.float(), weight_upper) @@ -3267,13 +3242,9 @@ def dflash_forward( q = attn.q_b_proj(attn.q_a_layernorm(attn.q_a_proj(hs_normed))) q = q.view(batch, block_size, num_heads, self.qk_nope_head_dim + self.qk_rope_head_dim) - # The 576-wide query is assembled in one buffer rather than - # cat([absorbed, rope]): q's rope half is a 64-slice of a 192-wide - # row, so that cat had a non-contiguous input and PyTorch dropped - # off CatArrayBatchedCopy_alignedK_contig onto the generic kernel -- - # 17.3 us vs 1.6 us per launch, five launches per iteration - # (measured job 2986420). Mirrors fused_q in modules/mla.py - # forward_absorption_generation. + # One buffer rather than cat([absorbed, rope]): q's rope half is a 64-slice + # of a 192-wide row, so cat fell off CatArrayBatchedCopy onto the generic + # kernel -- 17.3 us vs 1.6 us, five launches/iter (measured job 2986420). query = torch.empty( (batch, block_size, num_heads, self.kv_lora_rank + self.qk_rope_head_dim), dtype=q.dtype, @@ -3300,12 +3271,9 @@ def dflash_forward( ) query[..., self.kv_lora_rank :] = q[..., self.qk_nope_head_dim :] q_nope = q[..., : self.qk_nope_head_dim] - # Absorb the nope half into latent space so the block attends the - # stored latent directly (MQA), instead of expanding it per head. - # [h, b*s, nope] x [h, nope, kv_lora_rank] straight into the latent - # half. view(), never reshape(): both of these are strided views and - # reshape would silently stage a copy that bmm_out then writes into - # and throws away, whereas view() raises. + # Absorb the nope half into latent space so the block attends the stored + # latent directly (MQA) instead of expanding per head. view(), never + # reshape(): both are strided views and reshape would stage a silent copy. torch.ops.trtllm.bmm_out( q_nope.view(batch * block_size, num_heads, self.qk_nope_head_dim).transpose(0, 1), attn._k_b_proj.to(q.dtype), diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index c771cf2c9e0c..58376155f92f 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -1713,13 +1713,9 @@ def forward( self.self_attention_res_norm, ) if capture is not None: - # Which residual value the drafter was distilled against is a - # property of the DRAFTER checkpoint, not a tuning knob: a mismatch - # only lowers acceptance, silently. ``hidden_states`` here is the - # pre-norm attn_res mixture (SGLang's aggregate_stream); prefix_only - # wants the incoming running prefix, which vLLM builds as - # ``prefix_sum + pending_mlp_out`` on its deferred-add path and is - # already in hand as ``prefix_sum``. + # A property of the DRAFTER checkpoint, not a knob: a mismatch only lowers + # acceptance, silently. hidden_states is the pre-norm attn_res mixture; + # prefix_only wants the running prefix, already in hand as prefix_sum. tapped = hidden_states if _AUX_ATTN_RES_STREAM_ENABLED else prefix_sum capture[0].maybe_capture_hidden_states(capture[1], tapped, None) diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 8780e187bfc5..f4c4eebe8425 100644 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1427,17 +1427,9 @@ def external_drafter_config_kwargs(model_config, spec_config) -> dict: spec_config=None, # Avoid recursive spec-dec max_num_tokens=model_config.max_num_tokens, moe_max_num_tokens=model_config.moe_max_num_tokens, - # Bounds the drafter's position tables. Without it the field stays - # None and they fall back to the checkpoint's advertised - # max_position_embeddings -- 1,048,576 for K3, a ~256 MiB complex64 - # table per rank for a context the runtime bounds far below that. - # - # The user's value, NOT the engine's. py_executor_creator raises - # model_engine_max_seq_len past this and never writes it back, so a - # drafter that indexes absolute positions must read the raised value at - # runtime (DFlashWorker publishes it as _runtime_position_ceiling) - # rather than have this line predict it -- reproducing that arithmetic - # here is what let the two drift apart in the first place. + # The user's value, NOT the engine's: py_executor_creator raises it past + # this and never writes back, so drafters read _runtime_position_ceiling. + # None sizes position tables from max_position_embeddings: 1M for K3. max_seq_len=model_config.max_seq_len, ) # Only the embedded DSpark draft shares the target's EPLB namespace (its diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 89ec6ae2c605..e367a7768013 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -3721,14 +3721,10 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): for req in scheduled_batch.context_requests: kv_cache = self._mirror_draft_kv_cache(req) if kv_cache is None: - # Pre-existing behaviour, kept deliberately: skipping here - # does NOT buy a retry, because copy_batch_block_offsets() - # asserts on this id later in the same iteration (see - # _mirror_draft_kv_cache). Saturation needs a disagg worker - # whose cancelled / retired-session requests pile up past - # the 2x slack, since the normal path frees the slot before - # start_transfer -- so this is a loud symptom of that, not a - # recoverable state. + # Pre-existing behaviour, kept deliberately: skipping does NOT buy a + # retry, because copy_batch_block_offsets() asserts on this id later in + # the same iteration. Saturation needs cancelled requests piling past + # the 2x slack, so this is a loud symptom, not a recoverable state. logger.warning( f"Draft KV cache mirror has no free IndexMapper slot for " f"context request {req.py_request_id}; this iteration " diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 93febece0bbd..27761d187d55 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -301,13 +301,9 @@ def _lazy_init(self, draft_model, spec_metadata, attn_metadata=None) -> None: ) if not self._win_inited: - # Published before the RoPE table is built, which happens lazily on - # the first forward. The engine's max_seq_len is what positions - # actually reach and is strictly above model_config's whenever spec - # decoding is on, so the config-derived cap undersizes the table. - # Unlike DFlash, the DSv4 block path can advance by a fully accepted - # block and then index another full block, so it needs the - # DSpark-specific bound. + # Published before the RoPE table is built (lazily, on first forward). The + # engine's max_seq_len is what positions reach and exceeds model_config's + # whenever spec decoding is on, so the config cap undersizes the table. max_ctx = getattr(attn_metadata, "max_seq_len", None) if max_ctx is not None: ceiling = _dspark_position_ceiling(max_ctx, block_size, self.max_draft_len) @@ -466,13 +462,10 @@ def _advance_generation_state( input_positions: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """Bootstrap and advance per-slot decode state without host synchronization.""" - # ``prepare()`` runs before graph creation, not before every replay. A - # CUDA-graph padding / ADP-idle request therefore reuses this one - # scratch row across all capture shapes. It has no persistent request - # state: reset it in the captured forward so its absolute position - # always bootstraps from ``input_positions`` rather than accumulating - # until a RoPE-table gather goes out of range. The row is never assigned - # to a live request, so this cannot reset user-visible state. + # ``prepare()`` runs before graph creation, not before every replay, so a + # padding / ADP-idle request reuses this scratch row across capture shapes. + # Reset it in the captured forward so its position bootstraps from + # ``input_positions`` rather than accumulating past the RoPE table. scratch = self._scratch_slot self._ctx_len[scratch].zero_() self._valid_len[scratch].zero_() diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py index 96c06555a956..555fab8775cf 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py @@ -1095,14 +1095,9 @@ def get_cache_size_per_token(model_config, *args, **kwargs): c._mapping.has_cp_helix.return_value = False c._mapping.pp_layers.return_value = list(range(self.DRAFT_LAYERS)) c._mapping.is_last_pp_rank.return_value = True - # The real config, not Mock() and not a SimpleNamespace: a bare Mock - # answers True to EVERY predicate, so use_one_engine() and - # is_mtp_vanilla() both fire and the code walks branches an external - # drafter never takes; a SimpleNamespace instead needs a new attribute - # stubbed every time the cost path reads one more field of the spec - # config (max_total_draft_tokens, tokens_per_gen_step, ...). DSparkDecodingConfig - # derives all of them and satisfies is_external_drafter() via - # is_parallel_draft() and nothing else. + # The real config, not Mock() or SimpleNamespace: a bare Mock answers True + # to every predicate, and a SimpleNamespace needs a new attribute stubbed + # each time the cost path reads one more spec-config field. c._speculative_config = DSparkDecodingConfig(max_draft_len=4) c._model_engine = SimpleNamespace(model=SimpleNamespace(model_config=target_model_config)) c._draft_model_engine = None diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py index aa00232f1202..eb9aeff4be25 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py @@ -959,12 +959,8 @@ def _yarn_mscale_sq(): MLA_PAGE = 32 # Two pages so the batch crosses a page boundary -- a single page never -# exercises the addressing this path exists for. It is no longer a hard -# requirement: a (1, 1) page table used to be rejected by the CuTe DSL runtime -# with "Can't deduce the leading dimension from layout", and -# cute_dsl_custom_ops.py now hands that degenerate shape an already-marked -# tensor (see test_mla_dspark_cute_dsl_takes_a_single_page_batch_one). The -# shared CTX_LEN is 24 and is pinned by the SWA tests, so this one is local. +# exercises this addressing. No longer a hard requirement since +# cute_dsl_custom_ops.py pre-marks the degenerate (1, 1) table. MLA_CTX_LEN = 40 # ceil((40 + MLA_BLOCK) / MLA_PAGE) = 2 pages MLA_SHORT_CTX_LEN = 20 # one page, and the block stays inside it @@ -1067,6 +1063,11 @@ def test_mla_dspark_block_decode_matches_unabsorbed_reference(monkeypatch, paged comes from a separate fixup (TRTLLM) or from kv_bounds inside one pass (CUTEDSL), and that is precisely the part a reference can catch and acceptance cannot. + + Two facts say the batch is wired right rather than the tolerance being + generous: ctx 20 is IDENTICAL across the eager and paged paths (a wrong page + range or kv_bounds for the short request would not be), and ctx 40 is + bit-identical to the same request run alone. """ import tensorrt_llm._torch.models.modeling_dspark as md @@ -1125,18 +1126,9 @@ def _counted(*a, **kw): diff = (out - expected).abs().max().item() diff_swapped = (out - swapped).abs().max().item() - # Measured per request against the bf16 reference, job 3049452: - # ctx 40 0.023438 eager / 0.015625 paged - # ctx 20 0.015625 in both - # Every value is an exact multiple of 2^-6, which is one bf16 ULP at this - # output magnitude -- the residual is the last bit, not an algorithm gap. - # 0.03 leaves ~1.5 ULP of headroom, and the negative control lands 89-201x - # away, so the bound is nowhere near admitting a degenerate output. - # - # Two facts say the batch is wired right rather than the tolerance being - # generous: ctx 20 is IDENTICAL across the eager and paged paths (a wrong - # page range or kv_bounds for the short request would not be), and ctx 40 - # is bit-identical to the same request run alone. + # Job 3049452, per request vs the bf16 reference: ctx 40 is 0.023438 eager / + # 0.015625 paged, ctx 20 is 0.015625 in both. Every value is an exact + # multiple of 2^-6 -- one bf16 ULP here -- so 0.03 leaves ~1.5 ULP. assert diff < 0.03, f"MLA parity failed: max abs diff {diff}" assert diff_swapped > 4 * max(diff, 1e-4), ( f"negative control failed: kv_b K/V halves swapped is too close " @@ -1505,12 +1497,9 @@ def test_mla_drafter_rejects_a_checkpoint_missing_backbone_weights(): weights = _tiny_mla_weights() drafter = _build_mla_drafter(weights) - # A whole module, not one tensor: the check is module-granular on purpose, - # since a fused module (gate_up_proj) is stored unfused and dropping half of - # it is a parameter-level gap the loader's own naming cannot distinguish. - # Derive the prefix from the fixture rather than spelling it: DFlash - # checkpoints name layers without the `model.` prefix, so a hardcoded key - # silently matches nothing and the test passes by not truncating anything. + # Module-granular on purpose: a fused module (gate_up_proj) is stored + # unfused, and dropping half is a parameter-level gap the loader's naming + # cannot see. Derive the prefix -- DFlash checkpoints omit `model.`. victim = next(k for k in weights if k.endswith("self_attn.o_proj.weight")) truncated = {k: v for k, v in weights.items() if k != victim} assert len(truncated) == len(weights) - 1 From bb8785928edfc4f0731295c2952ddbb43363489a Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Wed, 16 Sep 2026 23:51:19 -0700 Subject: [PATCH 18/25] [None][chore] Cap docstrings and the backend description by point Twenty-four docstring paragraphs over five lines, cut to the facts a reader cannot re-derive; four split into two points rather than trimmed. attention_backend's description drops the per-backend kernel table, which MLADSparkForCausalLM's docstring already carries: 210 words to 96. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/models/modeling_dflash.py | 7 +- tensorrt_llm/_torch/models/modeling_dspark.py | 110 ++++++++---------- .../_torch/models/modeling_speculative.py | 11 +- tensorrt_llm/_torch/pyexecutor/_util.py | 24 ++-- .../kv_cache/kv_cache_manager_v2.py | 25 ++-- tensorrt_llm/_torch/speculative/interface.py | 5 +- tensorrt_llm/llmapi/llm_args.py | 29 ++--- .../kv_cache/test_kv_cache_budget_split.py | 15 +-- .../test_kimi_k3_dspark_semantics.py | 48 ++++---- .../test_dspark_cute_dsl_rmsnorm_rope.py | 9 +- 10 files changed, 122 insertions(+), 161 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index 338ae4be8f29..b87de23c0d13 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -946,10 +946,9 @@ def _assert_backbone_complete(self, weights: Dict, weight_mapper=None) -> None: Module granularity. A PLAIN module is the hole `allow_partial_loading` cannot close -- the loader skips one whose subtree filters to nothing - (modeling_utils.py `if module_weights:`) whatever the flag says. A FUSED - module `allow_partial_loading=False` would catch (linear.py asserts all - three shards), but the flag must stay True for the target-shared - modules, so this requires every component rather than any. + (modeling_utils.py `if module_weights:`). A FUSED module would be caught by + `allow_partial_loading=False`, but the flag must stay True for the + target-shared modules, so this requires every component rather than any. Missing parameters INSIDE a present component stay tolerated: all three weights but only `q_proj.bias` leaves the rest at `torch.empty`. diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 7db8d1c97e16..9c3f59cc6f0b 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -1097,12 +1097,11 @@ def has_heads(self) -> bool: def _runtime_position_cap(model_config, pretrained_config, slack: int = 0) -> int: """Fallback RoPE table length: the served length, not the advertised one. - Both drafter families build a position table, and both were sized from the - checkpoint's max_position_embeddings. K3 advertises 1,048,576, which costs - ~256 MiB per rank as complex64 -- built via full-size fp32 cos/sin, so - ~768 MiB transient -- while the context cache is bounded by the runtime - max_seq_len. Only the table length is affected; YaRN's correction range is - computed from original_max_position_embeddings and does not move. + Both drafter families size a position table from the checkpoint's + max_position_embeddings. K3 advertises 1,048,576: ~256 MiB per rank as + complex64, ~768 MiB transient via full-size fp32 cos/sin, while the context + cache is bounded by the runtime max_seq_len. Only the table length moves; + YaRN's correction range comes from original_max_position_embeddings. Shrinking the table this way is also what made its bound matter: sized from max_position_embeddings it could not be indexed past, sized from @@ -1343,13 +1342,11 @@ def cache_attn_weights_from_state_dict(self, weights: Dict) -> None: def _dspark_freqs_table(self, device: torch.device) -> torch.Tensor: """Return the plain-RoPE table cached for ``(device, cap)``. - Sized from the ceiling the worker publishes once it knows the runtime - one, because that is what the decode positions actually reach: they come - from the target's position_ids, and py_executor_creator raises the - engine's max_seq_len past model_config's value without writing it back. - ``_freqs_cap`` is the fallback for direct construction in tests, where no - worker runs. MLADSparkForCausalLM reads the same attribute in - ``_mla_freqs_cis``. + Sized from the ceiling the worker publishes once it knows the runtime one, + because that is what decode positions reach: they come from the target's + position_ids, and py_executor_creator raises the engine's max_seq_len past + model_config's without writing it back. ``_freqs_cap`` is the fallback for + direct construction in tests; _mla_freqs_cis reads the same attribute. """ cap = int(getattr(self, "_runtime_position_ceiling", None) or self._freqs_cap) key = (str(device), cap) @@ -2386,10 +2383,9 @@ class _MLABlockFixup(NamedTuple): The block's own latents are written into the pool at ``ctx_len + j``. Those slots belong to this request -- the draft KV manager adds - ``max(draft_len, max_total_draft_tokens)`` tokens to every generation request - each step (resource_manager.py:1060) -- and the accepted tokens overwrite - them next step, so the write is transient. Same thing the GQA drafter's - TRTLLM backend does (modeling_dflash.py append_paged_kv_cache). + ``max(draft_len, max_total_draft_tokens)`` tokens per generation request each + step (resource_manager.py:1060) -- and the accepted tokens overwrite them next + step. Same thing the GQA drafter's TRTLLM backend does. In-block causality then leaves query j seeing block keys 0..j, so only the strict upper triangle needs fixing up. Every field depends on ctx_len, the @@ -2502,12 +2498,11 @@ def _resolve_dspark_mla_rope_params(pretrained_config, model_config=None, slack: class _DSparkHeadMixin: """The DSpark head set, shared by the GQA and MLA drafter wrappers. - DSpark is DFlash plus three things, and only these three are common to both - backbone shapes: the vanilla Markov intra-block logit bias, the - ``shift_label`` slot convention (block slot j predicts draft token j+1, so - slot 0 holds the anchor) and the confidence head. The block decode itself - has nothing in common between the shapes, which is why this is a mixin and - not a base class. + DSpark is DFlash plus three things common to both backbone shapes: the + vanilla Markov intra-block logit bias, the ``shift_label`` slot convention + (block slot j predicts draft token j+1, so slot 0 holds the anchor) and the + confidence head. The block decode has nothing in common between the shapes, + which is why this is a mixin and not a base class. Confidence-scheduled verification is not implemented yet: ``confidence_proj`` is loaded but unused, and drafting always proposes the full K tokens. @@ -2647,13 +2642,14 @@ class GQADSparkForCausalLM(_DSparkHeadMixin, DFlashForCausalLM): Adds the DSpark head set (see :class:`_DSparkHeadMixin`) on top of the DFlash block decode. - Named for the attention shape, not for a model: the backbone is whatever - the drafter config resolves to through the model registry, and the - inherited block decode works for every GQA family the DFlash drafters - already cover (qwen3, llama, gpt_oss, ...). A per-model subclass would be - empty. The GQA precondition is inherited, not introduced here -- see - ``DFlashForCausalLM._validate_gqa_shape``. The MLA-backboned drafter is the - sibling :class:`MLADSparkForCausalLM`, not a subclass of this. + Named for the attention shape, not a model: the backbone is whatever the + drafter config resolves to, and the inherited block decode covers every GQA + family the DFlash drafters already do (qwen3, llama, gpt_oss, ...), so a + per-model subclass would be empty. The GQA precondition is inherited from + ``DFlashForCausalLM._validate_gqa_shape``, not introduced here. + + The MLA-backboned drafter is the sibling :class:`MLADSparkForCausalLM`, not a + subclass of this. Reference: arXiv 2607.05147; deepseek-ai/DeepSpec. """ @@ -2689,14 +2685,9 @@ class MLADSparkForCausalLM(_DSparkHeadMixin, DFlashForCausalLM): ``spec_config.attention_backend`` selects among this class's own three block-decode implementations, not among the shared DFlash worker op sets: - ``CUTEDSL`` one cute-dsl pass with ``kv_bounds``, no fixup - (``_cute_dsl_block_decode``). Paged only. What ``AUTO`` picks - on a build that can run it. - ``TRTLLM`` flashinfer's trtllm-gen absorbed-MLA paged decode plus the - upper-triangle fixup (``_mla_paged_attention``). - ``VANILLA`` the eager torch reference. It is also what the other two - degrade to on the unpaged arena, where there are no page - tables to hand a kernel. + ``CUTEDSL`` one cute-dsl pass with ``kv_bounds``, no fixup; paged only, AUTO's pick. + ``TRTLLM`` trtllm-gen absorbed-MLA paged decode + upper-triangle fixup. + ``VANILLA`` the eager torch reference; where the other two land unpaged. Reference: arXiv 2607.05147; deepseek-ai/DeepSpec. """ @@ -2851,11 +2842,10 @@ def _assert_cute_dsl_can_implement(self, batch, block_size, page_size, dtype) -> """Reject a shape the kernel cannot serve, here, naming the reason. Without this the rejection is silent until it is not: ``get_valid_tactics`` - returns [] (cute_dsl_custom_ops.py:8508), the AutoTuner falls back to its - -1 sentinel, and ``forward`` compiles an unvalidated config -- so a draft - pool with ``tokens_per_block=256`` (``128 % 256 != 0``, - mla_decode_fp16.py:3798) surfaces as a CUTLASS error out of the fifth - drafter layer rather than as the configuration problem it is. + returns [] (cute_dsl_custom_ops.py:8508), the AutoTuner falls back to its -1 + sentinel, and ``forward`` compiles an unvalidated config -- so a draft pool + with ``tokens_per_block=256`` (``128 % 256 != 0``, mla_decode_fp16.py:3798) + surfaces as a CUTLASS error out of the fifth drafter layer. Page size is the only operand here that is not fixed by the checkpoint, and it comes from the draft KV pool, so this cannot move to __init__. @@ -2904,13 +2894,11 @@ def _assert_cute_dsl_can_implement(self, batch, block_size, page_size, dtype) -> def _cute_dsl_block_decode(self, query, layer_cache, fixup, block_size): """One cute-dsl MLA decode over context AND block, no fixup. - The trtllm-gen path has to split the work: its q_len semantics are - causal, so query i cannot see block positions j > i, and the strict - upper triangle is recovered eagerly afterwards and merged by LSE. The - cute-dsl kernel takes ``kv_bounds`` -- a per-query-token KV length that - REPLACES the implicit causal bound (mla_decode_fp16.py:400-407) -- so - filling it with ctx_len + block_size makes every block position visible - to every other in one pass and deletes the fixup outright. + The trtllm-gen path splits the work: its q_len semantics are causal, so + query i cannot see block positions j > i, and the strict upper triangle is + recovered eagerly and merged by LSE. The cute-dsl kernel takes ``kv_bounds``, + a per-query-token KV length REPLACING the implicit causal bound + (mla_decode_fp16.py:400-407), so one pass sees all of it and the fixup goes. No fallback: what the kernel cannot serve raises. Build-level availability is checked at construction through @@ -2974,12 +2962,11 @@ def _mla_paged_attention(self, query, blk_kv, layer_cache, fixup, block_size, ma full visibility). Its LSE is base 2 (also measured: exact against ``log2(sum exp)``, 1.84 off against the natural log). - The block's own latents go into the pool first, at ``ctx_len + j`` - (see _MLABlockFixup). With ``seq_lens = ctx_len + block_size`` the kernel - covers context AND block, and in-block causality leaves only the block's - strict upper triangle -- 8 keys, attended eagerly against ``blk_kv`` and - merged by log-sum-exp. The room those writes need is reserved by - ``dflash_allocated_ctx_limit``. + The block's own latents go into the pool first, at ``ctx_len + j`` (see + _MLABlockFixup). With ``seq_lens = ctx_len + block_size`` the kernel covers + context AND block, and in-block causality leaves only the strict upper + triangle -- 8 keys, attended eagerly against ``blk_kv`` and merged by + log-sum-exp. ``dflash_allocated_ctx_limit`` reserves the room. Returns the latent-space output ``[B, block, heads, kv_lora_rank]``. """ @@ -3044,12 +3031,11 @@ def _mla_paged_attention(self, query, blk_kv, layer_cache, fixup, block_size, ma def _mla_freqs_cis(self, device): """Adjacent-pair YaRN table, cached. See build_dspark_mla_yarn_freqs_cis. - Sized from the ceiling the worker publishes once it knows the runtime - one (_runtime_position_ceiling, set in DFlashWorker._lazy_init_ctx_buffers - before any forward), because that is the value ctx_len is clamped to and - it is strictly above model_config.max_seq_len whenever spec decoding is - on. The config-derived cap is the fallback for direct construction in - tests, where no worker runs. + Sized from the ceiling the worker publishes once it knows the runtime one + (_runtime_position_ceiling, set in DFlashWorker._lazy_init_ctx_buffers before + any forward), because that is what ctx_len is clamped to and it exceeds + model_config.max_seq_len whenever spec decoding is on. The config cap is the + fallback for direct construction in tests. """ rope = dict(self._mla_rope_params) runtime_cap = self._runtime_position_ceiling diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index f4c4eebe8425..f25ed45a364d 100644 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -96,12 +96,11 @@ def markov_prev_embeddings(prev_tokens: torch.Tensor, markov_w1: torch.Tensor) -> torch.Tensor: """``markov_w1[prev_tokens]`` with out-of-vocab anchors masked to zero. - The anchor is the last accepted token, which on the one-model rejection - path comes from flashinfer's ``chain_speculative_sampling``: it pads - non-accepted positions with ``-1`` and returns an out-of-range id for a row - whose ``relu(target - draft)`` residual has no mass. Mask like - ``modules/embedding.py`` does for the target embedding, so such a row - contributes no bias instead of tripping a device-side assert. + The anchor is the last accepted token, which on the one-model rejection path + comes from flashinfer's ``chain_speculative_sampling``: it pads non-accepted + positions with ``-1`` and returns an out-of-range id for a row whose + ``relu(target - draft)`` residual has no mass. Mask like modules/embedding.py + does, so such a row contributes no bias instead of tripping a device assert. Args: prev_tokens: previous token ids (draft vocab), any shape. diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 032804ae3d43..94dc2c8ebfe6 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1678,19 +1678,17 @@ def _get_effective_draft_config(self) -> ModelConfig: def _get_draft_kv_model_config(self) -> ModelConfig: """The draft ModelConfig describing the KV pool as it is ALLOCATED. - The args-level ``kv_cache_config.dtype`` sync stamps the TARGET's fp8 - KV algo onto every loaded model, including a standalone drafter. The - drafter stores and reads its pool in its weights dtype (DFlash - validates a bf16 pool and otherwise falls back to the max_seq_len-dense - private arena, which OOMs at long context), so the pool dtype must - follow the drafter, not the target. - - Every consumer of draft KV bytes must go through here. If the budget - split and the allocation read different dtypes, the split charges fp8 - bytes for a bf16 pool and the draft manager gets HALF the target's - tokens. The capacity scheduler admits on the target pool alone, so the - draft pool cannot backpressure -- past ~50% target utilization it raises - "Draft KV cache context resize failed", fatal to every rank. + The args-level ``kv_cache_config.dtype`` sync stamps the TARGET's fp8 KV + algo onto every loaded model, including a standalone drafter. The drafter + stores and reads its pool in its weights dtype (DFlash validates a bf16 pool + and otherwise falls back to the max_seq_len-dense private arena, which OOMs at + long context), so the pool dtype must follow the drafter. + + Every consumer of draft KV bytes must go through here. If the budget split + and the allocation read different dtypes, the split charges fp8 bytes for a + bf16 pool and the draft manager gets HALF the target's tokens. The capacity + scheduler admits on the target pool alone, so past ~50% target utilization it + raises "Draft KV cache context resize failed", fatal to every rank. """ effective_draft_config = self._get_effective_draft_config() # Narrower than is_external_drafter(), matching diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index e367a7768013..f07008e86aac 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -3665,13 +3665,12 @@ def _mirror_draft_kv_cache(self, req: LlmRequest): worker receives prompt KV rather than prefilling it. Returns None when the IndexMapper is saturated, which neither branch - survives. copy_batch_block_offsets() runs later in the SAME iteration - and feeds every id in the batch -- context ids included, see - IndexMapper::getCopyIndex -- to getIndex(), which TLLM_CHECKs on an - unmapped id (kvCacheManagerV2Utils.cpp). So the caller's `continue` on - the context path does not defer the request to a later iteration; - there is no later iteration. That predates the mirror refactor and is - left as is here; only the claim about it is corrected. + survives. copy_batch_block_offsets() runs later in the SAME iteration and + feeds every id in the batch -- context ids included, IndexMapper::getCopyIndex + -- to getIndex(), which TLLM_CHECKs on an unmapped id + (kvCacheManagerV2Utils.cpp), so the caller's `continue` defers to nothing. + + Pre-existing, and left as is here; only the claim about it is corrected. """ kv_cache = self.kv_cache_map.get(req.py_request_id) if kv_cache is not None: @@ -3693,12 +3692,12 @@ def _draft_pool_diagnostic(self) -> str: The draft manager mirrors the target's tokens but is sized from its own byte budget, and the capacity scheduler admits on the TARGET pool alone - (`scheduler_v2` touches `draft_kv_cache_manager` only to suspend/free). - A draft pool smaller in tokens than the target therefore cannot - backpressure -- it can only raise, and the raise kills every rank. When - that happens the first question is always "how big was the draft pool - and how full was it", so answer it in the message rather than leaving - it to post-hoc arithmetic over the budget-split log line. + (`scheduler_v2` touches `draft_kv_cache_manager` only to suspend/free). A + draft pool smaller in tokens than the target cannot backpressure -- it can + only raise, and the raise kills every rank. + + The first question is always how big the draft pool was and how full, so the + message answers it rather than leaving post-hoc arithmetic over the split log. """ live = sum(c.capacity for c in self.kv_cache_map.values()) diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index b6bab90e3e5a..7007e000045d 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -2077,9 +2077,8 @@ def _zero_padding_rows(logits: torch.Tensor, spec_metadata, A padding request decodes from an uninitialized KV/hidden state, so its logits can be non-finite and its probs degenerate. flashinfer's ``chain_speculative_sampling`` then rejects at a position whose - ``relu(target - draft)`` residual has no mass, where it reads an - uninitialized shared-memory slot and emits an out-of-range token id. - Zeroed logits give a uniform -- and therefore legal -- distribution. + ``relu(target - draft)`` residual has no mass, reads an uninitialized + shared-memory slot and emits an out-of-range token id. Zeroed logits are legal. Padding requests are the ones ``populate_sampling_params_for_one_model`` routed to ``dummy_slot_row`` (``py_seq_slot is None``), so this needs no diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index f579e377c9e2..5e7a7417f406 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3136,26 +3136,15 @@ class DSparkDecodingConfig(DecodingBaseConfig): attention_backend: Literal["AUTO", "VANILLA", "TRTLLM", "CUTEDSL"] = Field( default="AUTO", description= - "Attention backend for the block decode of a standalone DSpark drafter " - "(one shipped as its own checkpoint rather than inside the target's " - "mtp.* namespace). Ignored by the embedded DeepSeek-V4-Pro draft, which " - "uses its own captured-context attention. This is independent of the " - "backend used to construct the drafter's standard attention modules. " - "AUTO picks per drafter family: TRTLLM for both a GQA- and an " - "MLA-backboned drafter, though the two resolve that name to different " - "kernels. A GQA drafter degrades to VANILLA when its kernel is " - "unavailable; an MLA drafter raises instead, since a build that cannot " - "run its kernel cannot hold the target either. " - "TRTLLM requires FlashInfer and an NVIDIA Blackwell " - "GPU with SM100 or SM103; for a GQA backbone it uses generated FMHA " - "kernels with a private paged context cache, and for an MLA backbone " - "the absorbed-MLA paged decode plus a block-local fixup. VANILLA uses " - "FlashAttention with a contiguous cache on a GQA backbone and the eager " - "torch reference on an MLA one. CUTEDSL is MLA-only: one cute-dsl pass " - "over context and block that replaces the fixup. It is selectable but " - "not a default, because it needs a cute-dsl MLA decode taking per-token " - "kv_bounds that is not upstream yet; requesting it on a build without " - "that kernel raises with the reason.") + "Block-decode attention backend for a standalone DSpark drafter (one " + "shipped as its own checkpoint, not inside the target's mtp.* " + "namespace). Ignored by the embedded DeepSeek-V4-Pro draft. Independent " + "of the backend that builds the drafter's own attention modules.\n\n" + "AUTO resolves per drafter family and is right unless you are pinning a " + "kernel: a GQA backbone degrades when its kernel is missing, an MLA one " + "raises. TRTLLM needs FlashInfer and SM100/SM103. CUTEDSL is MLA-only " + "and needs a cute-dsl MLA decode taking per-token kv_bounds that is not " + "upstream yet. Which kernel each name selects: MLADSparkForCausalLM.") @model_validator(mode="after") def set_max_total_draft_tokens(self): diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py index 555fab8775cf..9197f3dacc3b 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py @@ -1033,13 +1033,14 @@ class TestExternalDrafterKvDtype: """The draft budget must be charged at the dtype the draft pool is ALLOCATED in. ``kv_cache_config.dtype: fp8`` stamps the target's fp8 KV algo onto an - external drafter's ModelConfig, but the drafter keeps a bf16 pool (dflash - rejects an fp8 pool and falls back to a max_seq_len-dense private arena). - The allocation path dropped the inherited algo; the cost path did not, so - the split charged 2880 B/token for a pool costing 5760, handed the draft - manager half the tokens the target got, and the GEN worker died in - ``_prepare_draft_resources`` at ~50% target utilization -- fatal to every - rank, because the capacity scheduler admits on the target pool alone. + external drafter's ModelConfig, but the drafter keeps a bf16 pool. The + allocation path dropped the inherited algo; the cost path did not, so the + split charged 2880 B/token for a pool costing 5760 and handed the draft + manager half the tokens the target got. + + The GEN worker then died in ``_prepare_draft_resources`` at ~50% target + utilization, fatal to every rank: the capacity scheduler admits on the target + pool alone. """ # MLA drafter: kv_lora_rank 512 + qk_rope_head_dim 64 = 576, kv_factor 1, diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py index eb9aeff4be25..70a6688c4fa2 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py @@ -968,13 +968,11 @@ def _yarn_mscale_sq(): def _run_mla_block_decode(drafter, requests, paged=False): """Drive one block decode over a BATCH, optionally through the paged cache. - ``requests`` is a list of ``(captured, noise_embed)`` pairs; each one's - context length is its own ``captured.shape[0]``, and they are deliberately - unequal. B >= 2 with unequal lengths is what makes the paged path testable - at all: at B == 1 the page table's stride is 1 either way, so - ``.t()`` and ``.t().contiguous()`` are indistinguishable, and - ``kv_bounds`` built with ``repeat`` instead of ``repeat_interleave`` gives - the same vector. Both are silent at B == 1 and wrong at B == 2. + ``requests`` is a list of ``(captured, noise_embed)`` pairs with deliberately + unequal context lengths. B >= 2 with unequal lengths is what makes the paged + path testable: at B == 1 the page table's stride is 1 either way, so ``.t()`` + and ``.t().contiguous()`` are indistinguishable, and ``kv_bounds`` built with + ``repeat`` instead of ``repeat_interleave`` gives the same vector. ``paged=False`` keeps the dense arena the eager reference path uses. ``paged=True`` hands ``dflash_forward`` a page table, which is what selects @@ -1058,11 +1056,9 @@ def test_mla_dspark_block_decode_matches_unabsorbed_reference(monkeypatch, paged The paged variants write the draft block into the pool and read it back through a kernel, so a bug in that write shows up here as a parity failure - rather than as lower acceptance length. Both kernels are compared against - the same reference: they differ in whether the block's strict upper triangle - comes from a separate fixup (TRTLLM) or from kv_bounds inside one pass - (CUTEDSL), and that is precisely the part a reference can catch and - acceptance cannot. + rather than as lower acceptance length. Both kernels meet the same reference; + they differ in whether the strict upper triangle comes from a separate fixup + (TRTLLM) or from kv_bounds in one pass (CUTEDSL). Two facts say the batch is wired right rather than the tolerance being generous: ctx 20 is IDENTICAL across the eager and paged paths (a wrong page @@ -1179,10 +1175,9 @@ def test_mla_dspark_backend_selects_its_own_block_decode(backend, variant): Neither worker op set may be loaded -- otherwise a drafter that never calls them drags in an optional dependency, and the worker's per-backend shape - checks (which the absorbed 64:1 / head_dim 576 shape fails) would bind on a - path that does not use them. What the field DOES drive is - _mla_block_decode_variant, and getting that wrong is silent: every variant - returns the same shape and only acceptance would move. + checks (which the absorbed 64:1 / head_dim 576 shape fails) bind on a path + that does not use them. What the field DOES drive is _mla_block_decode_variant, + and getting that wrong is silent: every variant returns the same shape. """ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_dspark import ( @@ -1212,11 +1207,10 @@ def test_mla_dspark_cute_dsl_takes_a_single_page_batch_one(): """One page and one request -- the page table the runtime used to reject. (max_blocks, B) == (1, 1) has stride (1, 1) and no dimension with size > 1, - so CuTe DSL's layout deduction refused to pick a leading dimension even - though, with both extents 1, the choice cannot change an address. It is - reachable whenever the draft pool is one page deep and a single request is - resident, which is a first-draft-step shape, not an exotic one -- and since - AUTO resolves to CUTEDSL it would be the default path that died. + so CuTe DSL's layout deduction refused to pick a leading dimension even though + with both extents 1 the choice cannot change an address. Reachable whenever + the draft pool is one page deep and a single request is resident -- a + first-draft-step shape, and since AUTO resolves to CUTEDSL, the default path. The parity bound is the same one the batched test uses; a wrongly deduced axis would show up as a magnitude error, not a last-bit one. @@ -1366,13 +1360,11 @@ def test_mla_dspark_rope_conventions_agree_on_scores(): def test_mla_block_fixup_stays_inside_the_allocation(): """A context that fills its allocation must still leave the block room. - The bound comes from dflash.py, so it is called here rather than restated: - a test that computed `allocated - block_size` itself would still pass - against a production path truncating to `allocated`. The MLA writes the - block's own latents at ctx_len..ctx_len+block_size, so that regression puts - the first of them on the first unallocated page, whose block-table entry - _refresh_ctx_block_tables clamped from a negative placeholder to 0 -- - another request's block. Silent cross-request corruption, not a fault. + The bound comes from dflash.py, so it is called here rather than restated: a + test computing `allocated - block_size` itself would still pass against a path + truncating to `allocated`. The MLA writes the block's latents at + ctx_len..ctx_len+block_size, so that regression puts the first on the first + unallocated page, clamped from a negative placeholder to 0 -- another request's. """ from tensorrt_llm._torch.models.modeling_dspark import _build_mla_block_fixup from tensorrt_llm._torch.speculative.dflash import dflash_allocated_ctx_limit diff --git a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py index 494127a28ae3..74f1448e31ab 100644 --- a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py +++ b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py @@ -212,11 +212,10 @@ def test_fused_dspark_rmsnorm_rope_norm_dim(split_norm): latent, leave k_pe raw. Driven through ``_rmsnorm_rope_batched`` rather than the custom op, because - the plumbing under test is the dispatcher's -- it has to forward norm_dim to - the support predicate and to the kernel. The predicate is asserted first so - a fused path that silently stopped applying never reads as a pass; that is - what made the original home of this test (hw_agnostic, mapped to CPU and - H100 only) vacuous, since is_sm_100f() is false there. + the plumbing under test is the dispatcher's -- it forwards norm_dim to the + support predicate and to the kernel. The predicate is asserted first so a + fused path that silently stopped applying never reads as a pass; that is what + made this test's original home (hw_agnostic, CPU/H100 only) vacuous. """ from tensorrt_llm._torch.custom_ops.dspark_rmsnorm_rope_custom_op import ( is_fused_dspark_rmsnorm_rope_supported, From e54aea2c95e3ce0bf540efe13279babd948cfef5 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Thu, 17 Sep 2026 01:19:03 -0700 Subject: [PATCH 19/25] [None][fix] Derive KDA CuTe argument alignment Ported from c6c985dcd1 rather than cherry-picked: that commit sits on a branch whose copy of this op is 341 lines larger (quantize_output, fuse_output_norm, output_scale), so its hunks for those have nothing to land on here. The mechanism is verbatim; only the call-site list differs. A per-layer beta_cache view is not 16-byte aligned in general, and the CuTe bridge was told it was, so KDA CTX workers died at kda_mtp_decode with "Tensor data pointer is not aligned to 16 bytes". _beta_cache_assumed_align derives it from the layer span, int32 metadata declares 4, and assumed_align joins the compile cache key. Signed-off-by: Zhenhuan Chen --- .../cute_dsl_kimi_k3_kda_mtp_ops.py | 157 ++++++++++-------- .../test_kda_mtp_decode_cute_parity.py | 139 ++++++++++++++-- 2 files changed, 212 insertions(+), 84 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py index a5c2501bca90..c38c427eeb63 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py @@ -53,6 +53,7 @@ (2, 12, 32), num_spec == 2``; other shapes compile the general variant. """ +from math import gcd from typing import Optional, Tuple import torch @@ -302,30 +303,50 @@ def _fits_32bit_stride(tensor: torch.Tensor) -> bool: return True -def _from_dlpack_arg(tensor: torch.Tensor): +def _from_dlpack_arg(tensor: torch.Tensor, *, assumed_align: int = 16): return from_dlpack( tensor, - assumed_align=16, + assumed_align=assumed_align, use_32bit_stride=_fits_32bit_stride(tensor), ) -def _dlpack_arg(tensor: torch.Tensor): +def _beta_cache_assumed_align(beta_cache: torch.Tensor) -> int: + """Return the alignment shared by KDA per-layer beta-cache views. + + The producer allocates ``[layers, slots, num_spec, local_heads]``. After + selecting a layer, ``slots * stride(0)`` is the physical layer span even + when the head rows are padded. Combine that span with the current pointer + and dtype, capped at the CuTe bridge's useful 16-byte guarantee. + """ + layer_span_bytes = beta_cache.shape[0] * beta_cache.stride(0) * beta_cache.element_size() + return gcd(16, beta_cache.data_ptr(), layer_span_bytes) + + +def _dlpack_arg(tensor: torch.Tensor, *, assumed_align: int): + # Alignment is deliberately mandatory: layout dynamism does not imply + # arbitrary pointer alignment. Each call site must state the guarantee + # provided by that tensor's producer and view pattern. + arg = _from_dlpack_arg(tensor, assumed_align=assumed_align) for dim, stride in enumerate(tensor.stride()): if stride == 1: - return _from_dlpack_arg(tensor).mark_layout_dynamic(dim) - return _from_dlpack_arg(tensor).mark_layout_dynamic() + return arg.mark_layout_dynamic(dim) + return arg.mark_layout_dynamic() -def _layout_key(tensor: torch.Tensor, dynamic_layout: bool = False): - arg = _dlpack_arg(tensor) if dynamic_layout else _from_dlpack_arg(tensor) +def _layout_key(tensor: torch.Tensor, dynamic_layout: bool = False, *, assumed_align: int = 16): + arg = ( + _dlpack_arg(tensor, assumed_align=assumed_align) + if dynamic_layout + else _from_dlpack_arg(tensor, assumed_align=assumed_align) + ) shape_mask = arg.dynamic_shapes_mask stride_mask = arg.dynamic_strides_mask shape = tuple(None if dynamic else size for size, dynamic in zip(tensor.shape, shape_mask)) stride = tuple( None if dynamic else value for value, dynamic in zip(tensor.stride(), stride_mask) ) - return (tensor.dtype, shape, stride, _fits_32bit_stride(tensor)) + return (tensor.dtype, shape, stride, _fits_32bit_stride(tensor), assumed_align) # (device_index, enabled) -> persistent int32 [1] control tensor. Keys are @@ -460,9 +481,6 @@ def kda_mtp_decode_impl( out = torch.zeros(1, T_total, HV, V_dim, dtype=x_q.dtype, device=x_q.device) if num_accepted_tokens.dtype != torch.int32: num_accepted_tokens = num_accepted_tokens.to(torch.int32) - if ssm_state_indices.data_ptr() % 16 != 0: - raise ValueError("ssm_state_indices must be 16-byte aligned before CuTe DLPack conversion") - _require_stride_layout( x_q=x_q, x_k=x_k, @@ -531,6 +549,7 @@ def kda_mtp_decode_impl( "buffer before enabling PROFILE_STAGES" ) stage_timing_arg = out + beta_cache_assumed_align = _beta_cache_assumed_align(beta_cache) key = ( x_q.dtype, @@ -544,9 +563,9 @@ def kda_mtp_decode_impl( lower_bound, use_flat_layout, _layout_key(h0_arg), - _layout_key(x_q_arg, dynamic_layout=True), - _layout_key(x_k_arg, dynamic_layout=True), - _layout_key(x_v_arg, dynamic_layout=True), + _layout_key(x_q_arg, dynamic_layout=True, assumed_align=16), + _layout_key(x_k_arg, dynamic_layout=True, assumed_align=16), + _layout_key(x_v_arg, dynamic_layout=True, assumed_align=16), _layout_key(w_q), _layout_key(w_k), _layout_key(w_v), @@ -554,16 +573,16 @@ def kda_mtp_decode_impl( _layout_key(cs_k), _layout_key(cs_v), _layout_key(A_log), - _layout_key(g, dynamic_layout=True), + _layout_key(g, dynamic_layout=True, assumed_align=16), _layout_key(dt_bias), - _layout_key(beta, dynamic_layout=True), - _layout_key(out, dynamic_layout=True), + _layout_key(beta, dynamic_layout=True, assumed_align=16), + _layout_key(out, dynamic_layout=True, assumed_align=16), _layout_key(qkg_cache), _layout_key(v_cache), - _layout_key(beta_cache), - _layout_key(ssm_state_indices, dynamic_layout=True), - _layout_key(cu_seqlens, dynamic_layout=True), - _layout_key(num_accepted_tokens, dynamic_layout=True), + _layout_key(beta_cache, assumed_align=beta_cache_assumed_align), + _layout_key(ssm_state_indices, dynamic_layout=True, assumed_align=4), + _layout_key(cu_seqlens, dynamic_layout=True, assumed_align=4), + _layout_key(num_accepted_tokens, dynamic_layout=True, assumed_align=4), use_setmaxreg, use_regular_metadata, use_reg_q_weights, @@ -579,30 +598,30 @@ def kda_mtp_decode_impl( ) _compiled_cache[key] = cute.compile( _run_kda_decode_mtp, - _from_dlpack_arg(h0_arg), - _dlpack_arg(x_q_arg), - _dlpack_arg(x_k_arg), - _dlpack_arg(x_v_arg), - _from_dlpack_arg(w_q), - _from_dlpack_arg(w_k), - _from_dlpack_arg(w_v), - _from_dlpack_arg(cs_q), - _from_dlpack_arg(cs_k), - _from_dlpack_arg(cs_v), - _from_dlpack_arg(A_log), - _dlpack_arg(g), - _from_dlpack_arg(dt_bias), - _dlpack_arg(beta), - _dlpack_arg(out), - _from_dlpack_arg(h0_arg), - _from_dlpack_arg(qkg_cache), - _from_dlpack_arg(v_cache), - _from_dlpack_arg(beta_cache), - _dlpack_arg(stage_timing_arg), - _dlpack_arg(ssm_state_indices), - _dlpack_arg(cu_seqlens), - _dlpack_arg(num_accepted_tokens), - _from_dlpack_arg(precompute_control), + _from_dlpack_arg(h0_arg, assumed_align=16), + _dlpack_arg(x_q_arg, assumed_align=16), + _dlpack_arg(x_k_arg, assumed_align=16), + _dlpack_arg(x_v_arg, assumed_align=16), + _from_dlpack_arg(w_q, assumed_align=16), + _from_dlpack_arg(w_k, assumed_align=16), + _from_dlpack_arg(w_v, assumed_align=16), + _from_dlpack_arg(cs_q, assumed_align=16), + _from_dlpack_arg(cs_k, assumed_align=16), + _from_dlpack_arg(cs_v, assumed_align=16), + _from_dlpack_arg(A_log, assumed_align=16), + _dlpack_arg(g, assumed_align=16), + _from_dlpack_arg(dt_bias, assumed_align=16), + _dlpack_arg(beta, assumed_align=16), + _dlpack_arg(out, assumed_align=16), + _from_dlpack_arg(h0_arg, assumed_align=16), + _from_dlpack_arg(qkg_cache, assumed_align=16), + _from_dlpack_arg(v_cache, assumed_align=16), + _from_dlpack_arg(beta_cache, assumed_align=beta_cache_assumed_align), + _dlpack_arg(stage_timing_arg, assumed_align=16), + _dlpack_arg(ssm_state_indices, assumed_align=4), + _dlpack_arg(cu_seqlens, assumed_align=4), + _dlpack_arg(num_accepted_tokens, assumed_align=4), + _from_dlpack_arg(precompute_control, assumed_align=16), scale=scale, HV=HV, K=K, @@ -624,30 +643,30 @@ def kda_mtp_decode_impl( ) _compiled_cache[key]( - _dlpack_arg(h0_arg), - _dlpack_arg(x_q_arg), - _dlpack_arg(x_k_arg), - _dlpack_arg(x_v_arg), - _dlpack_arg(w_q), - _dlpack_arg(w_k), - _dlpack_arg(w_v), - _dlpack_arg(cs_q), - _dlpack_arg(cs_k), - _dlpack_arg(cs_v), - _dlpack_arg(A_log), - _dlpack_arg(g), - _dlpack_arg(dt_bias), - _dlpack_arg(beta), - _dlpack_arg(out), - _dlpack_arg(h0_arg), - _dlpack_arg(qkg_cache), - _dlpack_arg(v_cache), - _dlpack_arg(beta_cache), - _dlpack_arg(stage_timing_arg), - _dlpack_arg(ssm_state_indices), - _dlpack_arg(cu_seqlens), - _dlpack_arg(num_accepted_tokens), - _dlpack_arg(precompute_control), + _dlpack_arg(h0_arg, assumed_align=16), + _dlpack_arg(x_q_arg, assumed_align=16), + _dlpack_arg(x_k_arg, assumed_align=16), + _dlpack_arg(x_v_arg, assumed_align=16), + _dlpack_arg(w_q, assumed_align=16), + _dlpack_arg(w_k, assumed_align=16), + _dlpack_arg(w_v, assumed_align=16), + _dlpack_arg(cs_q, assumed_align=16), + _dlpack_arg(cs_k, assumed_align=16), + _dlpack_arg(cs_v, assumed_align=16), + _dlpack_arg(A_log, assumed_align=16), + _dlpack_arg(g, assumed_align=16), + _dlpack_arg(dt_bias, assumed_align=16), + _dlpack_arg(beta, assumed_align=16), + _dlpack_arg(out, assumed_align=16), + _dlpack_arg(h0_arg, assumed_align=16), + _dlpack_arg(qkg_cache, assumed_align=16), + _dlpack_arg(v_cache, assumed_align=16), + _dlpack_arg(beta_cache, assumed_align=beta_cache_assumed_align), + _dlpack_arg(stage_timing_arg, assumed_align=16), + _dlpack_arg(ssm_state_indices, assumed_align=4), + _dlpack_arg(cu_seqlens, assumed_align=4), + _dlpack_arg(num_accepted_tokens, assumed_align=4), + _dlpack_arg(precompute_control, assumed_align=16), N, stream, ) diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py index 004998a4fffb..64747b380a42 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py @@ -260,7 +260,7 @@ def cpu_reference(data): } -def cute_run(data, zero_accepted_hint=False): +def cute_run(data, zero_accepted_hint=False, beta_cache_override=None): """Run the in-tree op on cloned caches; return the drop-format dict.""" import tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops # noqa: F401 @@ -275,7 +275,7 @@ def cute_run(data, zero_accepted_hint=False): cs[name] = dst qkg_cache = data["qkg_cache"].clone() v_cache = data["v_cache"].clone() - beta_cache = data["beta_cache"].clone() + beta_cache = data["beta_cache"].clone() if beta_cache_override is None else beta_cache_override out = torch.ops.trtllm.kda_mtp_decode( x_q=data["x_q"], x_k=data["x_k"], @@ -534,22 +534,131 @@ def test_zero_accepted_hint_variant(B, H): _assert_close(f"{name}(fast vs general)", fast[name], general[name], atol=1e-5) -def test_misaligned_state_indices_rejected_after_aligned_warmup(): - """The op enforces the alignment contract supplied by metadata prep.""" - data = make_conv_data(B=1, H=6, M=M, seed=13) - cute_run(data) +@pytest.mark.parametrize("field", ["ssm_state_indices", "cu_seqlens", "num_accepted_tokens"]) +def test_misaligned_scalar_metadata_after_aligned_warmup(field): + """Every scalar CuTe metadata argument accepts an offset int32 view. - index_storage = torch.empty(2, dtype=torch.int32, device="cuda") - index_storage[0] = -1 - index_storage[1:].copy_(data["ssm_state_indices"]) - misaligned_indices = index_storage[1:] - assert misaligned_indices.is_contiguous() - assert misaligned_indices.data_ptr() % 16 != 0 + Mixed context/generation batches hand the op a view into the shared buffer + that starts after one to three context rows: contiguous, int32, and not + 16-byte aligned. The op states 4-byte alignment for these, so the CuTe + bridge no longer has to be told a guarantee the caller cannot make. + """ + data = make_conv_data(B=1, H=6, M=M, seed=17) + expected = cute_run(data) + + source = data[field] + storage = torch.empty(source.numel() + 1, dtype=torch.int32, device="cuda") + storage[0] = -1 + storage[1:].copy_(source) + misaligned = storage[1:] + assert misaligned.is_contiguous() + assert misaligned.data_ptr() % 16 != 0 misaligned_data = dict(data) - misaligned_data["ssm_state_indices"] = misaligned_indices - with pytest.raises(ValueError, match="16-byte aligned"): - cute_run(misaligned_data) + misaligned_data[field] = misaligned + actual = cute_run(misaligned_data) + + for name in ("out", "recurrent_state", "qkg_cache", "v_cache", "beta_cache"): + _assert_close( + f"{name}({field} misaligned vs aligned)", actual[name], expected[name], atol=1e-5 + ) + + +@pytest.mark.parametrize( + "attention_mode,parallel_size,dtype,pool_size,num_spec,expected", + ( + pytest.param("dep", 8, torch.float32, 1177, 5, 16, id="dep8-mtp5-fp32"), + pytest.param("dep", 16, torch.float32, 1177, 7, 16, id="dep16-mtp7-fp32"), + pytest.param("tep", 8, torch.float32, 1177, 7, 16, id="tep8-mtp7-fp32"), + pytest.param("tep", 16, torch.float32, 1177, 2, 16, id="tep16-mtp2-fp32"), + pytest.param("tep", 16, torch.float32, 1177, 5, 8, id="tep16-mtp5-fp32"), + pytest.param("tep", 16, torch.float32, 1177, 7, 8, id="tep16-mtp7-fp32"), + pytest.param("tep", 16, torch.float32, 248, 7, 16, id="tep16-mtp7-even-pool"), + pytest.param("tep", 32, torch.float32, 1177, 7, 4, id="tep32-mtp7-fp32"), + pytest.param("dep", 16, torch.bfloat16, 1177, 7, 16, id="dep16-mtp7-bf16"), + pytest.param("tep", 8, torch.bfloat16, 1177, 7, 8, id="tep8-mtp7-bf16"), + pytest.param("tep", 16, torch.bfloat16, 1177, 7, 4, id="tep16-mtp7-bf16"), + pytest.param("tep", 32, torch.bfloat16, 1177, 7, 2, id="tep32-mtp7-bf16"), + ), +) +def test_beta_cache_alignment_is_derived_from_layout_and_dtype( + attention_mode, parallel_size, dtype, pool_size, num_spec, expected +): + from tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops import ( + _beta_cache_assumed_align, + ) + + local_heads = 96 if attention_mode == "dep" else 96 // parallel_size + parent = torch.empty(2, pool_size, num_spec, local_heads, dtype=dtype, device="cuda") + beta_cache = parent[0] + + assert _beta_cache_assumed_align(beta_cache) == expected + + +def test_beta_cache_alignment_uses_padded_physical_stride(): + from tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops import ( + _beta_cache_assumed_align, + ) + + parent = torch.empty(2, 1177, 7, 8, dtype=torch.float32, device="cuda") + beta_cache = parent[0, ..., :6] + + assert beta_cache.shape == (1177, 7, 6) + assert beta_cache.stride() == (56, 8, 1) + assert _beta_cache_assumed_align(beta_cache) == 16 + + +@pytest.mark.parametrize( + "attention_mode,parallel_size,num_spec,expected_alignment", + ( + pytest.param("dep", 16, 7, 16, id="dep16-mtp7"), + pytest.param("tep", 8, 7, 16, id="tep8-mtp7"), + pytest.param("tep", 16, 5, 8, id="tep16-mtp5"), + pytest.param("tep", 16, 7, 8, id="tep16-mtp7"), + pytest.param("tep", 32, 7, 4, id="tep32-mtp7"), + ), +) +def test_beta_cache_sibling_layer_after_aligned_warmup( + attention_mode, parallel_size, num_spec, expected_alignment +): + """A cached kernel accepts the next real per-layer beta-cache slice.""" + from tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops import ( + _beta_cache_assumed_align, + ) + + local_heads = 96 if attention_mode == "dep" else 96 // parallel_size + data = make_conv_data(B=1, H=local_heads, M=num_spec, seed=29) + + parent = torch.zeros(2, 1177, num_spec, local_heads, dtype=torch.float32, device="cuda") + layer0, layer1 = parent.unbind(0) + layer0[0].copy_(data["beta_cache"][0]) + layer1[0].copy_(data["beta_cache"][0]) + + assert _beta_cache_assumed_align(layer0) == expected_alignment + assert layer0.data_ptr() % 16 == 0 + assert layer1.data_ptr() % expected_alignment == 0 + if expected_alignment < 16: + assert layer1.data_ptr() % (2 * expected_alignment) != 0 + + expected = cute_run(data, beta_cache_override=layer0) + actual = cute_run(data, beta_cache_override=layer1) + + for name in ( + "out", + "recurrent_state", + "qkg_cache", + "v_cache", + "beta_cache", + "cs_q", + "cs_k", + "cs_v", + ): + _assert_close( + f"{name}({attention_mode}{parallel_size}-mtp{num_spec})", + actual[name], + expected[name], + atol=1e-5, + ) if __name__ == "__main__": From 494b3cfe65d42697ab5b236e9e10eb0b240d1766 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Wed, 16 Sep 2026 21:13:31 -0700 Subject: [PATCH 20/25] [None][feat] Join DFlash/DSpark to the paired draft KV reuse protocol precompute_context_kv writes the drafter's own post-norm post-RoPE K/V one entry per context token, and the target hidden at i depends on exactly tokens [0, i] -- the same dependency face as the target's KV. So the span is 0 like PARD: raw-prompt keys describe the draft pool and no chunk-tail lookahead token is needed. _store_prefill_context then indexes the newly computed tail from first_pos instead of 0, because the matched blocks already hold the prefix's drafter K/V. Gated on the pool being the draft manager's and that manager being paired; a short block table raises rather than writing into a neighbour. _joint_reuse_supported only demands _supports_reuse_match_backoff for a non-zero span, which is what lets the K3 KDA hybrid target pair at span 0. Signed-off-by: Zhenhuan Chen (cherry picked from commit b90a1b2b9b6bb13d2334ad0ffb9058c94b291ad6) --- tensorrt_llm/_torch/pyexecutor/_util.py | 14 ++++++++++--- tensorrt_llm/_torch/speculative/dflash.py | 16 +++++++++++++++ tensorrt_llm/_torch/speculative/interface.py | 21 +++++++++++--------- 3 files changed, 39 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 94dc2c8ebfe6..a337c6f1a8f2 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1653,10 +1653,18 @@ def _joint_reuse_supported(self) -> bool: """ if not self._is_kv_cache_manager_v2: return False - if not getattr(self._kv_cache_manager_cls, - "_supports_reuse_match_backoff", False): + lookahead = draft_prompt_lookahead(self._speculative_config) + if lookahead is None: return False - return draft_prompt_lookahead(self._speculative_config) is not None + if lookahead > 0 and not getattr(self._kv_cache_manager_cls, + "_supports_reuse_match_backoff", + False): + # The opt-out is about backing the match off by `lookahead` tokens, + # which a specialized commit/history protocol (recurrent snapshots, + # DSA) cannot express. A zero span asks for no backoff at all, so + # every backoff-sized path stays a no-op and the pairing is safe. + return False + return True def _get_effective_draft_config(self) -> ModelConfig: """ diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index f03374914337..e65c6e3e4604 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -320,6 +320,11 @@ def __init__( self._ctx_page_size = 32 self._ctx_pages_per_slot = 0 self._ctx_paged_append = None + # True only when the drafter reads the draft KV cache manager's pool + # AND that manager joins the target's reuse protocol, which is what + # makes a reused prefix's drafter K/V still addressable. Set in + # _lazy_init_ctx_buffers; read by _store_prefill_context. + self._ctx_reuse_addressable = False self._dflash_attention_backend = spec_config.attention_backend # Slot management (Python, updated in prepare() and eager mode) @@ -739,6 +744,12 @@ def _lazy_init_ctx_buffers( # block table one iteration at a time, so the footprint follows # the sequences served rather than max_batch x max_seq_len. self._ctx_kv_buf = pool + # Only a manager that publishes and matches its own blocks can + # hand back a reused prefix's drafter K/V; the private arena and + # an unpaired draft pool both start every request at 0. + self._ctx_reuse_addressable = bool( + getattr(draft_kv_cache_manager, "enable_joint_kv_cache_reuse", False) + ) self._ctx_kv_last_page_len = torch.full( (num_slots,), page_size, dtype=torch.int32, device="cuda" ) @@ -996,6 +1007,11 @@ def _store_prefill_context( slot = self._req_to_slot[req_id] cur = ctx_len_updates.get(slot, self._ctx_len_host[slot]) cap = self._max_ctx if ctx_alloc is None else min(self._max_ctx, ctx_alloc[i]) + if cur == 0 and first_pos > 0 and self._ctx_reuse_addressable: + # Reuse let the target skip [0, first_pos), so only the tail + # reaches the capture buffer. The prefix K/V are still in the + # paired pool: precompute_context_kv is per-token in (hidden, p). + cur = first_pos if cur + slen > cap: # Request-level, like the no-free-slots path above: truncating # would silently draft from a stale prefix, but killing the diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 7007e000045d..d7d84e2e769c 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -148,19 +148,22 @@ def draft_prompt_lookahead(spec_config) -> Optional[int]: # Two-model drafting runs the draft as its own engine over its own # request view, so nothing here is shifted against the target prompt. return None - if spec_mode.is_mtp_vanilla(): - # Known limitation: an internal chunk lacks target_hidden[i + k], so reuse - # can attach draft KV an unchunked prefill would not produce. Acceptance only. - return spec_config.max_draft_len - if spec_mode.is_eagle_one_model(): - return 1 if spec_mode.is_pard(): # PARDWorker feeds the drafter input_ids[:num_ctx_tokens] verbatim, so # its draft state at i is a function of tokens [0, i] like the target's. return 0 - # Unestablished elsewhere: DraftTargetOneModel likely shifts by 1 but is - # unvalidated; DFlash/DSpark build context draft K/V from projected target - # hidden states into worker-owned buffers, which block reuse cannot restore. + if spec_mode.is_dflash() or spec_mode.is_dspark(): + # precompute_context_kv(projected_hidden, positions) writes the drafter's + # OWN post-norm post-RoPE K/V, one entry per context token, and the + # target hidden at i already depends on exactly tokens [0, i]. Same + # dependency face as the target's own KV, so raw-prompt keys describe it + # and no chunk-tail lookahead token is needed. The paged pool then makes + # a matched prefix addressable again (dflash.py _store_prefill_context). + return 0 + # A shift-by-1 (Eagle) or D-chained (vanilla MTP) span additionally needs the + # context-chunk lookahead token, whose plumbing this branch does not carry; + # keying those modes without it would attach draft KV a chunked prefill never + # produced. DraftTargetOneModel's span is unvalidated upstream. return None From b6148798581a58d1484856fa5c3a12e14253da2d Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Wed, 16 Sep 2026 21:45:48 -0700 Subject: [PATCH 21/25] [None][chore] Report whether the DFlash ctx cache can reach a reused prefix Every other _managed_ctx_pool fallback costs memory only. This one is silent: the scheduler keeps matching prefixes the drafter cannot read, and the run looks unpaired in acceptance length alone. Signed-off-by: Zhenhuan Chen (cherry picked from commit 00ad9def03385d6f32e0f80a592a27727a62ce7b) --- tensorrt_llm/_torch/speculative/dflash.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index e65c6e3e4604..6c7f13070510 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -764,10 +764,25 @@ def _lazy_init_ctx_buffers( ) self._ctx_buf_inited = True + if not self._ctx_reuse_addressable and getattr( + draft_kv_cache_manager, "enable_joint_kv_cache_reuse", False + ): + # Every fallback above (shape mismatch, stride mismatch, pool layer + # count) logs its own reason and then costs only memory -- except + # this one, where the scheduler goes on matching prefixes the + # drafter can no longer read. Silent, and it looks exactly like the + # unpaired build in every metric but acceptance length. + logger.warning( + "DFlash: the draft KV cache manager joins block reuse but the drafter " + "fell back to a private ctx arena, so a reused prefix's drafter K/V " + "stays unreachable and acceptance length drops as if unpaired." + ) + logger.info( f"DFlash: allocated ctx buffers: max_batch={max_batch}, " f"max_ctx={self._max_ctx}, dtype={dtype}, " - f"dflash_attention_backend={self._dflash_attention_backend}" + f"dflash_attention_backend={self._dflash_attention_backend}, " + f"reuse_addressable={self._ctx_reuse_addressable}" ) def _ctx_paged_index_args(self): From 1ead943463314037b5ef346b736f80066b7c88fd Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Thu, 17 Sep 2026 02:30:29 -0700 Subject: [PATCH 22/25] [None][fix] Give the EPLB config fixture a max_seq_len external_drafter_config_kwargs() forwards model_config.max_seq_len so the drafter sizes its position table from what the runtime serves. The fixture builds model_config as a SimpleNamespace and had no such field, so all seven tests raised AttributeError. Read unguarded on purpose: a real ModelConfig always carries it, and a getattr fallback would silently restore the max_position_embeddings sizing this exists to remove. Signed-off-by: Zhenhuan Chen --- .../_torch/speculative/hw_agnostic/test_dspark_eplb_config.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_eplb_config.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_eplb_config.py index 38ca76204105..193e88f84e0a 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_eplb_config.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_eplb_config.py @@ -64,6 +64,10 @@ def _model_config(lb_config=None, num_hidden_layers=NUM_HIDDEN_LAYERS): mapping=object(), max_num_tokens=8192, moe_max_num_tokens=8192, + # Forwarded to the drafter so its position table is sized by what the + # runtime serves rather than max_position_embeddings. A real + # ModelConfig always carries it, so the helper reads it unguarded. + max_seq_len=8192, pretrained_config=SimpleNamespace(num_hidden_layers=num_hidden_layers), ) From dbfadd00f082872a40ca06e92bb791ddfb058548 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Thu, 17 Sep 2026 07:08:12 -0700 Subject: [PATCH 23/25] [None][fix] Keep Eagle/MTP in draft_prompt_lookahead The reuse-protocol cherry-pick came off a branch predating #18093, so applying it cleanly replaced this function wholesale and dropped the is_mtp_vanilla and is_eagle_one_model branches with it. Eagle/MTP one-model fell through to None: reuse_match_backoff went 1 -> 0, the estimator stopped charging the reuse window (36864 -> 24576 B), and Eagle3 reuse acceptance regressed. Restored; the DFlash/DSpark branch this PR adds is unaffected. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/speculative/interface.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index d7d84e2e769c..da68041e5c6c 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -148,6 +148,12 @@ def draft_prompt_lookahead(spec_config) -> Optional[int]: # Two-model drafting runs the draft as its own engine over its own # request view, so nothing here is shifted against the target prompt. return None + if spec_mode.is_mtp_vanilla(): + # Known limitation: an internal chunk lacks target_hidden[i + k], so reuse + # can attach draft KV an unchunked prefill would not produce. Acceptance only. + return spec_config.max_draft_len + if spec_mode.is_eagle_one_model(): + return 1 if spec_mode.is_pard(): # PARDWorker feeds the drafter input_ids[:num_ctx_tokens] verbatim, so # its draft state at i is a function of tokens [0, i] like the target's. @@ -160,10 +166,7 @@ def draft_prompt_lookahead(spec_config) -> Optional[int]: # and no chunk-tail lookahead token is needed. The paged pool then makes # a matched prefix addressable again (dflash.py _store_prefill_context). return 0 - # A shift-by-1 (Eagle) or D-chained (vanilla MTP) span additionally needs the - # context-chunk lookahead token, whose plumbing this branch does not carry; - # keying those modes without it would attach draft KV a chunked prefill never - # produced. DraftTargetOneModel's span is unvalidated upstream. + # DraftTargetOneModel likely shifts by 1 but is unvalidated upstream. return None From f6f4e83e08aeaa36d2cffaa4252f98f7090e3b12 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Thu, 17 Sep 2026 21:55:33 -0700 Subject: [PATCH 24/25] [None][chore] Name both dummy_slot_row publishers in the padding-row guard prepare_penalty_buffers is the second place that publishes dummy_slot_row; it does not exist on rubin-advance, where this guard was written. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/speculative/interface.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index da68041e5c6c..3a0b4b9e45bc 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -2090,7 +2090,8 @@ def _zero_padding_rows(logits: torch.Tensor, spec_metadata, routed to ``dummy_slot_row`` (``py_seq_slot is None``), so this needs no extra host-side state. Shapes are static: CUDA-graph safe. No-op while ``dummy_slot_row`` is still 0, which is a live request's row rather than - a padding marker. + a padding marker; ``prepare_rejection_sampling_buffers`` and + ``prepare_penalty_buffers`` are the two places that publish it. Args: logits: ``[(batch_size - num_contexts) * rows_per_request, vocab]``; From 54fd15a5b12e91eeeaa43616f46f4f8a38470961 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Fri, 18 Sep 2026 01:17:18 -0700 Subject: [PATCH 25/25] [None][fix] Bound the DSv4 DSpark drafter's RoPE positions Shrinking the table to the runtime ceiling made its end reachable, and warmup advances a slot no completion frees. It arrives through a graph replay, which runs no Python, so only an in-graph bound holds; the ceiling was also one short. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/speculative/dspark.py | 87 ++++++++++++++----- .../hw_agnostic/test_dspark_worker.py | 60 ++++++++++++- 2 files changed, 120 insertions(+), 27 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 27761d187d55..d62b2750f61b 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -40,13 +40,17 @@ def _dspark_position_ceiling(max_ctx: int, block_size: int, max_draft_len: int) -> int: """Return the number of RoPE entries needed by the DSv4 block drafter. - A target verification can accept the target token plus every draft token, - so ``start_pos`` advances by ``max_draft_len + 1``. The block drafter then - indexes ``block_size`` positions beginning at that start. With ``old`` - bounded by ``max_ctx``, the largest index is - ``max_ctx + max_draft_len + 1 + block_size``; the length is one greater. + ``start_pos`` is a FRAME index, one above the absolute token position: the + prompt token at position p occupies frame p+1 (``_seed_context_windows``), so a + request served to ``max_ctx`` bootstraps at ``max_ctx + 1``. A verification then + accepts up to ``max_draft_len + 1`` tokens and the block drafter indexes + ``block_size`` further positions, making the largest index + ``max_ctx + 1 + max_draft_len + 1 + block_size``; the length is one greater. + + The frame +1 is load-bearing: without it the drafter reached start_pos 4112 at + max_ctx 4105 and max_draft_len 5, one past the table [measured job 3097207]. """ - return int(max_ctx) + int(max_draft_len) + int(block_size) + 2 + return int(max_ctx) + int(max_draft_len) + int(block_size) + 3 @dataclass @@ -262,6 +266,9 @@ def __init__( self._valid_len: Optional[torch.Tensor] = None # [max_batch] written window entries self._position_initialized: Optional[torch.Tensor] = None # [max_batch] bool self._win = 0 + # Set in _lazy_init from the RoPE table the drafter will build; None + # leaves positions unbounded (direct construction in tests). + self._position_cap: Optional[int] = None # Slot management. ``_req_to_slot`` (python dict) + ``_free_slots`` are the # source of truth, updated in prepare()/forward(); ``_batch_to_slot`` is the @@ -292,6 +299,42 @@ def __init__( def max_draft_len(self) -> int: return self.spec_config.max_draft_len + def _publish_position_ceiling(self, draft_model, attn_metadata, block_size: int) -> None: + """Size the drafter's RoPE table from the runtime bound, and bound positions by it. + + Runs on EVERY forward, not once behind ``_win_inited``. The first forward can + land on a KV-estimation probe manager whose ``max_seq_len`` is below the real + one, and a table pinned to that probe would index out of range for the rest + of the process; ``DFlashWorker._lazy_init_ctx_buffers`` re-publishes for that + reason. Grows only, and the table cache is keyed on the cap. + + Two bounds exist and neither dominates: ``attn_metadata`` carries the KV + manager's ``max_seq_len`` while ``_freqs_cap`` carries the user's, and + ``_create_cuda_graph_warmup_request`` sizes its dummy request at whichever is + larger. Covering both is what keeps warmup off the end of the table. + """ + inner_model = getattr(draft_model, "dspark_model", None) or draft_model + config_cap = int(getattr(inner_model, "_freqs_cap", 0) or 0) + # _freqs_cap is max_seq_len + block_size + 2; undo that so the config + # bound and the KV manager's go through one formula. + config_ctx = max(0, config_cap - block_size - 2) + max_ctx = getattr(attn_metadata, "max_seq_len", None) + if max_ctx is not None: + ceiling = max( + _dspark_position_ceiling( + max(int(max_ctx), config_ctx), block_size, self.max_draft_len + ), + int(getattr(inner_model, "_runtime_position_ceiling", 0) or 0), + ) + draft_model._runtime_position_ceiling = ceiling + if inner_model is not draft_model: + inner_model._runtime_position_ceiling = ceiling + # Largest absolute position the block drafter may hold: forward_batched + # gathers ``freqs[start_pos + block_size]`` and the interim back-fill + # ``freqs[old + block_size]``, so keep one block of headroom. + table_len = int(getattr(inner_model, "_runtime_position_ceiling", 0) or config_cap) + self._position_cap = (table_len - 1 - block_size) if table_len else None + def _lazy_init(self, draft_model, spec_metadata, attn_metadata=None) -> None: block_size = int(draft_model.block_size) if block_size != self.max_draft_len: @@ -300,21 +343,9 @@ def _lazy_init(self, draft_model, spec_metadata, attn_metadata=None) -> None: f"got block_size={block_size} and max_draft_len={self.max_draft_len}" ) + self._publish_position_ceiling(draft_model, attn_metadata, block_size) + if not self._win_inited: - # Published before the RoPE table is built (lazily, on first forward). The - # engine's max_seq_len is what positions reach and exceeds model_config's - # whenever spec decoding is on, so the config cap undersizes the table. - max_ctx = getattr(attn_metadata, "max_seq_len", None) - if max_ctx is not None: - ceiling = _dspark_position_ceiling(max_ctx, block_size, self.max_draft_len) - draft_model._runtime_position_ceiling = ceiling - # The worker owns the DSv4 wrapper, while the RoPE-table cache - # lives on its inner ``dspark_model``. Publish to both so the - # value that sizes the table is the runtime bound rather than - # the construction fallback. - inner_model = getattr(draft_model, "dspark_model", None) - if inner_model is not None: - inner_model._runtime_position_ceiling = ceiling max_batch = spec_metadata.max_num_requests num_stages = draft_model.num_stages self._win = int(draft_model._attn_params["window_size"]) @@ -472,6 +503,14 @@ def _advance_generation_state( self._position_initialized[scratch].zero_() old = torch.where(self._position_initialized[slots], self._ctx_len[slots], input_positions) start_pos = old + num_accepted_tokens + # Bound both positions by the RoPE table they index. Warmup reaches this + # through a captured graph replay, which runs no Python, so only an in-graph + # tensor op can hold the line. _position_cap sits exactly at what a real + # request reaches, and only a warmup row -- whose draft is discarded -- + # climbs past it, because no completion ever frees its slot. + if self._position_cap is not None: + old = torch.clamp(old, max=self._position_cap) + start_pos = torch.clamp(start_pos, max=self._position_cap) self._ctx_len[slots] = start_pos self._valid_len[slots] = torch.clamp( self._valid_len[slots] + num_accepted_tokens, max=self._win @@ -517,6 +556,10 @@ def _draft_gen_block_batched( # the gen tokens after the context tokens. gen_start = attn_metadata.num_ctx_tokens slots = self._batch_to_slot[num_contexts:batch_size] # [G] + # Bootstrap iterations can process one target token per request, while + # normal speculative verification processes K+1. Use the actual accepted + # row width to index both captured hidden states and position IDs. + target_width = accepted_tokens.shape[1] nacc = num_accepted_tokens[num_contexts:batch_size].long() # [G] gidx = nacc - 1 # [G] index of the bonus within each verified prefix @@ -525,10 +568,6 @@ def _draft_gen_block_batched( accepted_tokens[num_contexts:batch_size].gather(1, gidx.unsqueeze(1)).squeeze(1).long() ) # [G] - # Bootstrap iterations can process one target token per request, while - # normal speculative verification processes K+1. Use the actual accepted - # row width to index both captured hidden states and position IDs. - target_width = accepted_tokens.shape[1] arange_g = torch.arange(num_gens, device=device) base = gen_start + arange_g * target_width # [G] main_hidden = captured[base + gidx] # [G, ncap*hidden] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py index e32a24dc20f9..0c898b42288f 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -912,13 +912,67 @@ def test_lazy_init_publishes_the_runtime_position_ceiling(): worker._lazy_init(dm, _make_metadata(max_num_requests=2), attn_metadata) - # A full target verification accepts K+1 tokens, then DSpark indexes K - # block positions from that new start. The table length is max index + 1. - expected = 4096 + 5 + 1 + 5 + 1 + # start_pos is a frame index (position + 1); a full verification then accepts + # K+1 tokens and DSpark indexes K block positions from that start. Length is + # max index + 1. + expected = 4096 + 1 + 5 + 1 + 5 + 1 assert dm._runtime_position_ceiling == expected assert dm.dspark_model._runtime_position_ceiling == expected +def test_position_ceiling_covers_the_config_cap_and_grows_after_kv_estimation(): + """Publishing once, from one bound, undersizes the table for the rest of the run. + + KV-cache estimation drives drafter forwards against a probe manager whose + max_seq_len is below the real one, and _create_cuda_graph_warmup_request puts + its dummy request at max_seq_len - 1 -- within one block of the table's end. + """ + worker = _make_worker() + dm = _fake_draft_model() + dm.dspark_model = types.SimpleNamespace(_freqs_cap=9000) + meta = _make_metadata(max_num_requests=2) + + # The probe manager's bound is below the config cap: cover both, not one. + # _freqs_cap 9000 encodes max_seq_len 9000 - block_size - 2. + worker._lazy_init(dm, meta, types.SimpleNamespace(max_seq_len=512)) + assert dm.dspark_model._runtime_position_ceiling == (9000 - 5 - 2) + 1 + 5 + 1 + 5 + 1 + + # The real manager arrives after _win_inited is already set. + worker._lazy_init(dm, meta, types.SimpleNamespace(max_seq_len=16384)) + assert dm.dspark_model._runtime_position_ceiling == 16384 + 1 + 5 + 1 + 5 + 1 + + +def test_position_bound_holds_across_cuda_graph_replays(): + """The bound has to be an in-graph tensor op, not a host-side guard. + + Warmup reaches this through ``cuda_graph_runner.replay``, which runs no Python, + so a captured graph keeps advancing a slot that no completion ever frees. Left + unbounded it ran 53 entries past the RoPE table [measured job 3097207]; the cap + itself sits at what a full-length real request reaches, so serving is untouched. + """ + worker = _make_worker() + dm = _fake_draft_model(window_size=8) + worker._lazy_init(dm, _make_metadata(max_num_requests=2), types.SimpleNamespace(max_seq_len=64)) + cap = worker._position_cap + assert cap == 64 + 1 + 5 + 1 # max_ctx + frame + K + bonus + assert cap == (64 + 5 + 5 + 3) - 1 - 5 # == table length - 1 - block_size + + slots = torch.tensor([0], device="cuda", dtype=torch.long) + num_accepted = torch.tensor([6], device="cuda", dtype=torch.long) + input_positions = torch.tensor([cap - 3], device="cuda", dtype=torch.long) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + worker._advance_generation_state(slots, num_accepted, input_positions) + + worker._ctx_len[slots] = 0 + worker._position_initialized[slots] = False + for _ in range(5): + graph.replay() + assert worker._ctx_len[0].item() <= cap + assert worker._ctx_len[0].item() == cap + + def test_freqs_table_prefers_the_runtime_ceiling_over_the_config_cap(mocker): """The cache key has to carry the cap, not just the device.