From e27249f524b1211029b4f30346e7e0b7b57d2d1c Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Tue, 18 Aug 2026 22:18:27 -0700 Subject: [PATCH 01/21] [None][refactor] Dispatch draft models through a builder registry get_draft_model held every draft implementation in one if/elif chain, so the base file imported its own consumers -- DSpark had to import lazily to break the modeling_dspark -> modeling_deepseekv4 -> modeling_speculative cycle. Builders now register from the module defining their draft model, resolved through SPEC_MODE_TO_MODULE like the model zoo. Pure refactor. EAGLE3 keeps precedence over the _use_shared_kv_cache override, which the old branch order encoded implicitly. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/models/_arch_index.py | 61 ++++- tensorrt_llm/_torch/models/modeling_dspark.py | 28 +++ .../_torch/models/modeling_speculative.py | 176 +++++++++----- tensorrt_llm/_torch/models/modeling_utils.py | 122 +++++++++- .../hw_agnostic/test_draft_model_registry.py | 226 ++++++++++++++++++ tests/unittest/others/test_lazy_model_zoo.py | 68 ++++++ 6 files changed, 610 insertions(+), 71 deletions(-) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_draft_model_registry.py diff --git a/tensorrt_llm/_torch/models/_arch_index.py b/tensorrt_llm/_torch/models/_arch_index.py index 8915af4c6475..7c04d8d0a450 100644 --- a/tensorrt_llm/_torch/models/_arch_index.py +++ b/tensorrt_llm/_torch/models/_arch_index.py @@ -6,11 +6,13 @@ import side effect. The zoo is imported lazily, so these tables record, without importing anything, which ``modeling_*`` module provides which architecture (``MODEL_ARCH_TO_MODULE``), which public class (``MODEL_CLASS_TO_MODULE``), -and which multimodal ``model_type`` (``MULTIMODAL_MODEL_TYPE_TO_MODULE``). +which multimodal ``model_type`` (``MULTIMODAL_MODEL_TYPE_TO_MODULE``), and +which speculative-decoding mode (``SPEC_MODE_TO_MODULE``). Regenerate after adding/moving a model: add the new entry by hand next to its neighbors, mirroring the ``@register_auto_model("")`` / -``@register_input_processor(..., model_type="")`` decorators and the +``@register_input_processor(..., model_type="")`` / +``@register_draft_model(SpeculativeDecodingMode.)`` decorators and the public class name. ``test_lazy_model_zoo.py`` fails on any drift between these tables and the decorators. """ @@ -223,3 +225,58 @@ def is_builtin_zoo_module(module_name: str) -> bool: "step3p7": "modeling_step3p7vl", "whisper": "modeling_whisper", } + +# ``SpeculativeDecodingMode`` member name -> module providing that mode's draft +# model builder (registered via ``@register_draft_model``). Keyed by the enum +# member *name* rather than the enum itself so this module keeps importing +# nothing. Modes absent from this table have no one-engine draft model to build +# (two-model / drafter-loop modes such as NGRAM, SA and USER_PROVIDED). +# +# Adding a speculative decoding mode +# ---------------------------------- +# 1. Write the builder in *your own* ``modeling_*.py``, next to the draft model +# it constructs -- never in ``modeling_speculative.py``. Keeping builders out +# of the factory file is the entire point of this table: ``get_draft_model`` +# imports no concrete draft implementation, which is what let DSpark drop the +# lazy import it needed while ``modeling_dspark`` imports back into +# ``modeling_speculative`` through ``modeling_deepseekv4``. +# 2. Decorate it with ``@register_draft_model(SpeculativeDecodingMode.)``. +# Stack the decorator to serve several modes with one builder. +# 3. Add the ``"": "modeling_"`` row below. +# 4. ``test_lazy_model_zoo.py`` and +# ``tests/unittest/_torch/speculative/hw_agnostic/test_draft_model_registry.py`` +# fail in both directions on any drift between decorators and this table. +# +# The builder signature is fixed at +# ``(model_config, draft_config, lm_head, model) -> nn.Module`` -- byte-for-byte +# the arguments of ``get_draft_model``, so the factory is pure forwarding with +# no per-mode glue. Everything a builder needs is reachable from those four +# (``model_config.pretrained_config.num_hidden_layers``, ``model.aux_stream_dict``, +# ``model_config.spec_config.*``). Do not widen it: an extra parameter has to be +# populated by the factory, which puts mode-specific knowledge straight back +# into the shared file this table exists to keep generic. +# +# Three rules the registry inherits from ``register_auto_model`` (see +# ``modeling_utils.py``), each written down because it was learned the hard way: +# - Look up only through ``get_registered_draft_model_builder``. It triggers +# the on-demand import before reading the mapping; a raw ``.get()`` silently +# misses every provider that has not been imported yet, and the zoo is +# imported lazily. +# - Built-in builders only fill empty slots, never overwrite. Lazy loading +# means a built-in's decorator can run *after* an external registration +# (``--custom_module_dirs``), so last-wins would let the built-in clobber a +# user's drafter. +# - Map a builder back to its modes via its ``_registered_spec_modes`` +# attribute, not by scanning the mapping for it. A built-in that lost its +# slot to an external registration is absent from the mapping but still has +# the attribute, so an identity scan reports it as unregistered. +SPEC_MODE_TO_MODULE = { + "DFLASH": "modeling_speculative", + "DRAFT_TARGET_ONE_MODEL": "modeling_speculative", + "DSPARK": "modeling_dspark", + "EAGLE3_ONE_MODEL": "modeling_speculative", + "MTP": "modeling_speculative", + "MTP_EAGLE": "modeling_speculative", + "MTP_EAGLE_ONE_MODEL": "modeling_speculative", + "PARD": "modeling_speculative", +} diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 2d556af5e71e..3281336e1de6 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -51,6 +51,7 @@ from ..modules.linear import Linear from ..modules.mhc.hyper_connection import HCHead from ..modules.rms_norm import RMSNorm +from ..speculative.interface import SpeculativeDecodingMode from ..utils import AuxStreamType from .dspark.attention import ( _rmsnorm, @@ -71,6 +72,7 @@ _rename_deepseek_v4_attn_subkey, _rename_deepseek_v4_ffn_subkey, ) +from .modeling_utils import register_draft_model # Matches the draft namespace ``mtp..`` in the V4-Pro-DSpark # checkpoint. Each draft stage is a full DeepSeek-V4 block stored under this @@ -1241,6 +1243,32 @@ def load_weights_from_target_model(self, target_model): self.dspark_model.lm_head = target_model.lm_head +@register_draft_model(SpeculativeDecodingMode.DSPARK) +def _build_dspark_draft(model_config, draft_config, lm_head, model): + """Build the DSpark drafter, reusing the target's aux streams. + + The draft stage count (``n_mtp_layers``) is not in the HF config, so it is + derived from the checkpoint's ``mtp.*`` namespace. + + Args: + model_config: the target engine's ``ModelConfig``. + draft_config: the drafter's own ``ModelConfig``. + lm_head: unused; DSpark shares the target's head at weight-load time. + model: the target model, whose aux streams the draft stages reuse. + + Returns: + The ``DSparkForCausalLM`` draft module. + """ + num_stages = count_dspark_stages(model_config.spec_config.speculative_model) + validate_dspark_eplb_layer_base(model_config, draft_config) + return DSparkForCausalLM( + draft_config, + getattr(model, "aux_stream_dict", None), + num_stages=num_stages, + block_size=model_config.spec_config.block_size, + ) + + __all__ = [ "DSparkBlock", "DSparkDraftModel", diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index ab8522a54471..4f032756f906 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -39,11 +39,14 @@ should_use_separate_draft_kv_cache) from ..speculative.dflash_attention import (get_dflash_flash_attention, get_dflash_trtllm_gen_ops) +from ..speculative.interface import SpeculativeDecodingMode from ..utils import AuxStreamType from .checkpoints.base_weight_mapper import BaseWeightMapper from .modeling_auto import AutoModelForCausalLM from .modeling_utils import (DecoderModel, DecoderModelForCausalLM, TModel, - get_model_architecture, register_auto_model) + get_model_architecture, + get_registered_draft_model_builder, + register_auto_model, register_draft_model) _SPECULATIVE_POSITION_HEADROOM = "_speculative_position_headroom" @@ -2433,81 +2436,122 @@ def external_drafter_config_kwargs(model_config, spec_config) -> dict: return kwargs -def get_draft_model(model_config, draft_config, lm_head, model): - """Construct the draft model for the configured speculative-decoding mode - (EAGLE3 / MTP / PARD / DFlash). The DFlash branch selects the Laguna drafter - by detecting its architecture in the draft checkpoint's own config.""" - assert getattr(model_config, 'spec_config', None) is not None +@register_draft_model(SpeculativeDecodingMode.EAGLE3_ONE_MODEL) +def _build_eagle3_one_model_draft(model_config, draft_config, lm_head, model): + """Build the EAGLE3 one-model drafter for the configured draft arch.""" spec_dec_mode = model_config.spec_config.spec_dec_mode - if spec_dec_mode.is_eagle3_one_model(): - if model_config.spec_config.eagle3_model_arch == "llama3": - # Eagle3ForCausalLM handles both Llama3 and DeepSeekV3 architectures - return Eagle3ForCausalLM( - draft_config, model_config.pretrained_config.num_hidden_layers) - elif model_config.spec_config.eagle3_model_arch == "mistral_large3": - return MistralLarge3EagleForCausalLM( - draft_config, model_config.pretrained_config.num_hidden_layers, - model.aux_stream_dict) - else: - raise ValueError( - f"Unsupported eagle3 model architecture: {spec_dec_mode.eagle3_model_arch}" - ) + if model_config.spec_config.eagle3_model_arch == "llama3": + # Eagle3ForCausalLM handles both Llama3 and DeepSeekV3 architectures + return Eagle3ForCausalLM( + draft_config, model_config.pretrained_config.num_hidden_layers) + elif model_config.spec_config.eagle3_model_arch == "mistral_large3": + return MistralLarge3EagleForCausalLM( + draft_config, model_config.pretrained_config.num_hidden_layers, + model.aux_stream_dict) + else: + raise ValueError( + f"Unsupported eagle3 model architecture: {spec_dec_mode.eagle3_model_arch}" + ) + + +@register_draft_model(SpeculativeDecodingMode.MTP) +@register_draft_model(SpeculativeDecodingMode.MTP_EAGLE_ONE_MODEL) +def _build_mtp_one_model_draft(model_config, draft_config, lm_head, model): + """Build the one-model MTP drafter (vanilla MTP and MTP-Eagle share it).""" + return MTPForCausalLM(model_config, + model_config.pretrained_config.num_hidden_layers, + lm_head, model) + + +@register_draft_model(SpeculativeDecodingMode.MTP_EAGLE) +def _build_mtp_eagle_draft(model_config, draft_config, lm_head, model): + """Build the two-model MTP-Eagle drafter.""" + return MTPDraftModelForCausalLM(model_config) + + +@register_draft_model(SpeculativeDecodingMode.PARD) +def _build_pard_draft(model_config, draft_config, lm_head, model): + """Build the PARD drafter.""" + return PARDForCausalLM(draft_config) + + +@register_draft_model(SpeculativeDecodingMode.DFLASH) +def _build_dflash_draft(model_config, draft_config, lm_head, model): + """Build the DFlash drafter. + + Selects the Laguna variant by detecting its architecture in the draft + checkpoint's own config. + """ + draft_arches = getattr(draft_config.pretrained_config, "architectures", + None) or [] + dflash_attention_backend = model_config.spec_config.attention_backend + if any("Laguna" in arch for arch in draft_arches): + return DFlashLagunaForCausalLM( + draft_config, + dflash_attention_backend=dflash_attention_backend, + ) + return DFlashForCausalLM( + draft_config, + dflash_attention_backend=dflash_attention_backend, + ) + + +@register_draft_model(SpeculativeDecodingMode.DRAFT_TARGET_ONE_MODEL) +def _build_draft_target_one_model_draft(model_config, draft_config, lm_head, + model): + """Build the one-model draft-target drafter from its own checkpoint.""" + # Keep the draft LM head vocab-sharded so greedy draft sampling uses the + # lighter TP gather (see SpecWorkerBase.greedy_sample_draft_with_tp_gather). + was_frozen = draft_config._frozen + draft_config._frozen = False + draft_config.lm_head_gather_output = False + draft_config._frozen = was_frozen + return AutoModelForCausalLM.from_config(draft_config) + + +def get_draft_model(model_config, draft_config, lm_head, model): + """Construct the draft model for the configured speculative-decoding mode. + + Dispatch is registry-based: each mode's builder lives next to the draft + model it constructs and registers itself via ``@register_draft_model``, so + this function never imports a concrete draft implementation (which is what + used to force a lazy import for DSpark, whose provider imports back into + this module through modeling_deepseekv4). - elif model_config.spec_config.uses_external_draft_model: + Args: + model_config: the target engine's ``ModelConfig``, carrying spec_config. + draft_config: the drafter's own ``ModelConfig``, or None when the mode + builds its draft from the target config alone. + lm_head: the target's LM head, shared by the one-model MTP drafter. + model: the target model, for drafters reusing its aux streams. + + Returns: + The draft ``nn.Module`` for this mode. + """ + assert getattr(model_config, 'spec_config', None) is not None + spec_config = model_config.spec_config + spec_dec_mode = spec_config.spec_dec_mode + # An external draft model is loaded straight from its own checkpoint, so it + # has no mode-specific builder to register: this stays an explicit pre-check + # ahead of the registry lookup rather than becoming a registry key. + # + # No mode guard is needed. `uses_external_draft_model` already implies + # `is_mtp_one_model()` (llm_args), which is disjoint from every other mode, + # so this branch cannot divert a drafter that a builder would have claimed. + # Pinned by test_draft_model_registry.py:: + # test_external_draft_model_bypasses_the_registry. + if spec_config.uses_external_draft_model: if draft_config is None: raise ValueError( "MTP speculative decoding with an external draft model requires " "its model config.") return AutoModelForCausalLM.from_config(draft_config) - elif spec_dec_mode.is_mtp_one_model(): - return MTPForCausalLM(model_config, - model_config.pretrained_config.num_hidden_layers, - lm_head, model) - elif spec_dec_mode.is_mtp_eagle(): - return MTPDraftModelForCausalLM(model_config) - elif spec_dec_mode.is_pard(): - return PARDForCausalLM(draft_config) - elif spec_dec_mode.is_dflash(): - draft_arches = getattr(draft_config.pretrained_config, "architectures", - None) or [] - dflash_attention_backend = model_config.spec_config.attention_backend - if any("Laguna" in arch for arch in draft_arches): - return DFlashLagunaForCausalLM( - draft_config, - dflash_attention_backend=dflash_attention_backend, - ) - return DFlashForCausalLM( - draft_config, - dflash_attention_backend=dflash_attention_backend, - ) - elif spec_dec_mode.is_dspark(): - # Lazy import to avoid a cycle (modeling_dspark -> modeling_deepseekv4 -> - # modeling_speculative). The DSpark draft reuses the target's aux streams. - # The draft stage count (n_mtp_layers) is not in the HF config, so derive - # it from the checkpoint's mtp.* namespace. - from .modeling_dspark import (DSparkForCausalLM, count_dspark_stages, - validate_dspark_eplb_layer_base) - num_stages = count_dspark_stages( - model_config.spec_config.speculative_model) - validate_dspark_eplb_layer_base(model_config, draft_config) - return DSparkForCausalLM( - draft_config, - getattr(model, "aux_stream_dict", None), - num_stages=num_stages, - block_size=model_config.spec_config.block_size, - ) - elif spec_dec_mode.is_draft_target_one_model(): - # Keep the draft LM head vocab-sharded so greedy draft sampling uses the - # lighter TP gather (see SpecWorkerBase.greedy_sample_draft_with_tp_gather). - was_frozen = draft_config._frozen - draft_config._frozen = False - draft_config.lm_head_gather_output = False - draft_config._frozen = was_frozen - return AutoModelForCausalLM.from_config(draft_config) - else: + builder = get_registered_draft_model_builder(spec_dec_mode) + if builder is None: raise NotImplementedError( f"get_draft_model does not support speculative decoding mode {spec_dec_mode}." ) + return builder(model_config, draft_config, lm_head, model) class SpecDecOneEngineForCausalLM(DecoderModelForCausalLM[TModel, TConfig], diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 1a140ea15e03..6725b388eec2 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -8,8 +8,8 @@ import os import time from dataclasses import dataclass -from typing import (Any, Dict, Generic, Iterator, List, Literal, Optional, - Tuple, Type, TypeVar, Union) +from typing import (Any, Callable, Dict, Generic, Iterator, List, Literal, + Optional, Tuple, Type, TypeVar, Union) import torch from torch import nn @@ -33,7 +33,8 @@ from ..modules.logits_processor import LogitsProcessor from ..modules.rms_norm import RMSNorm from ..speculative import SpecMetadata -from ._arch_index import MODEL_ARCH_TO_MODULE, is_builtin_zoo_module +from ._arch_index import (MODEL_ARCH_TO_MODULE, SPEC_MODE_TO_MODULE, + is_builtin_zoo_module) @contextlib.contextmanager @@ -897,6 +898,7 @@ def infer_max_seq_len(self) -> int: MODEL_CLASS_MAPPING = {} +DRAFT_MODEL_BUILDER_MAPPING = {} MODEL_CLASS_VISION_ENCODER_MAPPING = {} MODEL_CLASS_MAPPER_MAPPING = {} MODEL_CLASS_CHECKPOINT_WEIGHT_LOADER_DEFAULT_MAPPING = {} @@ -999,6 +1001,120 @@ def get_registered_model_class(model_arch: str) -> Optional[Type[nn.Module]]: return MODEL_CLASS_MAPPING.get(model_arch) +def _is_builtin_draft_model_builder(builder) -> bool: + return is_builtin_zoo_module(getattr(builder, "__module__", "")) + + +# Speculative-decoding modes each decorated builder declared via +# ``register_draft_model``, kept per function (``__dict__``, never inherited). +# Recorded even when the builder loses its ``DRAFT_MODEL_BUILDER_MAPPING`` slot +# to an external registration, so a builder can be mapped back to its modes +# without scanning the mapping by identity (which would silently skip exactly +# those overridden built-ins). +_REGISTERED_SPEC_MODES_ATTR = "_registered_spec_modes" + + +def register_draft_model(mode): + """Register the draft-model builder for a speculative-decoding mode. + + The builder is a plain function + ``(model_config, draft_config, lm_head, model) -> nn.Module`` that owns + everything its mode needs to construct its draft model, so the generic + dispatcher never has to import a concrete draft implementation. Stack the + decorator to serve several modes with one builder (vanilla MTP and + MTP_EAGLE_ONE_MODEL share theirs, mirroring + ``SpeculativeDecodingMode.is_mtp_one_model()``). + + Same registration priority as ``register_auto_model``: built-in builders + only fill empty slots and never overwrite, because under lazy loading a + built-in module may run its decorators *after* an external registration + (e.g. a drafter supplied through ``--custom_module_dirs``) and must not + clobber it. External registrations always overwrite. + + The builder belongs in the ``modeling_*.py`` that defines the draft model, + never in the factory file, and needs an ``SPEC_MODE_TO_MODULE`` row in + ``_arch_index.py`` so it can be found without importing the zoo -- see the + "Adding a speculative decoding mode" notes there. Example (DSpark, whose + stage count is only knowable from the checkpoint):: + + @register_draft_model(SpeculativeDecodingMode.DSPARK) + def _build_dspark_draft(model_config, draft_config, lm_head, model): + num_stages = count_dspark_stages( + model_config.spec_config.speculative_model) + validate_dspark_eplb_layer_base(model_config, draft_config) + return DSparkForCausalLM( + draft_config, + getattr(model, "aux_stream_dict", None), + num_stages=num_stages, + block_size=model_config.spec_config.block_size, + ) + + Args: + mode: the ``SpeculativeDecodingMode`` member this builder serves. + + Returns: + The decorator binding a builder function to ``mode``. + """ + + def decorator(builder): + modes = builder.__dict__.get(_REGISTERED_SPEC_MODES_ATTR) + if modes is None: + modes = set() + setattr(builder, _REGISTERED_SPEC_MODES_ATTR, modes) + modes.add(mode) + + existing = DRAFT_MODEL_BUILDER_MAPPING.get(mode) + if (existing is not None and existing is not builder + and _is_builtin_draft_model_builder(builder)): + logger.info( + f"Keeping existing draft-model builder " + f"{existing.__module__}.{existing.__qualname__} for " + f"speculative decoding mode {mode.name}; built-in " + f"{builder.__module__}.{builder.__qualname__} not registered.") + return builder + DRAFT_MODEL_BUILDER_MAPPING[mode] = builder + return builder + + return decorator + + +def _ensure_draft_model_registered(mode) -> None: + """Import the module providing ``mode``'s builder, if not yet loaded. + + Mirrors ``_ensure_model_registered``: builders register as an import side + effect, and the model zoo is imported lazily, so this turns a mode into + "the decorator has run". Modes missing from the static index are left to + the caller's normal unsupported-mode handling. + """ + module_name = SPEC_MODE_TO_MODULE.get(mode.name) + if module_name is None: + return + full_name = f"tensorrt_llm._torch.models.{module_name}" + try: + importlib.import_module(full_name) + except ModuleNotFoundError as e: + # Only swallow "the providing module itself is missing" (stale index + # entry); a missing dependency *inside* the module is a real error and + # must not be masked as "unsupported speculative decoding mode". + if e.name != full_name: + raise + logger.warning(f"Lazy import of {module_name} for speculative " + f"decoding mode {mode.name} failed: {e!r}") + + +def get_registered_draft_model_builder(mode) -> Optional[Callable]: + """Resolve ``mode`` to its registered draft-model builder, or ``None``. + + The single entry point for builder lookups: the model zoo is imported + lazily, so this resolves the providing module on demand before reading the + registry. Do not read ``DRAFT_MODEL_BUILDER_MAPPING`` directly — a raw + ``.get()`` silently misses every not-yet-imported provider. + """ + if mode not in DRAFT_MODEL_BUILDER_MAPPING: + _ensure_draft_model_registered(mode) + return DRAFT_MODEL_BUILDER_MAPPING.get(mode) + + def get_registered_vision_encoder( model_arch: str) -> Optional[Tuple[Type[nn.Module], Optional[Type]]]: """Resolve ``model_arch`` to its ``(vision_encoder_cls, vlm_base_model)``. diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_model_registry.py b/tests/unittest/_torch/speculative/hw_agnostic/test_draft_model_registry.py new file mode 100644 index 000000000000..e05e82c079f9 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_draft_model_registry.py @@ -0,0 +1,226 @@ +# 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. +"""Dispatch contract of the draft-model builder registry. + +``get_draft_model`` used to be one if/elif chain whose *branch order* carried +unwritten rules. Registry dispatch has no inherent order, so what the chain +implied is pinned here explicitly — breaking it silently swaps the draft model +for ``AutoModelForCausalLM``, which no accuracy test would attribute back to +this function. + +The external-draft pre-check runs ahead of the registry without a mode guard, +which is only safe because ``uses_external_draft_model`` implies +``is_mtp_one_model()``. That mutual exclusion is an invariant of ``llm_args``, not +of this module, so it is asserted here rather than assumed. + +Everything here asserts *which builder is selected*, never the object it +builds: constructing a real drafter needs GPUs and checkpoints, and the +selection is the whole contract of this layer. +""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.models import modeling_speculative, modeling_utils +from tensorrt_llm._torch.models._arch_index import SPEC_MODE_TO_MODULE +from tensorrt_llm._torch.models.modeling_utils import ( + _REGISTERED_SPEC_MODES_ATTR, + DRAFT_MODEL_BUILDER_MAPPING, + get_registered_draft_model_builder, + register_draft_model, +) +from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode + +_EXTERNAL_DRAFT_SENTINEL = object() +_BUILDER_SENTINEL = object() + + +class _StubAutoModel: + """Stands in for ``AutoModelForCausalLM`` on the external-draft path.""" + + @staticmethod + def from_config(draft_config): + return _EXTERNAL_DRAFT_SENTINEL + + +def _model_config(mode, *, uses_external_draft_model=False, eagle3_model_arch="llama3"): + """Minimal duck-typed ``ModelConfig`` for dispatch-only assertions.""" + spec_config = SimpleNamespace( + spec_dec_mode=mode, + uses_external_draft_model=uses_external_draft_model, + eagle3_model_arch=eagle3_model_arch, + ) + return SimpleNamespace( + spec_config=spec_config, + pretrained_config=SimpleNamespace(num_hidden_layers=4), + ) + + +def _stub_builder(monkeypatch, mode): + """Replace ``mode``'s registered builder with a sentinel-returning stub.""" + monkeypatch.setitem( + DRAFT_MODEL_BUILDER_MAPPING, mode, lambda *args, **kwargs: _BUILDER_SENTINEL + ) + + +def test_external_draft_model_bypasses_the_registry(monkeypatch): + # An external draft model is loaded from its own checkpoint, so the + # pre-check must short-circuit before the registry is consulted at all. + monkeypatch.setattr(modeling_speculative, "AutoModelForCausalLM", _StubAutoModel) + monkeypatch.setattr( + modeling_speculative, + "get_registered_draft_model_builder", + lambda mode: pytest.fail(f"registry consulted for {mode.name} under external draft"), + ) + + result = modeling_speculative.get_draft_model( + _model_config(SpeculativeDecodingMode.MTP, uses_external_draft_model=True), + draft_config=object(), + lm_head=None, + model=None, + ) + + assert result is _EXTERNAL_DRAFT_SENTINEL + + +def test_eagle3_is_unaffected_by_the_external_draft_flag(monkeypatch): + # `uses_external_draft_model` implies `is_mtp_one_model()`, so it can never + # be true for EAGLE3. This pins the mutual exclusion that lets the pre-check run + # without a mode guard: were the property ever widened, EAGLE3 would start + # building an AutoModel drafter and this test would catch it. + monkeypatch.setattr(modeling_speculative, "AutoModelForCausalLM", _StubAutoModel) + _stub_builder(monkeypatch, SpeculativeDecodingMode.EAGLE3_ONE_MODEL) + + result = modeling_speculative.get_draft_model( + _model_config(SpeculativeDecodingMode.EAGLE3_ONE_MODEL, uses_external_draft_model=True), + draft_config=object(), + lm_head=None, + model=None, + ) + + assert result is _BUILDER_SENTINEL, "external-draft flag hijacked the EAGLE3 builder" + + +def test_external_draft_model_without_draft_config_raises(monkeypatch): + monkeypatch.setattr(modeling_speculative, "AutoModelForCausalLM", _StubAutoModel) + + with pytest.raises(ValueError, match="requires its model config"): + modeling_speculative.get_draft_model( + _model_config(SpeculativeDecodingMode.MTP, uses_external_draft_model=True), + draft_config=None, + lm_head=None, + model=None, + ) + + +def test_unregistered_mode_raises_not_implemented(): + # NGRAM is a drafter-loop mode with no one-engine draft model, so it is + # absent from both the index and the registry. + assert SpeculativeDecodingMode.NGRAM.name not in SPEC_MODE_TO_MODULE + + with pytest.raises(NotImplementedError, match="does not support speculative decoding mode"): + modeling_speculative.get_draft_model( + _model_config(SpeculativeDecodingMode.NGRAM), + draft_config=object(), + lm_head=None, + model=None, + ) + + +def test_every_indexed_mode_resolves_to_a_declaring_builder(): + # Index -> decorator direction: each indexed mode must resolve through the + # single entry point, and the builder must itself declare that mode. The + # declaration is read off the function, never by scanning the mapping by + # identity (a built-in overridden externally keeps the attribute but loses + # its slot). + for mode_name in SPEC_MODE_TO_MODULE: + mode = getattr(SpeculativeDecodingMode, mode_name, None) + assert mode is not None, f"{mode_name} is not a SpeculativeDecodingMode member" + builder = get_registered_draft_model_builder(mode) + assert builder is not None, f"no builder resolved for {mode_name}" + assert mode in getattr(builder, _REGISTERED_SPEC_MODES_ATTR, set()), ( + f"{builder.__module__}.{builder.__qualname__} is registered for " + f"{mode_name} but does not declare it" + ) + + +def test_no_builder_declares_a_mode_missing_from_the_index(): + # Decorator -> index direction: importing every indexed provider and + # walking its builders catches a mode added to an already-indexed module + # without its index entry. (A brand-new provider module is caught by the + # AST scan in tests/unittest/others/test_lazy_model_zoo.py, which needs no + # import and therefore sees modules this loop would never load.) + import importlib + + declared = set() + for module_name in set(SPEC_MODE_TO_MODULE.values()): + module = importlib.import_module(f"tensorrt_llm._torch.models.{module_name}") + for attr in vars(module).values(): + declared |= getattr(attr, _REGISTERED_SPEC_MODES_ATTR, set()) + + missing = {mode.name for mode in declared} - set(SPEC_MODE_TO_MODULE) + assert not missing, f"builders declare modes missing from _arch_index: {missing}" + + +def test_builtin_builder_does_not_override_external_registration(): + # Under lazy loading a built-in module may run its decorators *after* an + # external registration (e.g. --custom_module_dirs), so built-ins only fill + # empty slots. The reverse direction stays last-wins. + mode = SpeculativeDecodingMode.NGRAM + assert mode not in DRAFT_MODEL_BUILDER_MAPPING + + def external(model_config, draft_config, lm_head, model): + return "external" + + def builtin(model_config, draft_config, lm_head, model): + return "builtin" + + builtin.__module__ = "tensorrt_llm._torch.models.modeling_fake" + + try: + register_draft_model(mode)(external) + register_draft_model(mode)(builtin) + assert DRAFT_MODEL_BUILDER_MAPPING[mode] is external, ( + "built-in builder overrode an external registration" + ) + + del DRAFT_MODEL_BUILDER_MAPPING[mode] + register_draft_model(mode)(builtin) + register_draft_model(mode)(external) + assert DRAFT_MODEL_BUILDER_MAPPING[mode] is external + finally: + DRAFT_MODEL_BUILDER_MAPPING.pop(mode, None) + + +def test_stacked_decorators_share_one_builder(): + # Vanilla MTP and MTP_EAGLE_ONE_MODEL are one branch in + # SpeculativeDecodingMode.is_mtp_one_model(); the registry expresses that + # as two keys pointing at the same function. + mtp = get_registered_draft_model_builder(SpeculativeDecodingMode.MTP) + mtp_eagle_one = get_registered_draft_model_builder(SpeculativeDecodingMode.MTP_EAGLE_ONE_MODEL) + + assert mtp is mtp_eagle_one + declared = getattr(mtp, _REGISTERED_SPEC_MODES_ATTR, set()) + assert {SpeculativeDecodingMode.MTP, SpeculativeDecodingMode.MTP_EAGLE_ONE_MODEL} <= declared + + +def test_registry_module_placement_matches_index(): + # Builders live next to the draft model they construct, never in the + # factory file: that is what keeps get_draft_model free of concrete + # imports (and what removed the DSpark lazy import). + builder = get_registered_draft_model_builder(SpeculativeDecodingMode.DSPARK) + assert builder.__module__ == "tensorrt_llm._torch.models.modeling_dspark" + assert modeling_utils.DRAFT_MODEL_BUILDER_MAPPING is DRAFT_MODEL_BUILDER_MAPPING diff --git a/tests/unittest/others/test_lazy_model_zoo.py b/tests/unittest/others/test_lazy_model_zoo.py index e82ccf1b8e83..52d80a1fa460 100644 --- a/tests/unittest/others/test_lazy_model_zoo.py +++ b/tests/unittest/others/test_lazy_model_zoo.py @@ -179,6 +179,74 @@ def test_arch_index_matches_decorators(): assert not wrong, f"index points at the wrong module: {wrong}" +def _decorated_draft_model_registrations(): + """AST-scan the modeling files for ``@register_draft_model`` decorators. + + Source scan rather than a registry walk on purpose: the registry is only + populated by importing a provider, and a built-in builder that lost its + slot to an external registration would be missing from it entirely, so a + registry walk would pass while the index is stale. + """ + mode_to_modules = {} + for path in sorted(_MODELS_DIR.glob("*.py")): + tree = ast.parse(path.read_text()) + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + func = node.func + name = getattr(func, "id", None) or getattr(func, "attr", None) + if name != "register_draft_model" or not node.args: + continue + arg = node.args[0] + # ``register_draft_model(SpeculativeDecodingMode.DFLASH)`` + if isinstance(arg, ast.Attribute): + mode_to_modules.setdefault(arg.attr, set()).add(path.stem) + return mode_to_modules + + +def test_spec_mode_index_matches_decorators(): + from tensorrt_llm._torch.models._arch_index import SPEC_MODE_TO_MODULE + + mode_truth = _decorated_draft_model_registrations() + + missing = set(mode_truth) - set(SPEC_MODE_TO_MODULE) + assert not missing, f"spec modes missing from _arch_index: {missing}" + stale = set(SPEC_MODE_TO_MODULE) - set(mode_truth) + assert not stale, f"stale spec modes in _arch_index: {stale}" + wrong = { + mode: (SPEC_MODE_TO_MODULE[mode], mode_truth[mode]) + for mode in SPEC_MODE_TO_MODULE + if SPEC_MODE_TO_MODULE[mode] not in mode_truth[mode] + } + assert not wrong, f"index points at the wrong module: {wrong}" + + +def test_spec_mode_index_resolves_every_builder(): + # End-to-end check of the lazy path: every indexed mode must resolve to a + # builder through the single entry point, and that builder must itself + # declare the mode. The declaration is read off the function + # (``_registered_spec_modes``), never by scanning the mapping by identity: + # a built-in builder overridden by an external registration keeps its + # attribute but loses its mapping slot. + from tensorrt_llm._torch.models._arch_index import SPEC_MODE_TO_MODULE + from tensorrt_llm._torch.models.modeling_utils import ( + _REGISTERED_SPEC_MODES_ATTR, + get_registered_draft_model_builder, + ) + from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode + + for mode_name in SPEC_MODE_TO_MODULE: + mode = getattr(SpeculativeDecodingMode, mode_name, None) + assert mode is not None, f"{mode_name} is not a SpeculativeDecodingMode" + builder = get_registered_draft_model_builder(mode) + assert builder is not None, f"no draft-model builder resolved for {mode_name}" + declared = getattr(builder, _REGISTERED_SPEC_MODES_ATTR, set()) + assert mode in declared, ( + f"{builder.__module__}.{builder.__qualname__} is registered for " + f"{mode_name} but does not declare it" + ) + + def test_class_index_matches_package_all(): # MODEL_CLASS_TO_MODULE is the one table with no decorator to mirror: it # backs PEP 562 attribute access on the models package. Every name in the From 9d745fc1b965ffdf572569d04b6ddcf1a0f854d5 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Tue, 18 Aug 2026 22:26:11 -0700 Subject: [PATCH 02/21] [None][fix] Report the configured arch in the unsupported-EAGLE3 error The message read eagle3_model_arch off spec_dec_mode, a SpeculativeDecodingMode with no such attribute, so an unsupported arch raised AttributeError and the message was never shown. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/models/modeling_speculative.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 4f032756f906..94247ced2546 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -2439,19 +2439,18 @@ def external_drafter_config_kwargs(model_config, spec_config) -> dict: @register_draft_model(SpeculativeDecodingMode.EAGLE3_ONE_MODEL) def _build_eagle3_one_model_draft(model_config, draft_config, lm_head, model): """Build the EAGLE3 one-model drafter for the configured draft arch.""" - spec_dec_mode = model_config.spec_config.spec_dec_mode - if model_config.spec_config.eagle3_model_arch == "llama3": + eagle3_model_arch = model_config.spec_config.eagle3_model_arch + if eagle3_model_arch == "llama3": # Eagle3ForCausalLM handles both Llama3 and DeepSeekV3 architectures return Eagle3ForCausalLM( draft_config, model_config.pretrained_config.num_hidden_layers) - elif model_config.spec_config.eagle3_model_arch == "mistral_large3": + elif eagle3_model_arch == "mistral_large3": return MistralLarge3EagleForCausalLM( draft_config, model_config.pretrained_config.num_hidden_layers, model.aux_stream_dict) else: raise ValueError( - f"Unsupported eagle3 model architecture: {spec_dec_mode.eagle3_model_arch}" - ) + f"Unsupported eagle3 model architecture: {eagle3_model_arch}") @register_draft_model(SpeculativeDecodingMode.MTP) From 46f2d72d2758dac2d3c915f1f378313b4bbf4216 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Tue, 18 Aug 2026 22:34:58 -0700 Subject: [PATCH 03/21] [None][refactor] Move the DFlash draft model into modeling_dflash.py DFlash was 1326 of modeling_speculative.py's 2778 lines. Splitting it out leaves that module as generic speculative infrastructure plus Eagle3, PARD and MTP, and mirrors modeling_dspark.py so the two block-draft paths sit side by side. Pure move: the block references nothing else in modeling_speculative.py, and its builder registration moves with it, so neither module imports the other. The imports the block owned exclusively -- including the _flashinfer_rope try/except -- move too. The new file joins legacy-files.txt to keep the 80-column formatting the code already had, so this commit stays a move; graduating it to the ruff toolchain is a separate change. Signed-off-by: Zhenhuan Chen --- .pre-commit-config.yaml | 2 + legacy-files.txt | 1 + pyproject.toml | 1 + ruff-legacy.toml | 1 + tensorrt_llm/_torch/models/_arch_index.py | 2 +- tensorrt_llm/_torch/models/modeling_dflash.py | 1401 +++++++++++++++++ .../_torch/models/modeling_speculative.py | 1388 +--------------- .../modeling/test_modeling_speculative.py | 6 +- .../test_kimi_k3_dspark_semantics.py | 2 +- 9 files changed, 1412 insertions(+), 1392 deletions(-) create mode 100644 tensorrt_llm/_torch/models/modeling_dflash.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 2a259562ce95..182b085aed35 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -202,6 +202,7 @@ common-files: &common_files | tensorrt_llm/_torch/models/modeling_bert.py | tensorrt_llm/_torch/models/modeling_clip.py | tensorrt_llm/_torch/models/modeling_deepseekv3.py | + tensorrt_llm/_torch/models/modeling_dflash.py | tensorrt_llm/_torch/models/modeling_exaone4.py | tensorrt_llm/_torch/models/modeling_gemma3.py | tensorrt_llm/_torch/models/modeling_gemma3vl.py | @@ -968,6 +969,7 @@ legacy-files: &legacy_files | tensorrt_llm/_torch/models/modeling_bert.py | tensorrt_llm/_torch/models/modeling_clip.py | tensorrt_llm/_torch/models/modeling_deepseekv3.py | + tensorrt_llm/_torch/models/modeling_dflash.py | tensorrt_llm/_torch/models/modeling_exaone4.py | tensorrt_llm/_torch/models/modeling_gemma3.py | tensorrt_llm/_torch/models/modeling_gemma3vl.py | diff --git a/legacy-files.txt b/legacy-files.txt index 73fe3fee5899..b3c345ff143f 100644 --- a/legacy-files.txt +++ b/legacy-files.txt @@ -194,6 +194,7 @@ tensorrt_llm/_torch/models/modeling_auto.py tensorrt_llm/_torch/models/modeling_bert.py tensorrt_llm/_torch/models/modeling_clip.py tensorrt_llm/_torch/models/modeling_deepseekv3.py +tensorrt_llm/_torch/models/modeling_dflash.py tensorrt_llm/_torch/models/modeling_exaone4.py tensorrt_llm/_torch/models/modeling_gemma3.py tensorrt_llm/_torch/models/modeling_gemma3vl.py diff --git a/pyproject.toml b/pyproject.toml index f8e1447ca03b..f25e25193644 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -251,6 +251,7 @@ exclude = [ "tensorrt_llm/_torch/models/modeling_bert.py", "tensorrt_llm/_torch/models/modeling_clip.py", "tensorrt_llm/_torch/models/modeling_deepseekv3.py", + "tensorrt_llm/_torch/models/modeling_dflash.py", "tensorrt_llm/_torch/models/modeling_exaone4.py", "tensorrt_llm/_torch/models/modeling_gemma3.py", "tensorrt_llm/_torch/models/modeling_gemma3vl.py", diff --git a/ruff-legacy.toml b/ruff-legacy.toml index 612b315564ca..43715109db9b 100644 --- a/ruff-legacy.toml +++ b/ruff-legacy.toml @@ -211,6 +211,7 @@ include = [ "tensorrt_llm/_torch/models/modeling_bert.py", "tensorrt_llm/_torch/models/modeling_clip.py", "tensorrt_llm/_torch/models/modeling_deepseekv3.py", + "tensorrt_llm/_torch/models/modeling_dflash.py", "tensorrt_llm/_torch/models/modeling_exaone4.py", "tensorrt_llm/_torch/models/modeling_gemma3.py", "tensorrt_llm/_torch/models/modeling_gemma3vl.py", diff --git a/tensorrt_llm/_torch/models/_arch_index.py b/tensorrt_llm/_torch/models/_arch_index.py index 7c04d8d0a450..f2f278dea788 100644 --- a/tensorrt_llm/_torch/models/_arch_index.py +++ b/tensorrt_llm/_torch/models/_arch_index.py @@ -271,7 +271,7 @@ def is_builtin_zoo_module(module_name: str) -> bool: # slot to an external registration is absent from the mapping but still has # the attribute, so an identity scan reports it as unregistered. SPEC_MODE_TO_MODULE = { - "DFLASH": "modeling_speculative", + "DFLASH": "modeling_dflash", "DRAFT_TARGET_ONE_MODEL": "modeling_speculative", "DSPARK": "modeling_dspark", "EAGLE3_ONE_MODEL": "modeling_speculative", diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py new file mode 100644 index 000000000000..d6e1cb4dccf2 --- /dev/null +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -0,0 +1,1401 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +from dataclasses import replace +from typing import Dict, Optional + +import torch +import torch.nn.functional as F +from torch import nn +from transformers import PretrainedConfig + +from tensorrt_llm.logger import logger + +from ...functional import RotaryScalingType +from ..modules.rotary_embedding import RotaryEmbedding + +try: + from ..custom_ops import \ + flashinfer_apply_rope_with_cos_sin_cache_inplace as _flashinfer_rope +except ImportError: + _flashinfer_rope = None +from ..pyexecutor.config_utils import (_is_sliding_attention_layer, + get_layer_attention_window) +from ..speculative.dflash_attention import (get_dflash_flash_attention, + get_dflash_trtllm_gen_ops) +from ..speculative.interface import SpeculativeDecodingMode +from .modeling_utils import get_model_architecture, register_draft_model + + +def dspark_layer_window_size(use_swa: bool, swa_window: int, layer_types, + layer_idx: int) -> tuple[int, int]: + """flash-attn ``window_size`` for one draft layer of the block decode. + + DSpark drafters (deepseek-ai/DeepSpec) run the draft block through HF + attention with ``sliding_window`` set on 'sliding_attention' layers and + is_causal=False. HF's flash path + (transformers/modeling_flash_attention_utils.py) translates that to + ``window_size = (sliding_window - 1, sliding_window - 1)``, i.e. each + query attends keys within ``swa_window - 1`` KV-index distance on both + sides. In the DFlash pool layout KV index == token position, so this + limits draft queries to the most recent ``swa_window`` context tokens + plus the (nearby) draft block. Full-attention layers and non-dspark + drafters keep flash-attn's default ``(-1, -1)`` (no window). + """ + if not use_swa: + return (-1, -1) + if layer_types is not None and layer_idx < len(layer_types) and \ + layer_types[layer_idx] != 'sliding_attention': + return (-1, -1) + return (swa_window - 1, swa_window - 1) + + +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. + + Reference: DeepSpec ``VanillaMarkov`` (deepspec/modeling/dspark/ + markov_head.py): ``bias = markov_w2(markov_w1(prev_token))`` where + markov_w1 is nn.Embedding(vocab, rank) and markov_w2 is + nn.Linear(rank, vocab, bias=False). With both weights stored + [vocab, rank] this is ``markov_w1[prev] @ markov_w2.T``. + + Args: + prev_tokens: [B] long, previous token per request (draft vocab). + markov_w1: [vocab, rank]. + markov_w2: [vocab_or_shard, rank] (rows may be a TP vocab shard). + Returns: + [B, vocab_or_shard] bias in the markov weights' dtype. + """ + return F.linear(F.embedding(prev_tokens, markov_w1), markov_w2) + + +def dspark_markov_chain_logits( + base_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + markov_w1: torch.Tensor, + markov_w2: torch.Tensor, + argmax_fn=None, +) -> torch.Tensor: + """Apply the vanilla Markov intra-block bias across a drafted block. + + Reference: DeepSpec ``VanillaMarkov.sample_block_tokens`` at + temperature 0: for step i, ``logits_i += bias(prev_i)`` with + ``prev_0`` = the anchor token (last accepted token, block slot 0) and + ``prev_{i>0}`` = the greedy token from step i-1's *biased* logits. + llama.cpp PR #25173 implements the same greedy chain. + + Args: + base_logits: [B, K, vocab_or_shard] shared-lm_head logits. + first_prev_tokens: [B] long, anchor token ids (draft vocab). + markov_w1 / markov_w2: see :func:`dspark_markov_step_bias`. + argmax_fn: callable([B, vocab_or_shard]) -> [B] token ids in the + full draft vocab; defaults to plain argmax. Workers pass a + TP-aware argmax when the draft logits are vocab-sharded. + Returns: + [B, K, vocab_or_shard] biased logits. Greedy per-position argmax of + the result reproduces the reference sampled chain exactly. + """ + K = base_logits.shape[1] + if K == 0: + return base_logits + prev = first_prev_tokens.long() + steps = [] + for i in range(K): + bias = dspark_markov_step_bias(prev, markov_w1, markov_w2) + step_logits = base_logits[:, i] + bias.to(base_logits.dtype) + steps.append(step_logits) + if argmax_fn is not None: + prev = argmax_fn(step_logits).long() + else: + prev = torch.argmax(step_logits, dim=-1) + return torch.stack(steps, dim=1) + + +class DFlashForCausalLM(nn.Module): + """Draft model wrapper for DFlash speculative decoding. + + DFlash uses cross-attention where Q comes from noise/query tokens and K/V + come from the concatenation of target hidden states and noise hidden states. + The target_hidden stays CONSTANT across all layers (no input_layernorm applied). + + Reference: https://arxiv.org/pdf/2602.06036 + """ + + 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 + custom DFlash architecture label).""" + super().__init__() + + pretrained_cfg = draft_config.pretrained_config + try: + DraftModelClass, _ = get_model_architecture(pretrained_cfg) + except RuntimeError: + model_type = pretrained_cfg.model_type + arch_name = "".join(w.capitalize() + for w in model_type.split("_")) + "ForCausalLM" + logger.info( + f"DFlash: architecture {pretrained_cfg.architectures} not found, " + f"falling back to {arch_name} based on model_type={model_type}") + original_archs = pretrained_cfg.architectures + try: + pretrained_cfg.architectures = [arch_name] + DraftModelClass, _ = get_model_architecture(pretrained_cfg) + finally: + pretrained_cfg.architectures = original_archs + + # Remove spec_config to prevent recursive spec-dec initialization + draft_config_no_spec = replace(draft_config, + spec_config=None, + lm_head_gather_output=False) + + # Weights will be loaded later by ModelLoader.load_draft_weights() + self.draft_model_full = DraftModelClass(draft_config_no_spec) + self.model = self.draft_model_full.model + self.lm_head = self.draft_model_full.lm_head + + # Required by weight mappers + self.model_config = draft_config_no_spec + self.config = draft_config_no_spec.pretrained_config + + # Get mask_token_id from dflash_config + pretrained_config = draft_config.pretrained_config + dflash_config = getattr(pretrained_config, 'dflash_config', {}) + self.mask_token_id = dflash_config.get( + 'mask_token_id', + getattr(pretrained_config, 'mask_token_id', + pretrained_config.vocab_size)) + + self.target_layer_ids = dflash_config.get('target_layer_ids', None) + self.block_size = dflash_config.get( + 'block_size', getattr(pretrained_config, 'block_size', None)) + self.dflash_attention_backend = dflash_attention_backend + if 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() + else: + raise ValueError( + "DFlash attention backend must be VANILLA or TRTLLM, 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) + self.register_buffer("_dflash_block_offsets", None, persistent=False) + self._dflash_trtllm_gen_device = None + self._dflash_trtllm_gen_sm_count = None + logger.info( + f"DFlash draft model initialized with mask_token_id: {self.mask_token_id}, " + f"target_layer_ids: {self.target_layer_ids}, block_size: {self.block_size}, " + f"attention_backend: {self.dflash_attention_backend}") + + # DSpark drafters (DFlash + low-rank Markov head + confidence head, + # arXiv 2607.05147; reference: deepseek-ai/DeepSpec). The weights- + # independent drafter-forward semantics ARE implemented here: + # - vanilla Markov intra-block logit bias (applied by DFlashWorker + # through apply_markov_chain_logits), + # - sliding-window attention on 'sliding_attention' draft layers + # during the block decode (use_swa / swa_window_size), + # - the shift_label output convention (hidden state at block slot j + # predicts draft token j+1; slot 0 holds the anchor token). + # Confidence-scheduled verification is NOT implemented yet: the + # confidence_proj weights are loaded (for the follow-up MR) but never + # used, and drafting always proposes the full K tokens. + self._dspark_shift_label = bool(dflash_config.get('shift_label', False)) + self._dspark_use_swa = bool(dflash_config.get('use_swa', False)) + self._dspark_swa_window = int( + dflash_config.get('swa_window_size', 0) or 0) + self._dspark_markov_rank = int(dflash_config.get('markov_rank', 0) or 0) + self._dspark_markov_head_type = str( + dflash_config.get('markov_head_type', 'vanilla') + or 'vanilla').lower() + self._dspark_use_confidence_head = bool( + dflash_config.get('use_confidence_head', False)) + # Plain None placeholders rather than nn.Parameter/buffer: most + # DFlash checkpoints don't ship these heads, and their shapes + # ([vocab, rank]) are checkpoint-dependent, so nothing is + # pre-allocated. load_weights() fills them in only when the + # checkpoint ships them; consumers treat None as "head absent". + self.markov_w1 = None # [vocab, rank] (nn.Embedding weight layout) + self.markov_w2 = None # [vocab, rank] (nn.Linear(rank->vocab) weight) + self.confidence_proj_weight = None # loaded, unused (follow-up MR) + self.confidence_proj_bias = None + + if self._dspark_markov_rank > 0 and \ + self._dspark_markov_head_type != 'vanilla': + raise ValueError( + f"DFlash dspark drafter declares markov_head_type=" + f"'{self._dspark_markov_head_type}'; only 'vanilla' is " + "supported (gated/rnn heads need per-step hidden features).") + if self._dspark_use_swa and self._dspark_swa_window < 1: + raise ValueError( + "DFlash dspark drafter sets use_swa but swa_window_size=" + f"{dflash_config.get('swa_window_size')} is invalid.") + # causal=true is only invalid under the dspark convention. Legacy + # DFlash drafter configs (e.g. Laguna) also carry a causal field; + # their causality is handled by the legacy decode path + # (_sliding_layers_causal), so don't reject them here. + is_dspark = (str(dflash_config.get('projector_type', '') + or '').lower() == 'dspark' or self._dspark_shift_label + or self._dspark_use_swa or self._dspark_markov_rank > 0 + or self._dspark_use_confidence_head) + if is_dspark and dflash_config.get('causal'): + raise ValueError( + "DFlash dspark drafter sets causal=true; the block decode " + "only supports the non-causal dspark convention.") + # Per-layer flash-attn window for the block decode, resolved once. + num_draft_layers = getattr(pretrained_config, 'num_hidden_layers', 0) + layer_types = getattr(pretrained_config, 'layer_types', None) + self._dspark_layer_windows = [ + dspark_layer_window_size(self._dspark_use_swa, + self._dspark_swa_window, layer_types, i) + for i in range(num_draft_layers) + ] + if self._dspark_use_confidence_head: + logger.warning( + "DFlash dspark drafter declares use_confidence_head; " + "confidence-scheduled verification is not implemented yet " + "(confidence_proj weights are loaded but unused, drafting " + "always proposes the full K tokens).") + + self.logits_processor = None # Set by caller after construction + + # RoPE - lazily initialized from draft model's attention module + self._rope_initialized = False + self._rotary_cos_sin = None + self._is_neox = True + + self._cos_sin_cache_fp32 = None + self._rope_dummy_q = None + + # Lazy-built after weights load (see _build_fused_kv_buffers). + self._fused_kv_weight = None + self._fused_kv_bias = None + self._k_norm_stacked = None + self._k_norm_eps = None + self._num_attn_layers = 0 + self._num_heads = 0 + self._head_dim = 0 + self._num_kv_heads = 0 + self._has_qk_norm = False + self._use_fused_qk_norm_rope = False + # Laguna-specific draft-layer behaviors, disabled by default so generic + # DFlash drafters keep the original contract (no context input_layernorm, + # non-causal block attention). Subclasses opt in. + self._context_input_layernorm = False + self._sliding_layers_causal = False + self._warn_inferred_attention_windows() + + @staticmethod + def _rope_signature(attn): + """Return the effective RoPE configuration used by an attention layer.""" + if attn.rotary_emb is not None: + return ( + attn.rotary_emb.rope_params, + attn.rotary_emb.head_dim, + attn.rotary_emb.is_neox, + ) + if attn.pos_embd_params is not None: + return ( + attn.pos_embd_params.rope, + attn.head_dim, + attn.pos_embd_params.is_neox, + ) + return None + + def _validate_uniform_rope(self): + """Check that all draft layers can safely share one RoPE cache.""" + if len(self.model.layers) == 0: + raise ValueError("DFlash requires at least one draft model layer.") + + signatures = [ + self._rope_signature(layer.self_attn) for layer in self.model.layers + ] + + mismatched_layers = [ + layer_idx + for layer_idx, signature in enumerate(signatures[1:], start=1) + if signature != signatures[0] + ] + if mismatched_layers: + layer_types = getattr(self.config, 'layer_types', None) + raise ValueError( + "DFlash shares one RoPE cache across draft layers, but layers " + f"{mismatched_layers} have a different effective RoPE " + f"configuration from layer 0. layer_types={layer_types}.") + + def _init_rope(self): + """Initialize RoPE from the draft model's attention configuration. + + Reuses the existing RotaryEmbedding infrastructure which correctly + handles all RoPE variants (standard, YaRN, scaled, etc.). + """ + # The flattened context-KV path shares layer 0's RoPE cache. + self._validate_uniform_rope() + attn0 = self.model.layers[0].self_attn + + if attn0.rotary_emb is not None: + self._rotary_cos_sin = attn0.rotary_emb.rotary_cos_sin + self._is_neox = attn0.rotary_emb.is_neox + elif attn0.pos_embd_params is not None: + rope_emb = RotaryEmbedding( + attn0.pos_embd_params.rope, + head_dim=attn0.head_dim, + is_neox=attn0.pos_embd_params.is_neox, + ) + self._rotary_cos_sin = rope_emb.rotary_cos_sin + self._is_neox = rope_emb.is_neox + else: + # Fallback: basic NeoX-style RoPE + config = self.config + head_dim = getattr(config, 'head_dim', + config.hidden_size // config.num_attention_heads) + rope_theta = getattr(config, 'rope_theta', 1000000.0) + max_pos = getattr(config, 'max_position_embeddings', 32768) + + inv_freq = 1.0 / (rope_theta**(torch.arange( + 0, head_dim, 2, dtype=torch.float32, device='cuda') / head_dim)) + positions = torch.arange(max_pos, + dtype=torch.float32, + device='cuda') + freqs = torch.outer(positions, inv_freq) + rope_cos = freqs.cos().to(config.torch_dtype) + rope_sin = freqs.sin().to(config.torch_dtype) + # [max_pos, 2, rot_dim//2] to match RotaryEmbedding format + self._rotary_cos_sin = torch.stack([rope_cos, rope_sin], dim=1) + self._is_neox = True + + self._rope_initialized = True + + def project_target_hidden(self, + hidden_states: torch.Tensor) -> torch.Tensor: + """Project captured target hidden states into the draft hidden space. + + Generic DFlash: fc then hidden_norm. Subclasses (e.g. Laguna) may + normalize the per-aux features first by overriding this method. + """ + hidden_states = hidden_states.to(self.fc.weight.dtype) + return self.hidden_norm(self.fc(hidden_states)) + + @property + def has_markov_head(self) -> bool: + return self._dspark_markov_rank > 0 and self.markov_w1 is not None + + def apply_markov_chain_logits( + self, + base_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + argmax_fn=None, + vocab_slice: slice | None = None) -> torch.Tensor: + """Apply the dspark vanilla-Markov intra-block bias to block logits. + + No-op (returns ``base_logits`` unchanged) for non-dspark drafters. + See :func:`dspark_markov_chain_logits` for the semantics; when + ``base_logits`` is a TP vocab shard, the caller must pass this + rank's ``vocab_slice`` (to shard the markov_w2 rows identically) + and an ``argmax_fn`` returning full-vocab token ids — DFlashWorker + handles both. + """ + if not self.has_markov_head: + return base_logits + markov_w2 = self.markov_w2 if vocab_slice is None else \ + self.markov_w2[vocab_slice] + return dspark_markov_chain_logits(base_logits, + first_prev_tokens, + self.markov_w1, + markov_w2, + argmax_fn=argmax_fn) + + def _post_attention_gate(self, attn_output, gate_input, attn_mod, num_heads, + head_dim): + """Hook applied to the block-attention output before o_proj. + + No-op for generic DFlash; overridden by drafters that gate (e.g. Laguna). + """ + return attn_output + + def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): + """Load weights into the DFlash draft model. + + DFlash checkpoints differ from standard HF format: + - Layer weights lack the 'model.' prefix (e.g., 'layers.0...' not 'model.layers.0...') + - Extra DFlash-specific weights: 'fc.weight', 'hidden_norm.weight' + - Missing embed_tokens and lm_head (shared with target model) + """ + # Laguna DFlash checkpoints may ship a fused self_attn.qkv_proj; the draft + # loader expects split q/k/v (a fused key is silently dropped otherwise). + if any(k.endswith('self_attn.qkv_proj.weight') for k in weights): + for attr in ('num_attention_heads_per_layer', + 'num_key_value_heads_per_layer'): + per_layer = getattr(self.config, attr, None) + if per_layer is not None and len(set(per_layer)) > 1: + raise ValueError( + "DFlash load_weights() splits the fused qkv_proj using " + "the global head count, but the drafter has heterogeneous " + f"{attr} {sorted(set(per_layer))}; per-layer qkv splitting " + "is required for this checkpoint.") + head_dim = getattr( + self.config, 'head_dim', + self.config.hidden_size // self.config.num_attention_heads) + num_kv_heads = getattr(self.config, 'num_key_value_heads', + self.config.num_attention_heads) + q = self.config.num_attention_heads * head_dim + kv = num_kv_heads * head_dim + split = {} + for k, v in weights.items(): + if k.endswith('self_attn.qkv_proj.weight'): + b = k[:-len('qkv_proj.weight')] + split[b + 'q_proj.weight'] = v[:q] + split[b + 'k_proj.weight'] = v[q:q + kv] + split[b + 'v_proj.weight'] = v[q + kv:] + else: + split[k] = v + weights = split + + # DSpark head weights: keep them out of the backbone remap (they'd + # get a 'model.' prefix and be dropped by allow_partial_loading). + # markov_w1/markov_w2 drive the intra-block logit bias; the + # confidence_proj weights are loaded for the confidence-scheduling + # follow-up MR but are not used yet. + dspark_keys = ('markov_w1.weight', 'markov_w2.weight', + 'confidence_proj.weight', 'confidence_proj.bias') + dspark_weights = {k: weights[k] for k in dspark_keys if k in weights} + if dspark_weights: + weights = { + k: v + for k, v in weights.items() if k not in dspark_weights + } + if self._dspark_markov_rank > 0: + vocab = self.config.vocab_size + rank = self._dspark_markov_rank + for k in ('markov_w1.weight', 'markov_w2.weight'): + if k not in dspark_weights: + raise ValueError( + f"DFlash dspark drafter declares markov_rank=" + f"{self._dspark_markov_rank} but the checkpoint is " + f"missing {k}.") + if tuple(dspark_weights[k].shape) != (vocab, rank): + raise ValueError( + f"DFlash dspark {k} has shape " + f"{tuple(dspark_weights[k].shape)}, expected " + f"[vocab, markov_rank] = ({vocab}, {rank}).") + self.markov_w1 = dspark_weights['markov_w1.weight'].to('cuda') + self.markov_w2 = dspark_weights['markov_w2.weight'].to('cuda') + if 'confidence_proj.weight' in dspark_weights: + 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') + + # Remap: add 'model.' prefix where needed, and extract DFlash-specific weights + remapped = {} + for key, value in weights.items(): + if key in ('fc.weight', 'hidden_norm.weight'): + # DFlash-specific projection weights - store directly + remapped[key] = value + elif key == 'norm.weight': + remapped['model.norm.weight'] = value + elif not key.startswith('model.'): + remapped[f'model.{key}'] = value + else: + remapped[key] = value + + # Load DFlash-specific weights directly + if 'fc.weight' in remapped: + self.fc = nn.Linear(remapped['fc.weight'].shape[1], + remapped['fc.weight'].shape[0], + bias=False, + device='cuda', + dtype=remapped['fc.weight'].dtype) + self.fc.weight.data.copy_(remapped['fc.weight']) + del remapped['fc.weight'] + + if 'hidden_norm.weight' in remapped: + rms_norm_eps = getattr(self.config, 'rms_norm_eps', 1e-6) + self.hidden_norm = nn.RMSNorm( + remapped['hidden_norm.weight'].shape[0], + eps=rms_norm_eps, + device='cuda', + elementwise_affine=True, + dtype=remapped['hidden_norm.weight'].dtype) + 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. + self.draft_model_full.load_weights(weights=remapped, + weight_mapper=weight_mapper, + allow_partial_loading=True) + + def load_weights_from_target_model(self, + target_model: torch.nn.Module) -> None: + """Share embed_tokens and lm_head from the target model.""" + self.draft_model_full.model.embed_tokens = target_model.model.embed_tokens + self.draft_model_full.lm_head = target_model.lm_head + self.lm_head = target_model.lm_head + + def precompute_context_kv( + self, + projected_hidden: torch.Tensor, + positions: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Post-norm / post-RoPE K and V for ALL drafter layers in one fused GEMM. + + Args: + projected_hidden: [N, hidden_size], already fc + hidden_norm'd. + positions: [N] int32/64, RoPE positions for each entry. + Returns: + k: [N, L, nkv, hd] post k_norm and RoPE + v: [N, L, nkv, hd] post split only + """ + if self._fused_kv_weight is None: + self._build_fused_kv_buffers() + N = projected_hidden.shape[0] + L = self._num_attn_layers + nkv = self._num_kv_heads + hd = self._head_dim + weight_dtype = self._fused_kv_weight.dtype + if getattr(self, '_input_ln_eps', None) is not None: + ph = projected_hidden.float() + ph = ph * torch.rsqrt( + ph.pow(2).mean(-1, keepdim=True) + self._input_ln_eps) + projected_hidden = ph.to(weight_dtype) + elif projected_hidden.dtype != weight_dtype: + projected_hidden = projected_hidden.to(weight_dtype) + + kv_flat = F.linear(projected_hidden, self._fused_kv_weight, + self._fused_kv_bias) + # Per-layer layout [L0_K|L0_V|L1_K|L1_V|...] keeps K and V contiguous + # after the select() splits — no extra copy required. + kv = kv_flat.view(N, L, 2, nkv, hd) + k = kv[:, :, 0].contiguous() + v = kv[:, :, 1].contiguous() + + if self._k_norm_stacked is not None: + # Fuse L per-layer RMSNorms into one. k is [N, L, nkv, hd]; + # each layer has its own weight ([L, hd]) but shares eps. + k = F.rms_norm(k, (hd, ), eps=self._k_norm_eps) + k = k * self._k_norm_stacked.view(1, L, 1, hd) + + self._fused_rope_inplace(k.view(N * L, nkv * hd), positions, N, L) + return k, v + + def _get_cos_sin_cache(self) -> torch.Tensor: + """Return the flashinfer-style cos/sin cache for the drafter. + + Shape [max_positions, head_dim], fp32 — flashinfer's + apply_rope_with_cos_sin_cache_inplace requires fp32 regardless of + the query/key dtype. + """ + if self._cos_sin_cache_fp32 is not None: + return self._cos_sin_cache_fp32 + if not self._rope_initialized: + self._init_rope() + max_pos = self._rotary_cos_sin.shape[0] + self._cos_sin_cache_fp32 = self._rotary_cos_sin.view(max_pos, -1).to( + torch.float32).contiguous() + return self._cos_sin_cache_fp32 + + def _fused_rope_inplace( + self, + k_flat: torch.Tensor, + positions: torch.Tensor, + N: int, + L: int, + ) -> None: + """In-place fused RoPE over [N*L, nkv*hd] K values. + + Layout of k_flat: row (i*L + l) holds layer l of position i, so + positions must be repeat_interleaved by L to match. + """ + positions_int32 = positions.view(-1).to(torch.int32) + if L > 1: + positions_int32 = positions_int32.repeat_interleave(L) + + if _flashinfer_rope is not None: + # flashinfer requires a non-None query tensor; pass a single-head + # scratch so the extra rotate is negligible. + need_rows = k_flat.shape[0] + dummy_q = self._rope_dummy_q + if (dummy_q is None or dummy_q.dtype != k_flat.dtype + or dummy_q.shape[0] < need_rows): + dummy_q = k_flat.new_empty(need_rows, self._head_dim) + self._rope_dummy_q = dummy_q + _flashinfer_rope( + positions_int32, + dummy_q[:need_rows], + k_flat, + self._head_dim, + self._get_cos_sin_cache(), + self._is_neox, + ) + return + + # Pure-PyTorch fallback (older environments without flashinfer). + cos, sin = self._get_rope_cos_sin(positions_int32.view(1, -1), + dtype=k_flat.dtype) + k_roped = RotaryEmbedding.apply_rotary_pos_emb( + k_flat.view(k_flat.shape[0], -1, self._head_dim), + cos.squeeze(0), + sin.squeeze(0), + unsqueeze_dim=1, + is_neox=self._is_neox, + ) + k_flat.copy_(k_roped.view_as(k_flat)) + + def _build_fused_kv_buffers(self) -> None: + """Stack per-layer KV projection + k_norm weights for a single fused GEMM. + + Must run after weights are loaded. + """ + if self._fused_kv_weight is not None: + return + layers_attn = [layer.self_attn for layer in self.model.layers] + attn0 = layers_attn[0] + q_size = attn0.q_size + kv_size = attn0.kv_size + head_dim = attn0.head_dim + num_heads = attn0.num_heads + num_kv_heads = attn0.num_key_value_heads + # Head counts are read from layer 0 here and in dflash_forward; assert + # uniformity (the target uses per-layer heads, the drafter does not). + for a in layers_attn[1:]: + assert ( + a.q_size == q_size and a.kv_size == kv_size + and a.head_dim == head_dim and a.num_heads == num_heads + and a.num_key_value_heads == num_kv_heads), ( + "DFlash fused KV requires all drafter layers to share " + "q_size / kv_size / head_dim / num_heads / num_kv_heads.") + + has_k_norm = [hasattr(a, 'k_norm') for a in layers_attn] + assert all(has_k_norm) or not any(has_k_norm), ( + "DFlash fused KV requires either all or no drafter layers to have k_norm." + ) + + kv_weights = [ + a.qkv_proj.weight[q_size:q_size + 2 * kv_size] for a in layers_attn + ] + # Fold each drafter layer's input_layernorm weight into its KV projection + # so context K/V match the query path. vLLM laguna_dflash applies + # layer.input_layernorm to context states before KV; RMSNorm gives + # (x_hat * w) @ Wkv.T == x_hat @ (Wkv * w).T, and the shared 1/rms(x) is + # applied to projected_hidden in precompute_context_kv. + dlayers = self.model.layers + if self._context_input_layernorm and all( + hasattr(dl, 'input_layernorm') for dl in dlayers): + eps_set = { + getattr(dl.input_layernorm, 'variance_epsilon', + getattr(self.config, 'rms_norm_eps', 1e-6)) + for dl in dlayers + } + assert len(eps_set) == 1, ( + "DFlash fused context input_layernorm needs all drafter layers " + f"to share variance_epsilon; got {sorted(eps_set)}") + self._input_ln_eps = eps_set.pop() + folded = [] + for w, dl in zip(kv_weights, dlayers): + scale = dl.input_layernorm.weight.data + if getattr(dl.input_layernorm, 'use_gemma', False): + scale = scale + 1 + folded.append(w * scale[None, :].to(w.dtype)) + kv_weights = folded + else: + self._input_ln_eps = None + fused_kv_weight = torch.cat(kv_weights, dim=0).contiguous() + if attn0.qkv_proj.bias is not None: + kv_biases = [ + a.qkv_proj.bias[q_size:q_size + 2 * kv_size] + for a in layers_attn + ] + self._fused_kv_bias = torch.cat(kv_biases, dim=0).contiguous() + else: + self._fused_kv_bias = None + + if all(has_k_norm): + k_norm0 = layers_attn[0].k_norm + eps = k_norm0.variance_epsilon + eps_set = {a.k_norm.variance_epsilon for a in layers_attn} + assert len(eps_set) == 1, ( + f"DFlash fused k_norm requires all drafter layers to share " + f"variance_epsilon; got {sorted(eps_set)}.") + self._k_norm_stacked = torch.stack( + [a.k_norm.weight.data for a in layers_attn]) + self._k_norm_eps = eps + else: + self._k_norm_stacked = None + self._k_norm_eps = None + self._num_attn_layers = len(layers_attn) + self._num_heads = num_heads + self._head_dim = head_dim + self._num_kv_heads = num_kv_heads + self._fused_kv_weight = fused_kv_weight + + # fused_qk_norm_rope derives YaRN / partial-rotary frequencies on + # the fly, which can disagree with precompute_context_kv's cached + # cos/sin. Only enable it when the drafter uses plain RoPE. + self._has_qk_norm = (all(has_k_norm) + and all(hasattr(a, 'q_norm') for a in layers_attn)) + rope_params = getattr(getattr(attn0, 'pos_embd_params', None), 'rope', + None) + scale_type = getattr(rope_params, 'scale_type', None) + partial_rotary_factor = getattr( + getattr(attn0, 'pretrained_config', None), 'partial_rotary_factor', + 1.0) + self._use_fused_qk_norm_rope = (self._has_qk_norm + and hasattr(attn0, 'apply_qk_norm_rope') + and rope_params is not None + and scale_type + in (None, RotaryScalingType.none) + and partial_rotary_factor == 1.0) + + logger.debug( + f"DFlash: fused KV weights built for {self._num_attn_layers} layers " + f"(fused_kv_weight shape={tuple(self._fused_kv_weight.shape)})") + + def _get_rope_cos_sin(self, positions, dtype=None): + """Get cos/sin for given positions, suitable for apply_rotary_pos_emb. + + Args: + positions: [B, seq_len] + dtype: target dtype for cos/sin (default: keep original) + Returns: + rope_cos: [B, seq, rot_dim//2] (broadcastable with unsqueeze_dim=1) + rope_sin: [B, seq, rot_dim//2] + """ + if not self._rope_initialized: + self._init_rope() + + # rotary_cos_sin: [max_pos, 2, rot_dim//2] + rope_cache = self._rotary_cos_sin[positions] # [B, seq, 2, rot_dim//2] + rope_cos = rope_cache[..., 0, :] # [B, seq, rot_dim//2] + rope_sin = rope_cache[..., 1, :] + if dtype is not None: + rope_cos = rope_cos.to(dtype) + rope_sin = rope_sin.to(dtype) + return rope_cos, rope_sin + + def _warn_inferred_attention_windows(self) -> None: + """Warn once at initialization when checkpoint metadata enables SWA.""" + if getattr(self.config, 'use_sliding_window', None) is not None: + return + + num_hidden_layers = getattr(self.config, 'num_hidden_layers', None) + if num_hidden_layers is None: + num_hidden_layers = len(self.model.layers) + layers_by_window = {} + for layer_idx in range(num_hidden_layers): + window = get_layer_attention_window(self.config, layer_idx) + if window is not None: + layers_by_window.setdefault(window, []).append(layer_idx) + + for window, layer_indices in layers_by_window.items(): + logger.warning( + "DFlash inferred pooled-context sliding-window attention from " + f"checkpoint config for draft layers {layer_indices}: " + f"window={window}. Context attention is truncated to {window} " + "tokens for these layers; if the drafter expects full context, " + "acceptance rate may drop. Set use_sliding_window explicitly " + "to confirm or disable windowing.") + + def _get_attention_mask_args(self, layer_idx): + """Return FlashAttention causal and local-window arguments for a layer.""" + layer_types = getattr(self.config, 'layer_types', None) + is_sliding_layer = False + if layer_types: + layer_type = layer_types[layer_idx % len(layer_types)] + is_sliding_layer = _is_sliding_attention_layer(layer_type) + + sliding_window = get_layer_attention_window(self.config, layer_idx) + is_sliding_layer = is_sliding_layer or sliding_window is not None + if not is_sliding_layer: + return False, (-1, -1) + + causal = self._sliding_layers_causal or sliding_window is not None + if sliding_window is None: + # Legacy drafters without an explicit window preserve their prior + # non-windowed behavior. + return causal, (-1, -1) + # FlashAttention's bounds are inclusive: W tokens are current + W-1 left. + return causal, (sliding_window - 1, 0) + + def _prepare_dflash_trtllm_gen_buffers( + self, + dtype: torch.dtype, + device: torch.device, + max_batch_size: int, + block_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + ) -> None: + trtllm_gen_ops = self._dflash_trtllm_gen_ops + workspace_bytes = trtllm_gen_ops.get_workspace_size( + dtype=dtype, + num_tokens=max_batch_size * block_size, + num_gen_tokens=max_batch_size * block_size, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_size=head_dim, + max_num_requests=max_batch_size, + rotary_embedding_dim=0, + fp8_context_fmha=False, + ) + device = torch.device(device) + is_capturing = torch.cuda.is_current_stream_capturing() + if self._dflash_trtllm_gen_device != device: + if is_capturing: + raise RuntimeError( + "DFlash TRTLLM-Gen buffers must be prepared on the current " + "device before CUDA graph capture.") + self._dflash_trtllm_gen_device = device + self._dflash_trtllm_gen_sm_count = ( + torch.cuda.get_device_properties(device).multi_processor_count) + + workspace = self._dflash_trtllm_gen_workspace + workspace_needs_allocation = ( + workspace is None or workspace.device != device + or workspace.numel() * workspace.element_size() < workspace_bytes) + if workspace_needs_allocation: + if is_capturing: + raise RuntimeError( + "The DFlash TRTLLM-Gen workspace must be allocated at the " + "required size before CUDA graph capture.") + self._dflash_trtllm_gen_workspace = torch.empty(workspace_bytes, + dtype=torch.uint8, + device=device) + + sm_count = self._dflash_trtllm_gen_sm_count + counter_bytes = trtllm_gen_ops.get_multi_ctas_kv_counter_size( + num_heads, max_batch_size, sm_count) + counters = self._dflash_trtllm_gen_counters + counters_need_allocation = (counters is None + or counters.device != device + or counters.numel() * + counters.element_size() < counter_bytes) + if counters_need_allocation: + if is_capturing: + raise RuntimeError( + "The DFlash TRTLLM-Gen counter buffer must be allocated at " + "the required size before CUDA graph capture.") + self._dflash_trtllm_gen_counters = torch.zeros(counter_bytes, + dtype=torch.uint8, + device=device) + + append_batch_indices = self._dflash_batch_indices + block_offsets = self._dflash_block_offsets + static_indices_need_allocation = ( + append_batch_indices is None or block_offsets is None + or append_batch_indices.device != device + or block_offsets.device != device + or append_batch_indices.size(0) < max_batch_size + or append_batch_indices.size(1) != block_size + or block_offsets.numel() != block_size) + if static_indices_need_allocation: + if is_capturing: + raise RuntimeError( + "DFlash TRTLLM-Gen index buffers must be allocated at the " + "required size before CUDA graph capture.") + self._dflash_batch_indices = (torch.arange( + max_batch_size, dtype=torch.int32, + device=device).view(-1, 1).expand(-1, block_size).contiguous()) + self._dflash_block_offsets = torch.arange(block_size, + dtype=torch.int32, + device=device) + + def dflash_forward( + self, + noise_embedding: torch.Tensor, + query_positions: torch.Tensor, + num_ctx_per_req: torch.Tensor, + ctx_k_cache: torch.Tensor, + ctx_v_cache: torch.Tensor, + ctx_cache_batch_idx: torch.Tensor, + ctx_kv_cache: Optional[torch.Tensor] = None, + ctx_page_table: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """DFlash draft forward with cross-attention over a pooled K/V buffer. + + All shapes are fixed so the forward is CUDA-graph compatible. + + Args: + noise_embedding: [B, block_size, hidden_size] + query_positions: [B, block_size] + num_ctx_per_req: [B] — per-batch context length in the pool + ctx_k_cache: [pool_batch, L, max_ctx+block_size, nkv, hd] + ctx_v_cache: [pool_batch, L, max_ctx+block_size, nkv, hd] + ctx_cache_batch_idx: [B] — slot index into the pool per batch entry + Returns: + [B * block_size, hidden_size] + """ + if self.dflash_attention_backend == 'TRTLLM': + if ctx_kv_cache is None or ctx_page_table is None: + raise RuntimeError( + "DFlash TRTLLM-Gen requires a paged context cache and page table." + ) + trtllm_gen_ops = self._dflash_trtllm_gen_ops + elif self.dflash_attention_backend == 'VANILLA': + flash_attention = self._dflash_flash_attention + else: + raise ValueError( + "DFlash attention backend must be VANILLA or TRTLLM, got " + f"{self.dflash_attention_backend!r}.") + + if self._fused_kv_weight is None: + self._build_fused_kv_buffers() + + layer0 = self.model.layers[0] + attn0 = layer0.self_attn + q_size = attn0.q_size + kv_size = attn0.kv_size + head_dim = attn0.head_dim + # Uniformity across layers is asserted in _build_fused_kv_buffers (above). + num_heads_per_rank = attn0.num_heads + num_kv_heads_per_rank = attn0.num_key_value_heads + gqa_group_size = num_heads_per_rank // num_kv_heads_per_rank + + has_qk_norm = self._has_qk_norm + is_bf16 = noise_embedding.dtype == torch.bfloat16 + use_fused_qk_norm_rope = self._use_fused_qk_norm_rope and is_bf16 + use_fused_rope = (_flashinfer_rope is not None and has_qk_norm + and is_bf16 and not use_fused_qk_norm_rope) + + B = noise_embedding.shape[0] + block_size = noise_embedding.shape[1] + + hidden_states = noise_embedding # [B, block_size, hidden] + + # Precompute RoPE cos/sin for the pure-PyTorch fallback path only. + # The fused flashinfer path reads self._get_cos_sin_cache() inline. + rope_dtype = hidden_states.dtype + if not use_fused_rope: + q_rope_cos, q_rope_sin = self._get_rope_cos_sin(query_positions, + dtype=rope_dtype) + _rope = RotaryEmbedding.apply_rotary_pos_emb + + # cache_seqlens (BEFORE append). flash_attn appends block_size + # k/v at cache_seqlens[i]..+block_size for batch i. + cache_seqlens_i32 = num_ctx_per_req[:B].to(torch.int32) + cache_batch_idx_i32 = ctx_cache_batch_idx.to(torch.int32) + + if self.dflash_attention_backend == 'TRTLLM': + max_batch_size = ctx_page_table.size(0) + self._prepare_dflash_trtllm_gen_buffers( + hidden_states.dtype, + hidden_states.device, + max_batch_size, + block_size, + num_heads_per_rank, + num_kv_heads_per_rank, + head_dim, + ) + block_tables = ctx_page_table.index_select( + 0, cache_batch_idx_i32.long()) + pages_per_slot = block_tables.size(1) + page_size = ctx_kv_cache.size(-2) + kv_indices = block_tables.flatten() + kv_indptr = torch.arange( + 0, + (B + 1) * pages_per_slot, + pages_per_slot, + dtype=torch.int32, + device=hidden_states.device, + ) + seq_lens_after = cache_seqlens_i32 + block_size + kv_last_page_len = ((seq_lens_after - 1) % page_size) + 1 + batch_indices = self._dflash_batch_indices + append_batch_indices = batch_indices[:B].reshape(-1) + append_positions = ( + cache_seqlens_i32.view(-1, 1) + + self._dflash_block_offsets).reshape(-1).contiguous() + + # Flatten query positions once for the fused QK-norm-RoPE kernel. + query_positions_flat_i32 = query_positions.reshape(-1).to(torch.int32) + + residual = None + + for layer_idx, layer in enumerate(self.model.layers): + attn_mod = layer.self_attn + + # Apply input_layernorm (flatten to 2D for norm, reshape back) + hs_flat = hidden_states.reshape(-1, hidden_states.shape[-1]) + if residual is None: + residual = hidden_states.clone() + hs_normed_flat = layer.input_layernorm(hs_flat) + else: + res_flat = residual.reshape(-1, residual.shape[-1]) + hs_normed_flat, res_flat = layer.input_layernorm( + hs_flat, res_flat) + residual = res_flat.reshape(B, block_size, -1) + + # QKV projection on normed query tokens (2D) + qkv_query = attn_mod.qkv_proj(hs_normed_flat) # [B*blk, qkv_size] + + if use_fused_qk_norm_rope: + # One kernel does q_norm + k_norm + RoPE in-place on qkv. + # Only safe when the drafter's rope params don't use YaRN / + # long-rope / partial-rotary — otherwise fall back to the + # shared-cache path below. + attn_mod.apply_qk_norm_rope(qkv_query, query_positions_flat_i32) + q_all_2d = qkv_query[:, :q_size] + k_noise_2d = qkv_query[:, q_size:q_size + kv_size] + v_noise_2d = qkv_query[:, q_size + kv_size:] + Q_bshd = q_all_2d.reshape(B, block_size, num_heads_per_rank, + head_dim) + k_noise_bshd = k_noise_2d.reshape(B, block_size, + num_kv_heads_per_rank, + head_dim) + v_noise_bshd = v_noise_2d.reshape(B, block_size, + num_kv_heads_per_rank, + head_dim) + elif use_fused_rope: + # Per-head RMSNorm on q/k (returns new contiguous tensors), + # then flashinfer in-place RoPE sharing the same cos/sin cache + # as precompute_context_kv. + q = attn_mod.q_norm(qkv_query[:, :q_size].reshape( + -1, head_dim)).view(-1, q_size) + k = attn_mod.k_norm(qkv_query[:, + q_size:q_size + kv_size].reshape( + -1, + head_dim)).view(-1, kv_size) + _flashinfer_rope( + query_positions_flat_i32, + q, + k, + head_dim, + self._get_cos_sin_cache(), + self._is_neox, + ) + Q_bshd = q.view(B, block_size, num_heads_per_rank, head_dim) + k_noise_bshd = k.view(B, block_size, num_kv_heads_per_rank, + head_dim) + v_noise_bshd = qkv_query[:, q_size + kv_size:].reshape( + B, block_size, num_kv_heads_per_rank, head_dim) + else: + qkv_query_3d = qkv_query.reshape(B, block_size, -1) + q_all = qkv_query_3d[..., :q_size] + k_noise_all = qkv_query_3d[..., q_size:q_size + kv_size] + v_noise_all = qkv_query_3d[..., q_size + kv_size:] + if has_qk_norm: + q_for_rope = attn_mod.q_norm(q_all.reshape( + -1, head_dim)).reshape(B, block_size, q_size) + k_noise_for_rope = attn_mod.k_norm( + k_noise_all.reshape(-1, head_dim)).reshape( + B, block_size, kv_size) + else: + q_for_rope = q_all + k_noise_for_rope = k_noise_all + Q = _rope(q_for_rope.reshape(B, block_size, num_heads_per_rank, + head_dim).transpose(1, 2), + q_rope_cos, + q_rope_sin, + unsqueeze_dim=1, + is_neox=self._is_neox) + k_noise_rope = _rope(k_noise_for_rope.reshape( + B, block_size, num_kv_heads_per_rank, + head_dim).transpose(1, 2), + q_rope_cos, + q_rope_sin, + unsqueeze_dim=1, + is_neox=self._is_neox) + Q_bshd = Q.transpose(1, 2) + k_noise_bshd = k_noise_rope.transpose(1, 2) + v_noise_bshd = v_noise_all.reshape(B, block_size, + num_kv_heads_per_rank, + head_dim) + + # Per-layer view into the pooled ctx cache. + causal, window_size = self._get_attention_mask_args(layer_idx) + dspark_window = (self._dspark_layer_windows[layer_idx] if layer_idx + < len(self._dspark_layer_windows) else (-1, -1)) + if dspark_window != (-1, -1): + window_size = dspark_window + if self.dflash_attention_backend == 'TRTLLM': + layer_cache = ctx_kv_cache[layer_idx] + trtllm_gen_ops.append_paged_kv_cache( + append_key=k_noise_bshd.reshape(-1, num_kv_heads_per_rank, + head_dim).contiguous(), + append_value=v_noise_bshd.reshape(-1, num_kv_heads_per_rank, + head_dim).contiguous(), + batch_indices=append_batch_indices, + positions=append_positions, + paged_kv_cache=layer_cache, + kv_indices=kv_indices, + kv_indptr=kv_indptr, + kv_last_page_len=kv_last_page_len, + kv_layout="HND", + ) + out = torch.empty_like(Q_bshd) + q_flat = Q_bshd.reshape(-1, num_heads_per_rank, head_dim) + out_flat = out.reshape(-1, num_heads_per_rank, head_dim) + window_left = window_size[0] + if causal: + trtllm_gen_ops.batch_decode_with_kv_cache( + query=q_flat, + kv_cache=(layer_cache[:, 0], layer_cache[:, 1]), + workspace_buffer=self._dflash_trtllm_gen_workspace, + block_tables=block_tables, + seq_lens=seq_lens_after, + max_seq_len=pages_per_slot * page_size, + bmm1_scale=head_dim**-0.5, + bmm2_scale=1.0, + window_left=window_left, + out=out_flat, + sinks=None, + enable_pdl=False, + kv_layout="HND", + backend="trtllm-gen", + q_len_per_req=block_size, + max_q_len=None, + cum_seq_lens_q=None, + kv_cache_sf=None, + uses_shared_paged_kv_idx=True, + bmm1_scale_log2=None, + multi_ctas_kv_counter_buffer=self. + _dflash_trtllm_gen_counters, + ) + else: + cum_seq_lens_q = torch.arange( + 0, + (B + 1) * block_size, + block_size, + dtype=torch.int32, + device=hidden_states.device, + ) + cum_seq_lens_kv = torch.cat(( + torch.zeros(1, + dtype=torch.int32, + device=hidden_states.device), + seq_lens_after.cumsum(0, dtype=torch.int32), + )) + trtllm_gen_ops.batch_context_with_kv_cache( + query=q_flat, + kv_cache=(layer_cache[:, 0], layer_cache[:, 1]), + workspace_buffer=self._dflash_trtllm_gen_workspace, + block_tables=block_tables, + seq_lens=seq_lens_after, + max_q_len=block_size, + max_kv_len=pages_per_slot * page_size, + bmm1_scale=head_dim**-0.5, + bmm2_scale=1.0, + batch_size=B, + cum_seq_lens_q=cum_seq_lens_q, + cum_seq_lens_kv=cum_seq_lens_kv, + window_left=window_left, + out=out_flat, + sinks=None, + enable_pdl=False, + kv_layout="HND", + kv_cache_sf=None, + uses_shared_paged_kv_idx=True, + causal=False, + multi_ctas_kv_counter_buffer=self. + _dflash_trtllm_gen_counters, + ) + else: # VANILLA, validated before entering the layer loop. + layer_k_cache = ctx_k_cache[:, layer_idx] + layer_v_cache = ctx_v_cache[:, layer_idx] + + # Pack gqa_group_size query heads sharing a KV head into the + # row dimension: [B, blk, h_q, d] -> [B, group*blk, h_kv, d]. + # Each CTA owns a whole query-head group and streams KV head's context once + # instead of gqa_group_size CTAs each re-reading it. + # Exact only while every row of the block attends to the same + # key set, i.e. non-causal, unwindowed layers. Causal or + # windowed layers mask by row, so they stay unpacked. + pack_gqa = (gqa_group_size > 1 and not causal + and window_size == (-1, -1)) + if pack_gqa: + q_grouped = Q_bshd.reshape(B, block_size, + num_kv_heads_per_rank, + gqa_group_size, head_dim) + q_packed = q_grouped.permute(0, 3, 1, 2, 4) + q_in = q_packed.reshape(B, gqa_group_size * block_size, + num_kv_heads_per_rank, head_dim) + else: + q_in = Q_bshd + out = flash_attention( + q=q_in, + k_cache=layer_k_cache, + v_cache=layer_v_cache, + k=k_noise_bshd, + v=v_noise_bshd, + cache_seqlens=cache_seqlens_i32, + cache_batch_idx=cache_batch_idx_i32, + causal=causal, + window_size=window_size, + ) + if pack_gqa: + # Undo the packing: [B, group*blk, h_kv, d] -> [B, blk, h_q, d]. + out = out.view(B, gqa_group_size, block_size, + num_kv_heads_per_rank, + head_dim).permute(0, 2, 3, 1, 4) + + attn_output = out.reshape(B * block_size, q_size) + + # Per-drafter post-attention gate (no-op for generic DFlash; Laguna + # applies per-head softplus g_proj gating). gate input is the + # input_layernorm output (the attention input). + attn_output = self._post_attention_gate(attn_output, hs_normed_flat, + attn_mod, + num_heads_per_rank, + head_dim) + + # o_proj (flat 2D, handles all-reduce internally) + hidden_out = attn_mod.o_proj(attn_output) + + # Post-attention layernorm + MLP (flat 2D) + 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(B, block_size, -1) + residual = res_flat.reshape(B, block_size, -1) + + # Final norm + hidden_states_out, _ = self.model.norm( + hidden_states.reshape(-1, hidden_states.shape[-1]), + residual.reshape(-1, residual.shape[-1])) + return hidden_states_out + + def forward( + self, + attn_metadata, + input_ids: torch.LongTensor = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + return_context_logits: bool = False, + spec_metadata=None, + hidden_states: torch.Tensor | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Run the draft model and return (hidden_states, hidden_states) for the + speculative-decoding contract.""" + hidden_states_out = self.model( + input_ids=input_ids, + attn_metadata=attn_metadata, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + **kwargs, + ) + + return hidden_states_out, hidden_states_out + + +class DFlashLagunaForCausalLM(DFlashForCausalLM): + """Laguna DFlash drafter. + + The generic block decode lives in DFlashForCausalLM; this subclass supplies + the Laguna draft-layer specifics: per-head g_proj softplus gating and the + per-aux fc_norm applied to captured target features before fc. + """ + + @staticmethod + def _normalize_config(config: PretrainedConfig) -> None: + """Fill TRT-LLM Laguna defaults missing from dense DFlash drafts.""" + if getattr(config, "num_experts", None) is None: + config.num_experts = 0 + if getattr(config, "mlp_layer_types", None) is None: + config.mlp_layer_types = ["dense"] * config.num_hidden_layers + if getattr(config, "block_size", None) is None: + dflash_config = getattr(config, "dflash_config", {}) + if isinstance(dflash_config, dict): + config.block_size = dflash_config.get("block_size", None) + + def __init__(self, + draft_config, + *, + dflash_attention_backend: str = 'VANILLA'): + """Pin the Laguna draft-layer class and enable Laguna-specific behaviors + (context input_layernorm, causal sliding blocks); reject non-per-head + gating.""" + # The checkpoint labels itself with the vLLM name (model_type "llama"); + # remap to the Laguna architecture so TRT-LLM builds the Laguna layers. + draft_config.pretrained_config.architectures = ["LagunaForCausalLM"] + self._normalize_config(draft_config.pretrained_config) + super().__init__( + draft_config, + dflash_attention_backend=dflash_attention_backend, + ) + self._context_input_layernorm = True + self._sliding_layers_causal = True + gating = getattr(self.config, 'gating', True) + if gating not in (True, 'per-head'): + raise NotImplementedError( + f"Laguna DFlash drafter supports per-head gating only, " + f"got gating={gating!r}") + + def load_weights(self, weights, weight_mapper=None, **kwargs): + """Build the per-aux ``fc_norm`` from the drafter's ``aux_hidden_norms.*`` + weights, then defer the remaining weights to the base loader.""" + aux_keys = sorted( + (k for k in weights if k.startswith('aux_hidden_norms.')), + key=lambda k: int(k.split('.')[1])) + if not aux_keys: + raise ValueError( + "Laguna DFlash checkpoint is missing aux_hidden_norms.* weights" + ) + weights = dict(weights) + eps = getattr(self.config, 'rms_norm_eps', 1e-6) + norms = [] + for k in aux_keys: + w = weights.pop(k) + norm = nn.RMSNorm(w.shape[0], + eps=eps, + device='cuda', + elementwise_affine=True, + dtype=w.dtype) + norm.weight.data.copy_(w) + norms.append(norm) + self.fc_norm = nn.ModuleList(norms) + super().load_weights(weights, weight_mapper=weight_mapper, **kwargs) + + def project_target_hidden(self, hidden_states): + """Project captured target features to the draft width: apply the per-aux + ``fc_norm`` to each hidden chunk, then ``fc`` + ``hidden_norm``.""" + hidden_states = hidden_states.to(self.fc.weight.dtype) + fc_norm = getattr(self, 'fc_norm', None) + if fc_norm is not None: + chunks = hidden_states.chunk(len(fc_norm), dim=-1) + hidden_states = torch.cat( + [norm(chunk) for norm, chunk in zip(fc_norm, chunks)], dim=-1) + return self.hidden_norm(self.fc(hidden_states)) + + def _post_attention_gate(self, attn_output, gate_input, attn_mod, num_heads, + head_dim): + """Apply Laguna's per-head softplus output gate (``g_proj``) to the + attention output; a no-op when the layer has no ``g_proj``.""" + g_proj = getattr(attn_mod, 'g_proj', None) + if g_proj is None: + return attn_output + gate = F.softplus(g_proj(gate_input).float()).to(attn_output.dtype) + return (attn_output.unflatten(-1, (num_heads, head_dim)) * + gate.unsqueeze(-1)).flatten(-2) + + +@register_draft_model(SpeculativeDecodingMode.DFLASH) +def _build_dflash_draft(model_config, draft_config, lm_head, model): + """Build the DFlash drafter. + + Selects the Laguna variant by detecting its architecture in the draft + checkpoint's own config. + """ + draft_arches = getattr(draft_config.pretrained_config, "architectures", + None) or [] + dflash_attention_backend = model_config.spec_config.attention_backend + if any("Laguna" in arch for arch in draft_arches): + return DFlashLagunaForCausalLM( + draft_config, + dflash_attention_backend=dflash_attention_backend, + ) + return DFlashForCausalLM( + draft_config, + dflash_attention_backend=dflash_attention_backend, + ) diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 94247ced2546..eaf6605279e4 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -6,13 +6,12 @@ from typing import Dict, Generic, List, Optional, Tuple import torch -import torch.nn.functional as F from torch import nn from transformers import LlamaConfig, PretrainedConfig from tensorrt_llm.logger import logger -from ...functional import PositionEmbeddingType, RotaryScalingType +from ...functional import PositionEmbeddingType from ..attention_backend import AttentionMetadata from ..attention_backend.interface import PositionalEmbeddingParams, RopeParams from ..model_config import ModelConfig, TConfig @@ -25,20 +24,9 @@ WeightsLoadingConfig) from ..modules.mla import MLA from ..modules.rms_norm import RMSNorm -from ..modules.rotary_embedding import RotaryEmbedding - -try: - from ..custom_ops import \ - flashinfer_apply_rope_with_cos_sin_cache_inplace as _flashinfer_rope -except ImportError: - _flashinfer_rope = None -from ..pyexecutor.config_utils import (_is_sliding_attention_layer, - get_layer_attention_window) from ..pyexecutor.guided_decoder import CapturableGuidedDecoder from ..speculative import (SpecMetadata, get_spec_worker, should_use_separate_draft_kv_cache) -from ..speculative.dflash_attention import (get_dflash_flash_attention, - get_dflash_trtllm_gen_ops) from ..speculative.interface import SpeculativeDecodingMode from ..utils import AuxStreamType from .checkpoints.base_weight_mapper import BaseWeightMapper @@ -854,1359 +842,6 @@ def forward( return hidden_states_out, hidden_states_out -def dspark_layer_window_size(use_swa: bool, swa_window: int, layer_types, - layer_idx: int) -> tuple[int, int]: - """flash-attn ``window_size`` for one draft layer of the block decode. - - DSpark drafters (deepseek-ai/DeepSpec) run the draft block through HF - attention with ``sliding_window`` set on 'sliding_attention' layers and - is_causal=False. HF's flash path - (transformers/modeling_flash_attention_utils.py) translates that to - ``window_size = (sliding_window - 1, sliding_window - 1)``, i.e. each - query attends keys within ``swa_window - 1`` KV-index distance on both - sides. In the DFlash pool layout KV index == token position, so this - limits draft queries to the most recent ``swa_window`` context tokens - plus the (nearby) draft block. Full-attention layers and non-dspark - drafters keep flash-attn's default ``(-1, -1)`` (no window). - """ - if not use_swa: - return (-1, -1) - if layer_types is not None and layer_idx < len(layer_types) and \ - layer_types[layer_idx] != 'sliding_attention': - return (-1, -1) - return (swa_window - 1, swa_window - 1) - - -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. - - Reference: DeepSpec ``VanillaMarkov`` (deepspec/modeling/dspark/ - markov_head.py): ``bias = markov_w2(markov_w1(prev_token))`` where - markov_w1 is nn.Embedding(vocab, rank) and markov_w2 is - nn.Linear(rank, vocab, bias=False). With both weights stored - [vocab, rank] this is ``markov_w1[prev] @ markov_w2.T``. - - Args: - prev_tokens: [B] long, previous token per request (draft vocab). - markov_w1: [vocab, rank]. - markov_w2: [vocab_or_shard, rank] (rows may be a TP vocab shard). - Returns: - [B, vocab_or_shard] bias in the markov weights' dtype. - """ - return F.linear(F.embedding(prev_tokens, markov_w1), markov_w2) - - -def dspark_markov_chain_logits( - base_logits: torch.Tensor, - first_prev_tokens: torch.Tensor, - markov_w1: torch.Tensor, - markov_w2: torch.Tensor, - argmax_fn=None, -) -> torch.Tensor: - """Apply the vanilla Markov intra-block bias across a drafted block. - - Reference: DeepSpec ``VanillaMarkov.sample_block_tokens`` at - temperature 0: for step i, ``logits_i += bias(prev_i)`` with - ``prev_0`` = the anchor token (last accepted token, block slot 0) and - ``prev_{i>0}`` = the greedy token from step i-1's *biased* logits. - llama.cpp PR #25173 implements the same greedy chain. - - Args: - base_logits: [B, K, vocab_or_shard] shared-lm_head logits. - first_prev_tokens: [B] long, anchor token ids (draft vocab). - markov_w1 / markov_w2: see :func:`dspark_markov_step_bias`. - argmax_fn: callable([B, vocab_or_shard]) -> [B] token ids in the - full draft vocab; defaults to plain argmax. Workers pass a - TP-aware argmax when the draft logits are vocab-sharded. - Returns: - [B, K, vocab_or_shard] biased logits. Greedy per-position argmax of - the result reproduces the reference sampled chain exactly. - """ - K = base_logits.shape[1] - if K == 0: - return base_logits - prev = first_prev_tokens.long() - steps = [] - for i in range(K): - bias = dspark_markov_step_bias(prev, markov_w1, markov_w2) - step_logits = base_logits[:, i] + bias.to(base_logits.dtype) - steps.append(step_logits) - if argmax_fn is not None: - prev = argmax_fn(step_logits).long() - else: - prev = torch.argmax(step_logits, dim=-1) - return torch.stack(steps, dim=1) - - -class DFlashForCausalLM(nn.Module): - """Draft model wrapper for DFlash speculative decoding. - - DFlash uses cross-attention where Q comes from noise/query tokens and K/V - come from the concatenation of target hidden states and noise hidden states. - The target_hidden stays CONSTANT across all layers (no input_layernorm applied). - - Reference: https://arxiv.org/pdf/2602.06036 - """ - - 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 - custom DFlash architecture label).""" - super().__init__() - - pretrained_cfg = draft_config.pretrained_config - try: - DraftModelClass, _ = get_model_architecture(pretrained_cfg) - except RuntimeError: - model_type = pretrained_cfg.model_type - arch_name = "".join(w.capitalize() - for w in model_type.split("_")) + "ForCausalLM" - logger.info( - f"DFlash: architecture {pretrained_cfg.architectures} not found, " - f"falling back to {arch_name} based on model_type={model_type}") - original_archs = pretrained_cfg.architectures - try: - pretrained_cfg.architectures = [arch_name] - DraftModelClass, _ = get_model_architecture(pretrained_cfg) - finally: - pretrained_cfg.architectures = original_archs - - # Remove spec_config to prevent recursive spec-dec initialization - draft_config_no_spec = replace(draft_config, - spec_config=None, - lm_head_gather_output=False) - - # Weights will be loaded later by ModelLoader.load_draft_weights() - self.draft_model_full = DraftModelClass(draft_config_no_spec) - self.model = self.draft_model_full.model - self.lm_head = self.draft_model_full.lm_head - - # Required by weight mappers - self.model_config = draft_config_no_spec - self.config = draft_config_no_spec.pretrained_config - - # Get mask_token_id from dflash_config - pretrained_config = draft_config.pretrained_config - dflash_config = getattr(pretrained_config, 'dflash_config', {}) - self.mask_token_id = dflash_config.get( - 'mask_token_id', - getattr(pretrained_config, 'mask_token_id', - pretrained_config.vocab_size)) - - self.target_layer_ids = dflash_config.get('target_layer_ids', None) - self.block_size = dflash_config.get( - 'block_size', getattr(pretrained_config, 'block_size', None)) - self.dflash_attention_backend = dflash_attention_backend - if 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() - else: - raise ValueError( - "DFlash attention backend must be VANILLA or TRTLLM, 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) - self.register_buffer("_dflash_block_offsets", None, persistent=False) - self._dflash_trtllm_gen_device = None - self._dflash_trtllm_gen_sm_count = None - logger.info( - f"DFlash draft model initialized with mask_token_id: {self.mask_token_id}, " - f"target_layer_ids: {self.target_layer_ids}, block_size: {self.block_size}, " - f"attention_backend: {self.dflash_attention_backend}") - - # DSpark drafters (DFlash + low-rank Markov head + confidence head, - # arXiv 2607.05147; reference: deepseek-ai/DeepSpec). The weights- - # independent drafter-forward semantics ARE implemented here: - # - vanilla Markov intra-block logit bias (applied by DFlashWorker - # through apply_markov_chain_logits), - # - sliding-window attention on 'sliding_attention' draft layers - # during the block decode (use_swa / swa_window_size), - # - the shift_label output convention (hidden state at block slot j - # predicts draft token j+1; slot 0 holds the anchor token). - # Confidence-scheduled verification is NOT implemented yet: the - # confidence_proj weights are loaded (for the follow-up MR) but never - # used, and drafting always proposes the full K tokens. - self._dspark_shift_label = bool(dflash_config.get('shift_label', False)) - self._dspark_use_swa = bool(dflash_config.get('use_swa', False)) - self._dspark_swa_window = int( - dflash_config.get('swa_window_size', 0) or 0) - self._dspark_markov_rank = int(dflash_config.get('markov_rank', 0) or 0) - self._dspark_markov_head_type = str( - dflash_config.get('markov_head_type', 'vanilla') - or 'vanilla').lower() - self._dspark_use_confidence_head = bool( - dflash_config.get('use_confidence_head', False)) - # Plain None placeholders rather than nn.Parameter/buffer: most - # DFlash checkpoints don't ship these heads, and their shapes - # ([vocab, rank]) are checkpoint-dependent, so nothing is - # pre-allocated. load_weights() fills them in only when the - # checkpoint ships them; consumers treat None as "head absent". - self.markov_w1 = None # [vocab, rank] (nn.Embedding weight layout) - self.markov_w2 = None # [vocab, rank] (nn.Linear(rank->vocab) weight) - self.confidence_proj_weight = None # loaded, unused (follow-up MR) - self.confidence_proj_bias = None - - if self._dspark_markov_rank > 0 and \ - self._dspark_markov_head_type != 'vanilla': - raise ValueError( - f"DFlash dspark drafter declares markov_head_type=" - f"'{self._dspark_markov_head_type}'; only 'vanilla' is " - "supported (gated/rnn heads need per-step hidden features).") - if self._dspark_use_swa and self._dspark_swa_window < 1: - raise ValueError( - "DFlash dspark drafter sets use_swa but swa_window_size=" - f"{dflash_config.get('swa_window_size')} is invalid.") - # causal=true is only invalid under the dspark convention. Legacy - # DFlash drafter configs (e.g. Laguna) also carry a causal field; - # their causality is handled by the legacy decode path - # (_sliding_layers_causal), so don't reject them here. - is_dspark = (str(dflash_config.get('projector_type', '') - or '').lower() == 'dspark' or self._dspark_shift_label - or self._dspark_use_swa or self._dspark_markov_rank > 0 - or self._dspark_use_confidence_head) - if is_dspark and dflash_config.get('causal'): - raise ValueError( - "DFlash dspark drafter sets causal=true; the block decode " - "only supports the non-causal dspark convention.") - # Per-layer flash-attn window for the block decode, resolved once. - num_draft_layers = getattr(pretrained_config, 'num_hidden_layers', 0) - layer_types = getattr(pretrained_config, 'layer_types', None) - self._dspark_layer_windows = [ - dspark_layer_window_size(self._dspark_use_swa, - self._dspark_swa_window, layer_types, i) - for i in range(num_draft_layers) - ] - if self._dspark_use_confidence_head: - logger.warning( - "DFlash dspark drafter declares use_confidence_head; " - "confidence-scheduled verification is not implemented yet " - "(confidence_proj weights are loaded but unused, drafting " - "always proposes the full K tokens).") - - self.logits_processor = None # Set by caller after construction - - # RoPE - lazily initialized from draft model's attention module - self._rope_initialized = False - self._rotary_cos_sin = None - self._is_neox = True - - self._cos_sin_cache_fp32 = None - self._rope_dummy_q = None - - # Lazy-built after weights load (see _build_fused_kv_buffers). - self._fused_kv_weight = None - self._fused_kv_bias = None - self._k_norm_stacked = None - self._k_norm_eps = None - self._num_attn_layers = 0 - self._num_heads = 0 - self._head_dim = 0 - self._num_kv_heads = 0 - self._has_qk_norm = False - self._use_fused_qk_norm_rope = False - # Laguna-specific draft-layer behaviors, disabled by default so generic - # DFlash drafters keep the original contract (no context input_layernorm, - # non-causal block attention). Subclasses opt in. - self._context_input_layernorm = False - self._sliding_layers_causal = False - self._warn_inferred_attention_windows() - - @staticmethod - def _rope_signature(attn): - """Return the effective RoPE configuration used by an attention layer.""" - if attn.rotary_emb is not None: - return ( - attn.rotary_emb.rope_params, - attn.rotary_emb.head_dim, - attn.rotary_emb.is_neox, - ) - if attn.pos_embd_params is not None: - return ( - attn.pos_embd_params.rope, - attn.head_dim, - attn.pos_embd_params.is_neox, - ) - return None - - def _validate_uniform_rope(self): - """Check that all draft layers can safely share one RoPE cache.""" - if len(self.model.layers) == 0: - raise ValueError("DFlash requires at least one draft model layer.") - - signatures = [ - self._rope_signature(layer.self_attn) for layer in self.model.layers - ] - - mismatched_layers = [ - layer_idx - for layer_idx, signature in enumerate(signatures[1:], start=1) - if signature != signatures[0] - ] - if mismatched_layers: - layer_types = getattr(self.config, 'layer_types', None) - raise ValueError( - "DFlash shares one RoPE cache across draft layers, but layers " - f"{mismatched_layers} have a different effective RoPE " - f"configuration from layer 0. layer_types={layer_types}.") - - def _init_rope(self): - """Initialize RoPE from the draft model's attention configuration. - - Reuses the existing RotaryEmbedding infrastructure which correctly - handles all RoPE variants (standard, YaRN, scaled, etc.). - """ - # The flattened context-KV path shares layer 0's RoPE cache. - self._validate_uniform_rope() - attn0 = self.model.layers[0].self_attn - - if attn0.rotary_emb is not None: - self._rotary_cos_sin = attn0.rotary_emb.rotary_cos_sin - self._is_neox = attn0.rotary_emb.is_neox - elif attn0.pos_embd_params is not None: - rope_emb = RotaryEmbedding( - attn0.pos_embd_params.rope, - head_dim=attn0.head_dim, - is_neox=attn0.pos_embd_params.is_neox, - ) - self._rotary_cos_sin = rope_emb.rotary_cos_sin - self._is_neox = rope_emb.is_neox - else: - # Fallback: basic NeoX-style RoPE - config = self.config - head_dim = getattr(config, 'head_dim', - config.hidden_size // config.num_attention_heads) - rope_theta = getattr(config, 'rope_theta', 1000000.0) - max_pos = getattr(config, 'max_position_embeddings', 32768) - - inv_freq = 1.0 / (rope_theta**(torch.arange( - 0, head_dim, 2, dtype=torch.float32, device='cuda') / head_dim)) - positions = torch.arange(max_pos, - dtype=torch.float32, - device='cuda') - freqs = torch.outer(positions, inv_freq) - rope_cos = freqs.cos().to(config.torch_dtype) - rope_sin = freqs.sin().to(config.torch_dtype) - # [max_pos, 2, rot_dim//2] to match RotaryEmbedding format - self._rotary_cos_sin = torch.stack([rope_cos, rope_sin], dim=1) - self._is_neox = True - - self._rope_initialized = True - - def project_target_hidden(self, - hidden_states: torch.Tensor) -> torch.Tensor: - """Project captured target hidden states into the draft hidden space. - - Generic DFlash: fc then hidden_norm. Subclasses (e.g. Laguna) may - normalize the per-aux features first by overriding this method. - """ - hidden_states = hidden_states.to(self.fc.weight.dtype) - return self.hidden_norm(self.fc(hidden_states)) - - @property - def has_markov_head(self) -> bool: - return self._dspark_markov_rank > 0 and self.markov_w1 is not None - - def apply_markov_chain_logits( - self, - base_logits: torch.Tensor, - first_prev_tokens: torch.Tensor, - argmax_fn=None, - vocab_slice: slice | None = None) -> torch.Tensor: - """Apply the dspark vanilla-Markov intra-block bias to block logits. - - No-op (returns ``base_logits`` unchanged) for non-dspark drafters. - See :func:`dspark_markov_chain_logits` for the semantics; when - ``base_logits`` is a TP vocab shard, the caller must pass this - rank's ``vocab_slice`` (to shard the markov_w2 rows identically) - and an ``argmax_fn`` returning full-vocab token ids — DFlashWorker - handles both. - """ - if not self.has_markov_head: - return base_logits - markov_w2 = self.markov_w2 if vocab_slice is None else \ - self.markov_w2[vocab_slice] - return dspark_markov_chain_logits(base_logits, - first_prev_tokens, - self.markov_w1, - markov_w2, - argmax_fn=argmax_fn) - - def _post_attention_gate(self, attn_output, gate_input, attn_mod, num_heads, - head_dim): - """Hook applied to the block-attention output before o_proj. - - No-op for generic DFlash; overridden by drafters that gate (e.g. Laguna). - """ - return attn_output - - def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): - """Load weights into the DFlash draft model. - - DFlash checkpoints differ from standard HF format: - - Layer weights lack the 'model.' prefix (e.g., 'layers.0...' not 'model.layers.0...') - - Extra DFlash-specific weights: 'fc.weight', 'hidden_norm.weight' - - Missing embed_tokens and lm_head (shared with target model) - """ - # Laguna DFlash checkpoints may ship a fused self_attn.qkv_proj; the draft - # loader expects split q/k/v (a fused key is silently dropped otherwise). - if any(k.endswith('self_attn.qkv_proj.weight') for k in weights): - for attr in ('num_attention_heads_per_layer', - 'num_key_value_heads_per_layer'): - per_layer = getattr(self.config, attr, None) - if per_layer is not None and len(set(per_layer)) > 1: - raise ValueError( - "DFlash load_weights() splits the fused qkv_proj using " - "the global head count, but the drafter has heterogeneous " - f"{attr} {sorted(set(per_layer))}; per-layer qkv splitting " - "is required for this checkpoint.") - head_dim = getattr( - self.config, 'head_dim', - self.config.hidden_size // self.config.num_attention_heads) - num_kv_heads = getattr(self.config, 'num_key_value_heads', - self.config.num_attention_heads) - q = self.config.num_attention_heads * head_dim - kv = num_kv_heads * head_dim - split = {} - for k, v in weights.items(): - if k.endswith('self_attn.qkv_proj.weight'): - b = k[:-len('qkv_proj.weight')] - split[b + 'q_proj.weight'] = v[:q] - split[b + 'k_proj.weight'] = v[q:q + kv] - split[b + 'v_proj.weight'] = v[q + kv:] - else: - split[k] = v - weights = split - - # DSpark head weights: keep them out of the backbone remap (they'd - # get a 'model.' prefix and be dropped by allow_partial_loading). - # markov_w1/markov_w2 drive the intra-block logit bias; the - # confidence_proj weights are loaded for the confidence-scheduling - # follow-up MR but are not used yet. - dspark_keys = ('markov_w1.weight', 'markov_w2.weight', - 'confidence_proj.weight', 'confidence_proj.bias') - dspark_weights = {k: weights[k] for k in dspark_keys if k in weights} - if dspark_weights: - weights = { - k: v - for k, v in weights.items() if k not in dspark_weights - } - if self._dspark_markov_rank > 0: - vocab = self.config.vocab_size - rank = self._dspark_markov_rank - for k in ('markov_w1.weight', 'markov_w2.weight'): - if k not in dspark_weights: - raise ValueError( - f"DFlash dspark drafter declares markov_rank=" - f"{self._dspark_markov_rank} but the checkpoint is " - f"missing {k}.") - if tuple(dspark_weights[k].shape) != (vocab, rank): - raise ValueError( - f"DFlash dspark {k} has shape " - f"{tuple(dspark_weights[k].shape)}, expected " - f"[vocab, markov_rank] = ({vocab}, {rank}).") - self.markov_w1 = dspark_weights['markov_w1.weight'].to('cuda') - self.markov_w2 = dspark_weights['markov_w2.weight'].to('cuda') - if 'confidence_proj.weight' in dspark_weights: - 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') - - # Remap: add 'model.' prefix where needed, and extract DFlash-specific weights - remapped = {} - for key, value in weights.items(): - if key in ('fc.weight', 'hidden_norm.weight'): - # DFlash-specific projection weights - store directly - remapped[key] = value - elif key == 'norm.weight': - remapped['model.norm.weight'] = value - elif not key.startswith('model.'): - remapped[f'model.{key}'] = value - else: - remapped[key] = value - - # Load DFlash-specific weights directly - if 'fc.weight' in remapped: - self.fc = nn.Linear(remapped['fc.weight'].shape[1], - remapped['fc.weight'].shape[0], - bias=False, - device='cuda', - dtype=remapped['fc.weight'].dtype) - self.fc.weight.data.copy_(remapped['fc.weight']) - del remapped['fc.weight'] - - if 'hidden_norm.weight' in remapped: - rms_norm_eps = getattr(self.config, 'rms_norm_eps', 1e-6) - self.hidden_norm = nn.RMSNorm( - remapped['hidden_norm.weight'].shape[0], - eps=rms_norm_eps, - device='cuda', - elementwise_affine=True, - dtype=remapped['hidden_norm.weight'].dtype) - 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. - self.draft_model_full.load_weights(weights=remapped, - weight_mapper=weight_mapper, - allow_partial_loading=True) - - def load_weights_from_target_model(self, - target_model: torch.nn.Module) -> None: - """Share embed_tokens and lm_head from the target model.""" - self.draft_model_full.model.embed_tokens = target_model.model.embed_tokens - self.draft_model_full.lm_head = target_model.lm_head - self.lm_head = target_model.lm_head - - def precompute_context_kv( - self, - projected_hidden: torch.Tensor, - positions: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Post-norm / post-RoPE K and V for ALL drafter layers in one fused GEMM. - - Args: - projected_hidden: [N, hidden_size], already fc + hidden_norm'd. - positions: [N] int32/64, RoPE positions for each entry. - Returns: - k: [N, L, nkv, hd] post k_norm and RoPE - v: [N, L, nkv, hd] post split only - """ - if self._fused_kv_weight is None: - self._build_fused_kv_buffers() - N = projected_hidden.shape[0] - L = self._num_attn_layers - nkv = self._num_kv_heads - hd = self._head_dim - weight_dtype = self._fused_kv_weight.dtype - if getattr(self, '_input_ln_eps', None) is not None: - ph = projected_hidden.float() - ph = ph * torch.rsqrt( - ph.pow(2).mean(-1, keepdim=True) + self._input_ln_eps) - projected_hidden = ph.to(weight_dtype) - elif projected_hidden.dtype != weight_dtype: - projected_hidden = projected_hidden.to(weight_dtype) - - kv_flat = F.linear(projected_hidden, self._fused_kv_weight, - self._fused_kv_bias) - # Per-layer layout [L0_K|L0_V|L1_K|L1_V|...] keeps K and V contiguous - # after the select() splits — no extra copy required. - kv = kv_flat.view(N, L, 2, nkv, hd) - k = kv[:, :, 0].contiguous() - v = kv[:, :, 1].contiguous() - - if self._k_norm_stacked is not None: - # Fuse L per-layer RMSNorms into one. k is [N, L, nkv, hd]; - # each layer has its own weight ([L, hd]) but shares eps. - k = F.rms_norm(k, (hd, ), eps=self._k_norm_eps) - k = k * self._k_norm_stacked.view(1, L, 1, hd) - - self._fused_rope_inplace(k.view(N * L, nkv * hd), positions, N, L) - return k, v - - def _get_cos_sin_cache(self) -> torch.Tensor: - """Return the flashinfer-style cos/sin cache for the drafter. - - Shape [max_positions, head_dim], fp32 — flashinfer's - apply_rope_with_cos_sin_cache_inplace requires fp32 regardless of - the query/key dtype. - """ - if self._cos_sin_cache_fp32 is not None: - return self._cos_sin_cache_fp32 - if not self._rope_initialized: - self._init_rope() - max_pos = self._rotary_cos_sin.shape[0] - self._cos_sin_cache_fp32 = self._rotary_cos_sin.view(max_pos, -1).to( - torch.float32).contiguous() - return self._cos_sin_cache_fp32 - - def _fused_rope_inplace( - self, - k_flat: torch.Tensor, - positions: torch.Tensor, - N: int, - L: int, - ) -> None: - """In-place fused RoPE over [N*L, nkv*hd] K values. - - Layout of k_flat: row (i*L + l) holds layer l of position i, so - positions must be repeat_interleaved by L to match. - """ - positions_int32 = positions.view(-1).to(torch.int32) - if L > 1: - positions_int32 = positions_int32.repeat_interleave(L) - - if _flashinfer_rope is not None: - # flashinfer requires a non-None query tensor; pass a single-head - # scratch so the extra rotate is negligible. - need_rows = k_flat.shape[0] - dummy_q = self._rope_dummy_q - if (dummy_q is None or dummy_q.dtype != k_flat.dtype - or dummy_q.shape[0] < need_rows): - dummy_q = k_flat.new_empty(need_rows, self._head_dim) - self._rope_dummy_q = dummy_q - _flashinfer_rope( - positions_int32, - dummy_q[:need_rows], - k_flat, - self._head_dim, - self._get_cos_sin_cache(), - self._is_neox, - ) - return - - # Pure-PyTorch fallback (older environments without flashinfer). - cos, sin = self._get_rope_cos_sin(positions_int32.view(1, -1), - dtype=k_flat.dtype) - k_roped = RotaryEmbedding.apply_rotary_pos_emb( - k_flat.view(k_flat.shape[0], -1, self._head_dim), - cos.squeeze(0), - sin.squeeze(0), - unsqueeze_dim=1, - is_neox=self._is_neox, - ) - k_flat.copy_(k_roped.view_as(k_flat)) - - def _build_fused_kv_buffers(self) -> None: - """Stack per-layer KV projection + k_norm weights for a single fused GEMM. - - Must run after weights are loaded. - """ - if self._fused_kv_weight is not None: - return - layers_attn = [layer.self_attn for layer in self.model.layers] - attn0 = layers_attn[0] - q_size = attn0.q_size - kv_size = attn0.kv_size - head_dim = attn0.head_dim - num_heads = attn0.num_heads - num_kv_heads = attn0.num_key_value_heads - # Head counts are read from layer 0 here and in dflash_forward; assert - # uniformity (the target uses per-layer heads, the drafter does not). - for a in layers_attn[1:]: - assert ( - a.q_size == q_size and a.kv_size == kv_size - and a.head_dim == head_dim and a.num_heads == num_heads - and a.num_key_value_heads == num_kv_heads), ( - "DFlash fused KV requires all drafter layers to share " - "q_size / kv_size / head_dim / num_heads / num_kv_heads.") - - has_k_norm = [hasattr(a, 'k_norm') for a in layers_attn] - assert all(has_k_norm) or not any(has_k_norm), ( - "DFlash fused KV requires either all or no drafter layers to have k_norm." - ) - - kv_weights = [ - a.qkv_proj.weight[q_size:q_size + 2 * kv_size] for a in layers_attn - ] - # Fold each drafter layer's input_layernorm weight into its KV projection - # so context K/V match the query path. vLLM laguna_dflash applies - # layer.input_layernorm to context states before KV; RMSNorm gives - # (x_hat * w) @ Wkv.T == x_hat @ (Wkv * w).T, and the shared 1/rms(x) is - # applied to projected_hidden in precompute_context_kv. - dlayers = self.model.layers - if self._context_input_layernorm and all( - hasattr(dl, 'input_layernorm') for dl in dlayers): - eps_set = { - getattr(dl.input_layernorm, 'variance_epsilon', - getattr(self.config, 'rms_norm_eps', 1e-6)) - for dl in dlayers - } - assert len(eps_set) == 1, ( - "DFlash fused context input_layernorm needs all drafter layers " - f"to share variance_epsilon; got {sorted(eps_set)}") - self._input_ln_eps = eps_set.pop() - folded = [] - for w, dl in zip(kv_weights, dlayers): - scale = dl.input_layernorm.weight.data - if getattr(dl.input_layernorm, 'use_gemma', False): - scale = scale + 1 - folded.append(w * scale[None, :].to(w.dtype)) - kv_weights = folded - else: - self._input_ln_eps = None - fused_kv_weight = torch.cat(kv_weights, dim=0).contiguous() - if attn0.qkv_proj.bias is not None: - kv_biases = [ - a.qkv_proj.bias[q_size:q_size + 2 * kv_size] - for a in layers_attn - ] - self._fused_kv_bias = torch.cat(kv_biases, dim=0).contiguous() - else: - self._fused_kv_bias = None - - if all(has_k_norm): - k_norm0 = layers_attn[0].k_norm - eps = k_norm0.variance_epsilon - eps_set = {a.k_norm.variance_epsilon for a in layers_attn} - assert len(eps_set) == 1, ( - f"DFlash fused k_norm requires all drafter layers to share " - f"variance_epsilon; got {sorted(eps_set)}.") - self._k_norm_stacked = torch.stack( - [a.k_norm.weight.data for a in layers_attn]) - self._k_norm_eps = eps - else: - self._k_norm_stacked = None - self._k_norm_eps = None - self._num_attn_layers = len(layers_attn) - self._num_heads = num_heads - self._head_dim = head_dim - self._num_kv_heads = num_kv_heads - self._fused_kv_weight = fused_kv_weight - - # fused_qk_norm_rope derives YaRN / partial-rotary frequencies on - # the fly, which can disagree with precompute_context_kv's cached - # cos/sin. Only enable it when the drafter uses plain RoPE. - self._has_qk_norm = (all(has_k_norm) - and all(hasattr(a, 'q_norm') for a in layers_attn)) - rope_params = getattr(getattr(attn0, 'pos_embd_params', None), 'rope', - None) - scale_type = getattr(rope_params, 'scale_type', None) - partial_rotary_factor = getattr( - getattr(attn0, 'pretrained_config', None), 'partial_rotary_factor', - 1.0) - self._use_fused_qk_norm_rope = (self._has_qk_norm - and hasattr(attn0, 'apply_qk_norm_rope') - and rope_params is not None - and scale_type - in (None, RotaryScalingType.none) - and partial_rotary_factor == 1.0) - - logger.debug( - f"DFlash: fused KV weights built for {self._num_attn_layers} layers " - f"(fused_kv_weight shape={tuple(self._fused_kv_weight.shape)})") - - def _get_rope_cos_sin(self, positions, dtype=None): - """Get cos/sin for given positions, suitable for apply_rotary_pos_emb. - - Args: - positions: [B, seq_len] - dtype: target dtype for cos/sin (default: keep original) - Returns: - rope_cos: [B, seq, rot_dim//2] (broadcastable with unsqueeze_dim=1) - rope_sin: [B, seq, rot_dim//2] - """ - if not self._rope_initialized: - self._init_rope() - - # rotary_cos_sin: [max_pos, 2, rot_dim//2] - rope_cache = self._rotary_cos_sin[positions] # [B, seq, 2, rot_dim//2] - rope_cos = rope_cache[..., 0, :] # [B, seq, rot_dim//2] - rope_sin = rope_cache[..., 1, :] - if dtype is not None: - rope_cos = rope_cos.to(dtype) - rope_sin = rope_sin.to(dtype) - return rope_cos, rope_sin - - def _warn_inferred_attention_windows(self) -> None: - """Warn once at initialization when checkpoint metadata enables SWA.""" - if getattr(self.config, 'use_sliding_window', None) is not None: - return - - num_hidden_layers = getattr(self.config, 'num_hidden_layers', None) - if num_hidden_layers is None: - num_hidden_layers = len(self.model.layers) - layers_by_window = {} - for layer_idx in range(num_hidden_layers): - window = get_layer_attention_window(self.config, layer_idx) - if window is not None: - layers_by_window.setdefault(window, []).append(layer_idx) - - for window, layer_indices in layers_by_window.items(): - logger.warning( - "DFlash inferred pooled-context sliding-window attention from " - f"checkpoint config for draft layers {layer_indices}: " - f"window={window}. Context attention is truncated to {window} " - "tokens for these layers; if the drafter expects full context, " - "acceptance rate may drop. Set use_sliding_window explicitly " - "to confirm or disable windowing.") - - def _get_attention_mask_args(self, layer_idx): - """Return FlashAttention causal and local-window arguments for a layer.""" - layer_types = getattr(self.config, 'layer_types', None) - is_sliding_layer = False - if layer_types: - layer_type = layer_types[layer_idx % len(layer_types)] - is_sliding_layer = _is_sliding_attention_layer(layer_type) - - sliding_window = get_layer_attention_window(self.config, layer_idx) - is_sliding_layer = is_sliding_layer or sliding_window is not None - if not is_sliding_layer: - return False, (-1, -1) - - causal = self._sliding_layers_causal or sliding_window is not None - if sliding_window is None: - # Legacy drafters without an explicit window preserve their prior - # non-windowed behavior. - return causal, (-1, -1) - # FlashAttention's bounds are inclusive: W tokens are current + W-1 left. - return causal, (sliding_window - 1, 0) - - def _prepare_dflash_trtllm_gen_buffers( - self, - dtype: torch.dtype, - device: torch.device, - max_batch_size: int, - block_size: int, - num_heads: int, - num_kv_heads: int, - head_dim: int, - ) -> None: - trtllm_gen_ops = self._dflash_trtllm_gen_ops - workspace_bytes = trtllm_gen_ops.get_workspace_size( - dtype=dtype, - num_tokens=max_batch_size * block_size, - num_gen_tokens=max_batch_size * block_size, - num_heads=num_heads, - num_kv_heads=num_kv_heads, - head_size=head_dim, - max_num_requests=max_batch_size, - rotary_embedding_dim=0, - fp8_context_fmha=False, - ) - device = torch.device(device) - is_capturing = torch.cuda.is_current_stream_capturing() - if self._dflash_trtllm_gen_device != device: - if is_capturing: - raise RuntimeError( - "DFlash TRTLLM-Gen buffers must be prepared on the current " - "device before CUDA graph capture.") - self._dflash_trtllm_gen_device = device - self._dflash_trtllm_gen_sm_count = ( - torch.cuda.get_device_properties(device).multi_processor_count) - - workspace = self._dflash_trtllm_gen_workspace - workspace_needs_allocation = ( - workspace is None or workspace.device != device - or workspace.numel() * workspace.element_size() < workspace_bytes) - if workspace_needs_allocation: - if is_capturing: - raise RuntimeError( - "The DFlash TRTLLM-Gen workspace must be allocated at the " - "required size before CUDA graph capture.") - self._dflash_trtllm_gen_workspace = torch.empty(workspace_bytes, - dtype=torch.uint8, - device=device) - - sm_count = self._dflash_trtllm_gen_sm_count - counter_bytes = trtllm_gen_ops.get_multi_ctas_kv_counter_size( - num_heads, max_batch_size, sm_count) - counters = self._dflash_trtllm_gen_counters - counters_need_allocation = (counters is None - or counters.device != device - or counters.numel() * - counters.element_size() < counter_bytes) - if counters_need_allocation: - if is_capturing: - raise RuntimeError( - "The DFlash TRTLLM-Gen counter buffer must be allocated at " - "the required size before CUDA graph capture.") - self._dflash_trtllm_gen_counters = torch.zeros(counter_bytes, - dtype=torch.uint8, - device=device) - - append_batch_indices = self._dflash_batch_indices - block_offsets = self._dflash_block_offsets - static_indices_need_allocation = ( - append_batch_indices is None or block_offsets is None - or append_batch_indices.device != device - or block_offsets.device != device - or append_batch_indices.size(0) < max_batch_size - or append_batch_indices.size(1) != block_size - or block_offsets.numel() != block_size) - if static_indices_need_allocation: - if is_capturing: - raise RuntimeError( - "DFlash TRTLLM-Gen index buffers must be allocated at the " - "required size before CUDA graph capture.") - self._dflash_batch_indices = (torch.arange( - max_batch_size, dtype=torch.int32, - device=device).view(-1, 1).expand(-1, block_size).contiguous()) - self._dflash_block_offsets = torch.arange(block_size, - dtype=torch.int32, - device=device) - - def dflash_forward( - self, - noise_embedding: torch.Tensor, - query_positions: torch.Tensor, - num_ctx_per_req: torch.Tensor, - ctx_k_cache: torch.Tensor, - ctx_v_cache: torch.Tensor, - ctx_cache_batch_idx: torch.Tensor, - ctx_kv_cache: Optional[torch.Tensor] = None, - ctx_page_table: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """DFlash draft forward with cross-attention over a pooled K/V buffer. - - All shapes are fixed so the forward is CUDA-graph compatible. - - Args: - noise_embedding: [B, block_size, hidden_size] - query_positions: [B, block_size] - num_ctx_per_req: [B] — per-batch context length in the pool - ctx_k_cache: [pool_batch, L, max_ctx+block_size, nkv, hd] - ctx_v_cache: [pool_batch, L, max_ctx+block_size, nkv, hd] - ctx_cache_batch_idx: [B] — slot index into the pool per batch entry - Returns: - [B * block_size, hidden_size] - """ - if self.dflash_attention_backend == 'TRTLLM': - if ctx_kv_cache is None or ctx_page_table is None: - raise RuntimeError( - "DFlash TRTLLM-Gen requires a paged context cache and page table." - ) - trtllm_gen_ops = self._dflash_trtllm_gen_ops - elif self.dflash_attention_backend == 'VANILLA': - flash_attention = self._dflash_flash_attention - else: - raise ValueError( - "DFlash attention backend must be VANILLA or TRTLLM, got " - f"{self.dflash_attention_backend!r}.") - - if self._fused_kv_weight is None: - self._build_fused_kv_buffers() - - layer0 = self.model.layers[0] - attn0 = layer0.self_attn - q_size = attn0.q_size - kv_size = attn0.kv_size - head_dim = attn0.head_dim - # Uniformity across layers is asserted in _build_fused_kv_buffers (above). - num_heads_per_rank = attn0.num_heads - num_kv_heads_per_rank = attn0.num_key_value_heads - gqa_group_size = num_heads_per_rank // num_kv_heads_per_rank - - has_qk_norm = self._has_qk_norm - is_bf16 = noise_embedding.dtype == torch.bfloat16 - use_fused_qk_norm_rope = self._use_fused_qk_norm_rope and is_bf16 - use_fused_rope = (_flashinfer_rope is not None and has_qk_norm - and is_bf16 and not use_fused_qk_norm_rope) - - B = noise_embedding.shape[0] - block_size = noise_embedding.shape[1] - - hidden_states = noise_embedding # [B, block_size, hidden] - - # Precompute RoPE cos/sin for the pure-PyTorch fallback path only. - # The fused flashinfer path reads self._get_cos_sin_cache() inline. - rope_dtype = hidden_states.dtype - if not use_fused_rope: - q_rope_cos, q_rope_sin = self._get_rope_cos_sin(query_positions, - dtype=rope_dtype) - _rope = RotaryEmbedding.apply_rotary_pos_emb - - # cache_seqlens (BEFORE append). flash_attn appends block_size - # k/v at cache_seqlens[i]..+block_size for batch i. - cache_seqlens_i32 = num_ctx_per_req[:B].to(torch.int32) - cache_batch_idx_i32 = ctx_cache_batch_idx.to(torch.int32) - - if self.dflash_attention_backend == 'TRTLLM': - max_batch_size = ctx_page_table.size(0) - self._prepare_dflash_trtllm_gen_buffers( - hidden_states.dtype, - hidden_states.device, - max_batch_size, - block_size, - num_heads_per_rank, - num_kv_heads_per_rank, - head_dim, - ) - block_tables = ctx_page_table.index_select( - 0, cache_batch_idx_i32.long()) - pages_per_slot = block_tables.size(1) - page_size = ctx_kv_cache.size(-2) - kv_indices = block_tables.flatten() - kv_indptr = torch.arange( - 0, - (B + 1) * pages_per_slot, - pages_per_slot, - dtype=torch.int32, - device=hidden_states.device, - ) - seq_lens_after = cache_seqlens_i32 + block_size - kv_last_page_len = ((seq_lens_after - 1) % page_size) + 1 - batch_indices = self._dflash_batch_indices - append_batch_indices = batch_indices[:B].reshape(-1) - append_positions = ( - cache_seqlens_i32.view(-1, 1) + - self._dflash_block_offsets).reshape(-1).contiguous() - - # Flatten query positions once for the fused QK-norm-RoPE kernel. - query_positions_flat_i32 = query_positions.reshape(-1).to(torch.int32) - - residual = None - - for layer_idx, layer in enumerate(self.model.layers): - attn_mod = layer.self_attn - - # Apply input_layernorm (flatten to 2D for norm, reshape back) - hs_flat = hidden_states.reshape(-1, hidden_states.shape[-1]) - if residual is None: - residual = hidden_states.clone() - hs_normed_flat = layer.input_layernorm(hs_flat) - else: - res_flat = residual.reshape(-1, residual.shape[-1]) - hs_normed_flat, res_flat = layer.input_layernorm( - hs_flat, res_flat) - residual = res_flat.reshape(B, block_size, -1) - - # QKV projection on normed query tokens (2D) - qkv_query = attn_mod.qkv_proj(hs_normed_flat) # [B*blk, qkv_size] - - if use_fused_qk_norm_rope: - # One kernel does q_norm + k_norm + RoPE in-place on qkv. - # Only safe when the drafter's rope params don't use YaRN / - # long-rope / partial-rotary — otherwise fall back to the - # shared-cache path below. - attn_mod.apply_qk_norm_rope(qkv_query, query_positions_flat_i32) - q_all_2d = qkv_query[:, :q_size] - k_noise_2d = qkv_query[:, q_size:q_size + kv_size] - v_noise_2d = qkv_query[:, q_size + kv_size:] - Q_bshd = q_all_2d.reshape(B, block_size, num_heads_per_rank, - head_dim) - k_noise_bshd = k_noise_2d.reshape(B, block_size, - num_kv_heads_per_rank, - head_dim) - v_noise_bshd = v_noise_2d.reshape(B, block_size, - num_kv_heads_per_rank, - head_dim) - elif use_fused_rope: - # Per-head RMSNorm on q/k (returns new contiguous tensors), - # then flashinfer in-place RoPE sharing the same cos/sin cache - # as precompute_context_kv. - q = attn_mod.q_norm(qkv_query[:, :q_size].reshape( - -1, head_dim)).view(-1, q_size) - k = attn_mod.k_norm(qkv_query[:, - q_size:q_size + kv_size].reshape( - -1, - head_dim)).view(-1, kv_size) - _flashinfer_rope( - query_positions_flat_i32, - q, - k, - head_dim, - self._get_cos_sin_cache(), - self._is_neox, - ) - Q_bshd = q.view(B, block_size, num_heads_per_rank, head_dim) - k_noise_bshd = k.view(B, block_size, num_kv_heads_per_rank, - head_dim) - v_noise_bshd = qkv_query[:, q_size + kv_size:].reshape( - B, block_size, num_kv_heads_per_rank, head_dim) - else: - qkv_query_3d = qkv_query.reshape(B, block_size, -1) - q_all = qkv_query_3d[..., :q_size] - k_noise_all = qkv_query_3d[..., q_size:q_size + kv_size] - v_noise_all = qkv_query_3d[..., q_size + kv_size:] - if has_qk_norm: - q_for_rope = attn_mod.q_norm(q_all.reshape( - -1, head_dim)).reshape(B, block_size, q_size) - k_noise_for_rope = attn_mod.k_norm( - k_noise_all.reshape(-1, head_dim)).reshape( - B, block_size, kv_size) - else: - q_for_rope = q_all - k_noise_for_rope = k_noise_all - Q = _rope(q_for_rope.reshape(B, block_size, num_heads_per_rank, - head_dim).transpose(1, 2), - q_rope_cos, - q_rope_sin, - unsqueeze_dim=1, - is_neox=self._is_neox) - k_noise_rope = _rope(k_noise_for_rope.reshape( - B, block_size, num_kv_heads_per_rank, - head_dim).transpose(1, 2), - q_rope_cos, - q_rope_sin, - unsqueeze_dim=1, - is_neox=self._is_neox) - Q_bshd = Q.transpose(1, 2) - k_noise_bshd = k_noise_rope.transpose(1, 2) - v_noise_bshd = v_noise_all.reshape(B, block_size, - num_kv_heads_per_rank, - head_dim) - - # Per-layer view into the pooled ctx cache. - causal, window_size = self._get_attention_mask_args(layer_idx) - dspark_window = (self._dspark_layer_windows[layer_idx] if layer_idx - < len(self._dspark_layer_windows) else (-1, -1)) - if dspark_window != (-1, -1): - window_size = dspark_window - if self.dflash_attention_backend == 'TRTLLM': - layer_cache = ctx_kv_cache[layer_idx] - trtllm_gen_ops.append_paged_kv_cache( - append_key=k_noise_bshd.reshape(-1, num_kv_heads_per_rank, - head_dim).contiguous(), - append_value=v_noise_bshd.reshape(-1, num_kv_heads_per_rank, - head_dim).contiguous(), - batch_indices=append_batch_indices, - positions=append_positions, - paged_kv_cache=layer_cache, - kv_indices=kv_indices, - kv_indptr=kv_indptr, - kv_last_page_len=kv_last_page_len, - kv_layout="HND", - ) - out = torch.empty_like(Q_bshd) - q_flat = Q_bshd.reshape(-1, num_heads_per_rank, head_dim) - out_flat = out.reshape(-1, num_heads_per_rank, head_dim) - window_left = window_size[0] - if causal: - trtllm_gen_ops.batch_decode_with_kv_cache( - query=q_flat, - kv_cache=(layer_cache[:, 0], layer_cache[:, 1]), - workspace_buffer=self._dflash_trtllm_gen_workspace, - block_tables=block_tables, - seq_lens=seq_lens_after, - max_seq_len=pages_per_slot * page_size, - bmm1_scale=head_dim**-0.5, - bmm2_scale=1.0, - window_left=window_left, - out=out_flat, - sinks=None, - enable_pdl=False, - kv_layout="HND", - backend="trtllm-gen", - q_len_per_req=block_size, - max_q_len=None, - cum_seq_lens_q=None, - kv_cache_sf=None, - uses_shared_paged_kv_idx=True, - bmm1_scale_log2=None, - multi_ctas_kv_counter_buffer=self. - _dflash_trtllm_gen_counters, - ) - else: - cum_seq_lens_q = torch.arange( - 0, - (B + 1) * block_size, - block_size, - dtype=torch.int32, - device=hidden_states.device, - ) - cum_seq_lens_kv = torch.cat(( - torch.zeros(1, - dtype=torch.int32, - device=hidden_states.device), - seq_lens_after.cumsum(0, dtype=torch.int32), - )) - trtllm_gen_ops.batch_context_with_kv_cache( - query=q_flat, - kv_cache=(layer_cache[:, 0], layer_cache[:, 1]), - workspace_buffer=self._dflash_trtllm_gen_workspace, - block_tables=block_tables, - seq_lens=seq_lens_after, - max_q_len=block_size, - max_kv_len=pages_per_slot * page_size, - bmm1_scale=head_dim**-0.5, - bmm2_scale=1.0, - batch_size=B, - cum_seq_lens_q=cum_seq_lens_q, - cum_seq_lens_kv=cum_seq_lens_kv, - window_left=window_left, - out=out_flat, - sinks=None, - enable_pdl=False, - kv_layout="HND", - kv_cache_sf=None, - uses_shared_paged_kv_idx=True, - causal=False, - multi_ctas_kv_counter_buffer=self. - _dflash_trtllm_gen_counters, - ) - else: # VANILLA, validated before entering the layer loop. - layer_k_cache = ctx_k_cache[:, layer_idx] - layer_v_cache = ctx_v_cache[:, layer_idx] - - # Pack gqa_group_size query heads sharing a KV head into the - # row dimension: [B, blk, h_q, d] -> [B, group*blk, h_kv, d]. - # Each CTA owns a whole query-head group and streams KV head's context once - # instead of gqa_group_size CTAs each re-reading it. - # Exact only while every row of the block attends to the same - # key set, i.e. non-causal, unwindowed layers. Causal or - # windowed layers mask by row, so they stay unpacked. - pack_gqa = (gqa_group_size > 1 and not causal - and window_size == (-1, -1)) - if pack_gqa: - q_grouped = Q_bshd.reshape(B, block_size, - num_kv_heads_per_rank, - gqa_group_size, head_dim) - q_packed = q_grouped.permute(0, 3, 1, 2, 4) - q_in = q_packed.reshape(B, gqa_group_size * block_size, - num_kv_heads_per_rank, head_dim) - else: - q_in = Q_bshd - out = flash_attention( - q=q_in, - k_cache=layer_k_cache, - v_cache=layer_v_cache, - k=k_noise_bshd, - v=v_noise_bshd, - cache_seqlens=cache_seqlens_i32, - cache_batch_idx=cache_batch_idx_i32, - causal=causal, - window_size=window_size, - ) - if pack_gqa: - # Undo the packing: [B, group*blk, h_kv, d] -> [B, blk, h_q, d]. - out = out.view(B, gqa_group_size, block_size, - num_kv_heads_per_rank, - head_dim).permute(0, 2, 3, 1, 4) - - attn_output = out.reshape(B * block_size, q_size) - - # Per-drafter post-attention gate (no-op for generic DFlash; Laguna - # applies per-head softplus g_proj gating). gate input is the - # input_layernorm output (the attention input). - attn_output = self._post_attention_gate(attn_output, hs_normed_flat, - attn_mod, - num_heads_per_rank, - head_dim) - - # o_proj (flat 2D, handles all-reduce internally) - hidden_out = attn_mod.o_proj(attn_output) - - # Post-attention layernorm + MLP (flat 2D) - 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(B, block_size, -1) - residual = res_flat.reshape(B, block_size, -1) - - # Final norm - hidden_states_out, _ = self.model.norm( - hidden_states.reshape(-1, hidden_states.shape[-1]), - residual.reshape(-1, residual.shape[-1])) - return hidden_states_out - - def forward( - self, - attn_metadata, - input_ids: torch.LongTensor = None, - position_ids: torch.LongTensor | None = None, - inputs_embeds: torch.FloatTensor | None = None, - return_context_logits: bool = False, - spec_metadata=None, - hidden_states: torch.Tensor | None = None, - **kwargs, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Run the draft model and return (hidden_states, hidden_states) for the - speculative-decoding contract.""" - hidden_states_out = self.model( - input_ids=input_ids, - attn_metadata=attn_metadata, - position_ids=position_ids, - inputs_embeds=inputs_embeds, - spec_metadata=spec_metadata, - **kwargs, - ) - - return hidden_states_out, hidden_states_out - - -class DFlashLagunaForCausalLM(DFlashForCausalLM): - """Laguna DFlash drafter. - - The generic block decode lives in DFlashForCausalLM; this subclass supplies - the Laguna draft-layer specifics: per-head g_proj softplus gating and the - per-aux fc_norm applied to captured target features before fc. - """ - - @staticmethod - def _normalize_config(config: PretrainedConfig) -> None: - """Fill TRT-LLM Laguna defaults missing from dense DFlash drafts.""" - if getattr(config, "num_experts", None) is None: - config.num_experts = 0 - if getattr(config, "mlp_layer_types", None) is None: - config.mlp_layer_types = ["dense"] * config.num_hidden_layers - if getattr(config, "block_size", None) is None: - dflash_config = getattr(config, "dflash_config", {}) - if isinstance(dflash_config, dict): - config.block_size = dflash_config.get("block_size", None) - - def __init__(self, - draft_config, - *, - dflash_attention_backend: str = 'VANILLA'): - """Pin the Laguna draft-layer class and enable Laguna-specific behaviors - (context input_layernorm, causal sliding blocks); reject non-per-head - gating.""" - # The checkpoint labels itself with the vLLM name (model_type "llama"); - # remap to the Laguna architecture so TRT-LLM builds the Laguna layers. - draft_config.pretrained_config.architectures = ["LagunaForCausalLM"] - self._normalize_config(draft_config.pretrained_config) - super().__init__( - draft_config, - dflash_attention_backend=dflash_attention_backend, - ) - self._context_input_layernorm = True - self._sliding_layers_causal = True - gating = getattr(self.config, 'gating', True) - if gating not in (True, 'per-head'): - raise NotImplementedError( - f"Laguna DFlash drafter supports per-head gating only, " - f"got gating={gating!r}") - - def load_weights(self, weights, weight_mapper=None, **kwargs): - """Build the per-aux ``fc_norm`` from the drafter's ``aux_hidden_norms.*`` - weights, then defer the remaining weights to the base loader.""" - aux_keys = sorted( - (k for k in weights if k.startswith('aux_hidden_norms.')), - key=lambda k: int(k.split('.')[1])) - if not aux_keys: - raise ValueError( - "Laguna DFlash checkpoint is missing aux_hidden_norms.* weights" - ) - weights = dict(weights) - eps = getattr(self.config, 'rms_norm_eps', 1e-6) - norms = [] - for k in aux_keys: - w = weights.pop(k) - norm = nn.RMSNorm(w.shape[0], - eps=eps, - device='cuda', - elementwise_affine=True, - dtype=w.dtype) - norm.weight.data.copy_(w) - norms.append(norm) - self.fc_norm = nn.ModuleList(norms) - super().load_weights(weights, weight_mapper=weight_mapper, **kwargs) - - def project_target_hidden(self, hidden_states): - """Project captured target features to the draft width: apply the per-aux - ``fc_norm`` to each hidden chunk, then ``fc`` + ``hidden_norm``.""" - hidden_states = hidden_states.to(self.fc.weight.dtype) - fc_norm = getattr(self, 'fc_norm', None) - if fc_norm is not None: - chunks = hidden_states.chunk(len(fc_norm), dim=-1) - hidden_states = torch.cat( - [norm(chunk) for norm, chunk in zip(fc_norm, chunks)], dim=-1) - return self.hidden_norm(self.fc(hidden_states)) - - def _post_attention_gate(self, attn_output, gate_input, attn_mod, num_heads, - head_dim): - """Apply Laguna's per-head softplus output gate (``g_proj``) to the - attention output; a no-op when the layer has no ``g_proj``.""" - g_proj = getattr(attn_mod, 'g_proj', None) - if g_proj is None: - return attn_output - gate = F.softplus(g_proj(gate_input).float()).to(attn_output.dtype) - return (attn_output.unflatten(-1, (num_heads, head_dim)) * - gate.unsqueeze(-1)).flatten(-2) - - class MTPForCausalLM(nn.Module): def __init__( @@ -2474,27 +1109,6 @@ def _build_pard_draft(model_config, draft_config, lm_head, model): return PARDForCausalLM(draft_config) -@register_draft_model(SpeculativeDecodingMode.DFLASH) -def _build_dflash_draft(model_config, draft_config, lm_head, model): - """Build the DFlash drafter. - - Selects the Laguna variant by detecting its architecture in the draft - checkpoint's own config. - """ - draft_arches = getattr(draft_config.pretrained_config, "architectures", - None) or [] - dflash_attention_backend = model_config.spec_config.attention_backend - if any("Laguna" in arch for arch in draft_arches): - return DFlashLagunaForCausalLM( - draft_config, - dflash_attention_backend=dflash_attention_backend, - ) - return DFlashForCausalLM( - draft_config, - dflash_attention_backend=dflash_attention_backend, - ) - - @register_draft_model(SpeculativeDecodingMode.DRAFT_TARGET_ONE_MODEL) def _build_draft_target_one_model_draft(model_config, draft_config, lm_head, model): diff --git a/tests/unittest/_torch/modeling/test_modeling_speculative.py b/tests/unittest/_torch/modeling/test_modeling_speculative.py index db9aaa2782a1..bb3236abe243 100644 --- a/tests/unittest/_torch/modeling/test_modeling_speculative.py +++ b/tests/unittest/_torch/modeling/test_modeling_speculative.py @@ -25,8 +25,8 @@ from tensorrt_llm._torch.attention_backend.interface import RopeParams from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM from tensorrt_llm._torch.models.modeling_speculative import ( - DFlashForCausalLM, Eagle3ForCausalLM, SpecDecOneEngineForCausalLM, ) @@ -301,7 +301,7 @@ def test_dflash_attention_mask_args(): assert wrapper._get_attention_mask_args(1) == (False, (-1, -1)) assert wrapper._get_attention_mask_args(2) == (True, (4095, 0)) - with patch("tensorrt_llm._torch.models.modeling_speculative.logger.warning") as warning: + with patch("tensorrt_llm._torch.models.modeling_dflash.logger.warning") as warning: wrapper._warn_inferred_attention_windows() warning.assert_not_called() @@ -344,7 +344,7 @@ def test_dflash_attention_mask_args(): for layer_idx in range(5): assert laguna_wrapper._get_attention_mask_args(layer_idx) == (True, (511, 0)) - with patch("tensorrt_llm._torch.models.modeling_speculative.logger.warning") as warning: + with patch("tensorrt_llm._torch.models.modeling_dflash.logger.warning") as warning: laguna_wrapper._warn_inferred_attention_windows() warning.assert_called_once_with( "DFlash inferred pooled-context sliding-window attention from checkpoint " 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 a2cbae410d27..61c1fda9432a 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 @@ -23,7 +23,7 @@ import torch import torch.nn.functional as F -from tensorrt_llm._torch.models.modeling_speculative import ( +from tensorrt_llm._torch.models.modeling_dflash import ( DFlashForCausalLM, dspark_layer_window_size, dspark_markov_chain_logits, From 3dd1018a3d97911b29cfaf562d7abad11997961c Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Tue, 18 Aug 2026 22:36:16 -0700 Subject: [PATCH 04/21] [None][chore] Graduate modeling_dflash.py to the ruff toolchain The file only carried 80-column formatting because it was split out of a legacy module; nothing else in models/ is still on yapf. Moving it to Group A rewraps it to 100 columns and puts it under the full ruff rule set, which it passes with no remaining violations. Signed-off-by: Zhenhuan Chen --- .pre-commit-config.yaml | 2 - legacy-files.txt | 1 - pyproject.toml | 1 - ruff-legacy.toml | 1 - tensorrt_llm/_torch/models/modeling_dflash.py | 687 +++++++++--------- 5 files changed, 344 insertions(+), 348 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 182b085aed35..2a259562ce95 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -202,7 +202,6 @@ common-files: &common_files | tensorrt_llm/_torch/models/modeling_bert.py | tensorrt_llm/_torch/models/modeling_clip.py | tensorrt_llm/_torch/models/modeling_deepseekv3.py | - tensorrt_llm/_torch/models/modeling_dflash.py | tensorrt_llm/_torch/models/modeling_exaone4.py | tensorrt_llm/_torch/models/modeling_gemma3.py | tensorrt_llm/_torch/models/modeling_gemma3vl.py | @@ -969,7 +968,6 @@ legacy-files: &legacy_files | tensorrt_llm/_torch/models/modeling_bert.py | tensorrt_llm/_torch/models/modeling_clip.py | tensorrt_llm/_torch/models/modeling_deepseekv3.py | - tensorrt_llm/_torch/models/modeling_dflash.py | tensorrt_llm/_torch/models/modeling_exaone4.py | tensorrt_llm/_torch/models/modeling_gemma3.py | tensorrt_llm/_torch/models/modeling_gemma3vl.py | diff --git a/legacy-files.txt b/legacy-files.txt index b3c345ff143f..73fe3fee5899 100644 --- a/legacy-files.txt +++ b/legacy-files.txt @@ -194,7 +194,6 @@ tensorrt_llm/_torch/models/modeling_auto.py tensorrt_llm/_torch/models/modeling_bert.py tensorrt_llm/_torch/models/modeling_clip.py tensorrt_llm/_torch/models/modeling_deepseekv3.py -tensorrt_llm/_torch/models/modeling_dflash.py tensorrt_llm/_torch/models/modeling_exaone4.py tensorrt_llm/_torch/models/modeling_gemma3.py tensorrt_llm/_torch/models/modeling_gemma3vl.py diff --git a/pyproject.toml b/pyproject.toml index f25e25193644..f8e1447ca03b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -251,7 +251,6 @@ exclude = [ "tensorrt_llm/_torch/models/modeling_bert.py", "tensorrt_llm/_torch/models/modeling_clip.py", "tensorrt_llm/_torch/models/modeling_deepseekv3.py", - "tensorrt_llm/_torch/models/modeling_dflash.py", "tensorrt_llm/_torch/models/modeling_exaone4.py", "tensorrt_llm/_torch/models/modeling_gemma3.py", "tensorrt_llm/_torch/models/modeling_gemma3vl.py", diff --git a/ruff-legacy.toml b/ruff-legacy.toml index 43715109db9b..612b315564ca 100644 --- a/ruff-legacy.toml +++ b/ruff-legacy.toml @@ -211,7 +211,6 @@ include = [ "tensorrt_llm/_torch/models/modeling_bert.py", "tensorrt_llm/_torch/models/modeling_clip.py", "tensorrt_llm/_torch/models/modeling_deepseekv3.py", - "tensorrt_llm/_torch/models/modeling_dflash.py", "tensorrt_llm/_torch/models/modeling_exaone4.py", "tensorrt_llm/_torch/models/modeling_gemma3.py", "tensorrt_llm/_torch/models/modeling_gemma3vl.py", diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index d6e1cb4dccf2..c3e4152904dc 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -15,20 +15,18 @@ from ..modules.rotary_embedding import RotaryEmbedding try: - from ..custom_ops import \ - flashinfer_apply_rope_with_cos_sin_cache_inplace as _flashinfer_rope + from ..custom_ops import flashinfer_apply_rope_with_cos_sin_cache_inplace as _flashinfer_rope except ImportError: _flashinfer_rope = None -from ..pyexecutor.config_utils import (_is_sliding_attention_layer, - get_layer_attention_window) -from ..speculative.dflash_attention import (get_dflash_flash_attention, - get_dflash_trtllm_gen_ops) +from ..pyexecutor.config_utils import _is_sliding_attention_layer, get_layer_attention_window +from ..speculative.dflash_attention import get_dflash_flash_attention, get_dflash_trtllm_gen_ops from ..speculative.interface import SpeculativeDecodingMode from .modeling_utils import get_model_architecture, register_draft_model -def dspark_layer_window_size(use_swa: bool, swa_window: int, layer_types, - layer_idx: int) -> tuple[int, int]: +def dspark_layer_window_size( + use_swa: bool, swa_window: int, layer_types, layer_idx: int +) -> tuple[int, int]: """flash-attn ``window_size`` for one draft layer of the block decode. DSpark drafters (deepseek-ai/DeepSpec) run the draft block through HF @@ -44,14 +42,18 @@ def dspark_layer_window_size(use_swa: bool, swa_window: int, layer_types, """ if not use_swa: return (-1, -1) - if layer_types is not None and layer_idx < len(layer_types) and \ - layer_types[layer_idx] != 'sliding_attention': + if ( + layer_types is not None + and layer_idx < len(layer_types) + and layer_types[layer_idx] != "sliding_attention" + ): return (-1, -1) return (swa_window - 1, swa_window - 1) -def dspark_markov_step_bias(prev_tokens: torch.Tensor, markov_w1: torch.Tensor, - markov_w2: torch.Tensor) -> torch.Tensor: +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. Reference: DeepSpec ``VanillaMarkov`` (deepspec/modeling/dspark/ @@ -122,10 +124,7 @@ class DFlashForCausalLM(nn.Module): Reference: https://arxiv.org/pdf/2602.06036 """ - def __init__(self, - draft_config, - *, - dflash_attention_backend: str = 'VANILLA'): + 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 custom DFlash architecture label).""" @@ -136,11 +135,11 @@ def __init__(self, DraftModelClass, _ = get_model_architecture(pretrained_cfg) except RuntimeError: model_type = pretrained_cfg.model_type - arch_name = "".join(w.capitalize() - for w in model_type.split("_")) + "ForCausalLM" + arch_name = "".join(w.capitalize() for w in model_type.split("_")) + "ForCausalLM" logger.info( f"DFlash: architecture {pretrained_cfg.architectures} not found, " - f"falling back to {arch_name} based on model_type={model_type}") + f"falling back to {arch_name} based on model_type={model_type}" + ) original_archs = pretrained_cfg.architectures try: pretrained_cfg.architectures = [arch_name] @@ -149,9 +148,7 @@ def __init__(self, pretrained_cfg.architectures = original_archs # Remove spec_config to prevent recursive spec-dec initialization - draft_config_no_spec = replace(draft_config, - spec_config=None, - lm_head_gather_output=False) + draft_config_no_spec = replace(draft_config, spec_config=None, lm_head_gather_output=False) # Weights will be loaded later by ModelLoader.load_draft_weights() self.draft_model_full = DraftModelClass(draft_config_no_spec) @@ -164,24 +161,26 @@ def __init__(self, # Get mask_token_id from dflash_config pretrained_config = draft_config.pretrained_config - dflash_config = getattr(pretrained_config, 'dflash_config', {}) + dflash_config = getattr(pretrained_config, "dflash_config", {}) self.mask_token_id = dflash_config.get( - 'mask_token_id', - getattr(pretrained_config, 'mask_token_id', - pretrained_config.vocab_size)) + "mask_token_id", + getattr(pretrained_config, "mask_token_id", pretrained_config.vocab_size), + ) - self.target_layer_ids = dflash_config.get('target_layer_ids', None) + self.target_layer_ids = dflash_config.get("target_layer_ids", None) self.block_size = dflash_config.get( - 'block_size', getattr(pretrained_config, 'block_size', None)) + "block_size", getattr(pretrained_config, "block_size", None) + ) self.dflash_attention_backend = dflash_attention_backend - if self.dflash_attention_backend == 'VANILLA': + if self.dflash_attention_backend == "VANILLA": self._dflash_flash_attention = get_dflash_flash_attention() - elif self.dflash_attention_backend == 'TRTLLM': + elif self.dflash_attention_backend == "TRTLLM": self._dflash_trtllm_gen_ops = get_dflash_trtllm_gen_ops() else: raise ValueError( "DFlash attention backend must be VANILLA or TRTLLM, got " - f"{self.dflash_attention_backend!r}.") + 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) @@ -191,7 +190,8 @@ def __init__(self, logger.info( f"DFlash draft model initialized with mask_token_id: {self.mask_token_id}, " f"target_layer_ids: {self.target_layer_ids}, block_size: {self.block_size}, " - f"attention_backend: {self.dflash_attention_backend}") + f"attention_backend: {self.dflash_attention_backend}" + ) # DSpark drafters (DFlash + low-rank Markov head + confidence head, # arXiv 2607.05147; reference: deepseek-ai/DeepSpec). The weights- @@ -205,16 +205,14 @@ def __init__(self, # Confidence-scheduled verification is NOT implemented yet: the # confidence_proj weights are loaded (for the follow-up MR) but never # used, and drafting always proposes the full K tokens. - self._dspark_shift_label = bool(dflash_config.get('shift_label', False)) - self._dspark_use_swa = bool(dflash_config.get('use_swa', False)) - self._dspark_swa_window = int( - dflash_config.get('swa_window_size', 0) or 0) - self._dspark_markov_rank = int(dflash_config.get('markov_rank', 0) or 0) + self._dspark_shift_label = bool(dflash_config.get("shift_label", False)) + self._dspark_use_swa = bool(dflash_config.get("use_swa", False)) + self._dspark_swa_window = int(dflash_config.get("swa_window_size", 0) or 0) + self._dspark_markov_rank = int(dflash_config.get("markov_rank", 0) or 0) self._dspark_markov_head_type = str( - dflash_config.get('markov_head_type', 'vanilla') - or 'vanilla').lower() - self._dspark_use_confidence_head = bool( - dflash_config.get('use_confidence_head', False)) + dflash_config.get("markov_head_type", "vanilla") or "vanilla" + ).lower() + self._dspark_use_confidence_head = bool(dflash_config.get("use_confidence_head", False)) # Plain None placeholders rather than nn.Parameter/buffer: most # DFlash checkpoints don't ship these heads, and their shapes # ([vocab, rank]) are checkpoint-dependent, so nothing is @@ -225,34 +223,38 @@ def __init__(self, self.confidence_proj_weight = None # loaded, unused (follow-up MR) self.confidence_proj_bias = None - if self._dspark_markov_rank > 0 and \ - self._dspark_markov_head_type != 'vanilla': + if self._dspark_markov_rank > 0 and self._dspark_markov_head_type != "vanilla": raise ValueError( f"DFlash dspark drafter declares markov_head_type=" f"'{self._dspark_markov_head_type}'; only 'vanilla' is " - "supported (gated/rnn heads need per-step hidden features).") + "supported (gated/rnn heads need per-step hidden features)." + ) if self._dspark_use_swa and self._dspark_swa_window < 1: raise ValueError( "DFlash dspark drafter sets use_swa but swa_window_size=" - f"{dflash_config.get('swa_window_size')} is invalid.") + f"{dflash_config.get('swa_window_size')} is invalid." + ) # causal=true is only invalid under the dspark convention. Legacy # DFlash drafter configs (e.g. Laguna) also carry a causal field; # their causality is handled by the legacy decode path # (_sliding_layers_causal), so don't reject them here. - is_dspark = (str(dflash_config.get('projector_type', '') - or '').lower() == 'dspark' or self._dspark_shift_label - or self._dspark_use_swa or self._dspark_markov_rank > 0 - or self._dspark_use_confidence_head) - if is_dspark and dflash_config.get('causal'): + is_dspark = ( + str(dflash_config.get("projector_type", "") or "").lower() == "dspark" + or self._dspark_shift_label + or self._dspark_use_swa + or self._dspark_markov_rank > 0 + or self._dspark_use_confidence_head + ) + if is_dspark and dflash_config.get("causal"): raise ValueError( "DFlash dspark drafter sets causal=true; the block decode " - "only supports the non-causal dspark convention.") + "only supports the non-causal dspark convention." + ) # Per-layer flash-attn window for the block decode, resolved once. - num_draft_layers = getattr(pretrained_config, 'num_hidden_layers', 0) - layer_types = getattr(pretrained_config, 'layer_types', None) + num_draft_layers = getattr(pretrained_config, "num_hidden_layers", 0) + layer_types = getattr(pretrained_config, "layer_types", None) self._dspark_layer_windows = [ - dspark_layer_window_size(self._dspark_use_swa, - self._dspark_swa_window, layer_types, i) + dspark_layer_window_size(self._dspark_use_swa, self._dspark_swa_window, layer_types, i) for i in range(num_draft_layers) ] if self._dspark_use_confidence_head: @@ -260,7 +262,8 @@ def __init__(self, "DFlash dspark drafter declares use_confidence_head; " "confidence-scheduled verification is not implemented yet " "(confidence_proj weights are loaded but unused, drafting " - "always proposes the full K tokens).") + "always proposes the full K tokens)." + ) self.logits_processor = None # Set by caller after construction @@ -312,9 +315,7 @@ def _validate_uniform_rope(self): if len(self.model.layers) == 0: raise ValueError("DFlash requires at least one draft model layer.") - signatures = [ - self._rope_signature(layer.self_attn) for layer in self.model.layers - ] + signatures = [self._rope_signature(layer.self_attn) for layer in self.model.layers] mismatched_layers = [ layer_idx @@ -322,11 +323,12 @@ def _validate_uniform_rope(self): if signature != signatures[0] ] if mismatched_layers: - layer_types = getattr(self.config, 'layer_types', None) + layer_types = getattr(self.config, "layer_types", None) raise ValueError( "DFlash shares one RoPE cache across draft layers, but layers " f"{mismatched_layers} have a different effective RoPE " - f"configuration from layer 0. layer_types={layer_types}.") + f"configuration from layer 0. layer_types={layer_types}." + ) def _init_rope(self): """Initialize RoPE from the draft model's attention configuration. @@ -352,16 +354,15 @@ def _init_rope(self): else: # Fallback: basic NeoX-style RoPE config = self.config - head_dim = getattr(config, 'head_dim', - config.hidden_size // config.num_attention_heads) - rope_theta = getattr(config, 'rope_theta', 1000000.0) - max_pos = getattr(config, 'max_position_embeddings', 32768) - - inv_freq = 1.0 / (rope_theta**(torch.arange( - 0, head_dim, 2, dtype=torch.float32, device='cuda') / head_dim)) - positions = torch.arange(max_pos, - dtype=torch.float32, - device='cuda') + head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + rope_theta = getattr(config, "rope_theta", 1000000.0) + max_pos = getattr(config, "max_position_embeddings", 32768) + + inv_freq = 1.0 / ( + rope_theta + ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device="cuda") / head_dim) + ) + positions = torch.arange(max_pos, dtype=torch.float32, device="cuda") freqs = torch.outer(positions, inv_freq) rope_cos = freqs.cos().to(config.torch_dtype) rope_sin = freqs.sin().to(config.torch_dtype) @@ -371,8 +372,7 @@ def _init_rope(self): self._rope_initialized = True - def project_target_hidden(self, - hidden_states: torch.Tensor) -> torch.Tensor: + def project_target_hidden(self, hidden_states: torch.Tensor) -> torch.Tensor: """Project captured target hidden states into the draft hidden space. Generic DFlash: fc then hidden_norm. Subclasses (e.g. Laguna) may @@ -386,11 +386,12 @@ def has_markov_head(self) -> bool: return self._dspark_markov_rank > 0 and self.markov_w1 is not None def apply_markov_chain_logits( - self, - base_logits: torch.Tensor, - first_prev_tokens: torch.Tensor, - argmax_fn=None, - vocab_slice: slice | None = None) -> torch.Tensor: + self, + base_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + argmax_fn=None, + vocab_slice: slice | None = None, + ) -> torch.Tensor: """Apply the dspark vanilla-Markov intra-block bias to block logits. No-op (returns ``base_logits`` unchanged) for non-dspark drafters. @@ -402,16 +403,12 @@ def apply_markov_chain_logits( """ if not self.has_markov_head: return base_logits - markov_w2 = self.markov_w2 if vocab_slice is None else \ - self.markov_w2[vocab_slice] - return dspark_markov_chain_logits(base_logits, - first_prev_tokens, - self.markov_w1, - markov_w2, - argmax_fn=argmax_fn) - - def _post_attention_gate(self, attn_output, gate_input, attn_mod, num_heads, - head_dim): + markov_w2 = self.markov_w2 if vocab_slice is None else self.markov_w2[vocab_slice] + return dspark_markov_chain_logits( + base_logits, first_prev_tokens, self.markov_w1, markov_w2, argmax_fn=argmax_fn + ) + + def _post_attention_gate(self, attn_output, gate_input, attn_mod, num_heads, head_dim): """Hook applied to the block-attention output before o_proj. No-op for generic DFlash; overridden by drafters that gate (e.g. Laguna). @@ -428,30 +425,31 @@ def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): """ # Laguna DFlash checkpoints may ship a fused self_attn.qkv_proj; the draft # loader expects split q/k/v (a fused key is silently dropped otherwise). - if any(k.endswith('self_attn.qkv_proj.weight') for k in weights): - for attr in ('num_attention_heads_per_layer', - 'num_key_value_heads_per_layer'): + if any(k.endswith("self_attn.qkv_proj.weight") for k in weights): + for attr in ("num_attention_heads_per_layer", "num_key_value_heads_per_layer"): per_layer = getattr(self.config, attr, None) if per_layer is not None and len(set(per_layer)) > 1: raise ValueError( "DFlash load_weights() splits the fused qkv_proj using " "the global head count, but the drafter has heterogeneous " f"{attr} {sorted(set(per_layer))}; per-layer qkv splitting " - "is required for this checkpoint.") + "is required for this checkpoint." + ) head_dim = getattr( - self.config, 'head_dim', - self.config.hidden_size // self.config.num_attention_heads) - num_kv_heads = getattr(self.config, 'num_key_value_heads', - self.config.num_attention_heads) + self.config, "head_dim", self.config.hidden_size // self.config.num_attention_heads + ) + num_kv_heads = getattr( + self.config, "num_key_value_heads", self.config.num_attention_heads + ) q = self.config.num_attention_heads * head_dim kv = num_kv_heads * head_dim split = {} for k, v in weights.items(): - if k.endswith('self_attn.qkv_proj.weight'): - b = k[:-len('qkv_proj.weight')] - split[b + 'q_proj.weight'] = v[:q] - split[b + 'k_proj.weight'] = v[q:q + kv] - split[b + 'v_proj.weight'] = v[q + kv:] + if k.endswith("self_attn.qkv_proj.weight"): + b = k[: -len("qkv_proj.weight")] + split[b + "q_proj.weight"] = v[:q] + split[b + "k_proj.weight"] = v[q : q + kv] + split[b + "v_proj.weight"] = v[q + kv :] else: split[k] = v weights = split @@ -461,80 +459,83 @@ def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): # markov_w1/markov_w2 drive the intra-block logit bias; the # confidence_proj weights are loaded for the confidence-scheduling # follow-up MR but are not used yet. - dspark_keys = ('markov_w1.weight', 'markov_w2.weight', - 'confidence_proj.weight', 'confidence_proj.bias') + dspark_keys = ( + "markov_w1.weight", + "markov_w2.weight", + "confidence_proj.weight", + "confidence_proj.bias", + ) dspark_weights = {k: weights[k] for k in dspark_keys if k in weights} if dspark_weights: - weights = { - k: v - for k, v in weights.items() if k not in dspark_weights - } + weights = {k: v for k, v in weights.items() if k not in dspark_weights} if self._dspark_markov_rank > 0: vocab = self.config.vocab_size rank = self._dspark_markov_rank - for k in ('markov_w1.weight', 'markov_w2.weight'): + for k in ("markov_w1.weight", "markov_w2.weight"): if k not in dspark_weights: raise ValueError( f"DFlash dspark drafter declares markov_rank=" f"{self._dspark_markov_rank} but the checkpoint is " - f"missing {k}.") + f"missing {k}." + ) if tuple(dspark_weights[k].shape) != (vocab, rank): raise ValueError( f"DFlash dspark {k} has shape " f"{tuple(dspark_weights[k].shape)}, expected " - f"[vocab, markov_rank] = ({vocab}, {rank}).") - self.markov_w1 = dspark_weights['markov_w1.weight'].to('cuda') - self.markov_w2 = dspark_weights['markov_w2.weight'].to('cuda') - if 'confidence_proj.weight' in dspark_weights: - 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') + f"[vocab, markov_rank] = ({vocab}, {rank})." + ) + self.markov_w1 = dspark_weights["markov_w1.weight"].to("cuda") + self.markov_w2 = dspark_weights["markov_w2.weight"].to("cuda") + if "confidence_proj.weight" in dspark_weights: + 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") # Remap: add 'model.' prefix where needed, and extract DFlash-specific weights remapped = {} for key, value in weights.items(): - if key in ('fc.weight', 'hidden_norm.weight'): + if key in ("fc.weight", "hidden_norm.weight"): # DFlash-specific projection weights - store directly remapped[key] = value - elif key == 'norm.weight': - remapped['model.norm.weight'] = value - elif not key.startswith('model.'): - remapped[f'model.{key}'] = value + elif key == "norm.weight": + remapped["model.norm.weight"] = value + elif not key.startswith("model."): + remapped[f"model.{key}"] = value else: remapped[key] = value # Load DFlash-specific weights directly - if 'fc.weight' in remapped: - self.fc = nn.Linear(remapped['fc.weight'].shape[1], - remapped['fc.weight'].shape[0], - bias=False, - device='cuda', - dtype=remapped['fc.weight'].dtype) - self.fc.weight.data.copy_(remapped['fc.weight']) - del remapped['fc.weight'] - - if 'hidden_norm.weight' in remapped: - rms_norm_eps = getattr(self.config, 'rms_norm_eps', 1e-6) + if "fc.weight" in remapped: + self.fc = nn.Linear( + remapped["fc.weight"].shape[1], + remapped["fc.weight"].shape[0], + bias=False, + device="cuda", + dtype=remapped["fc.weight"].dtype, + ) + self.fc.weight.data.copy_(remapped["fc.weight"]) + del remapped["fc.weight"] + + if "hidden_norm.weight" in remapped: + rms_norm_eps = getattr(self.config, "rms_norm_eps", 1e-6) self.hidden_norm = nn.RMSNorm( - remapped['hidden_norm.weight'].shape[0], + remapped["hidden_norm.weight"].shape[0], eps=rms_norm_eps, - device='cuda', + device="cuda", elementwise_affine=True, - dtype=remapped['hidden_norm.weight'].dtype) - self.hidden_norm.weight.data.copy_(remapped['hidden_norm.weight']) - del remapped['hidden_norm.weight'] + dtype=remapped["hidden_norm.weight"].dtype, + ) + 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. - self.draft_model_full.load_weights(weights=remapped, - weight_mapper=weight_mapper, - allow_partial_loading=True) + self.draft_model_full.load_weights( + weights=remapped, weight_mapper=weight_mapper, allow_partial_loading=True + ) - def load_weights_from_target_model(self, - target_model: torch.nn.Module) -> None: + def load_weights_from_target_model(self, target_model: torch.nn.Module) -> None: """Share embed_tokens and lm_head from the target model.""" self.draft_model_full.model.embed_tokens = target_model.model.embed_tokens self.draft_model_full.lm_head = target_model.lm_head @@ -561,16 +562,14 @@ def precompute_context_kv( nkv = self._num_kv_heads hd = self._head_dim weight_dtype = self._fused_kv_weight.dtype - if getattr(self, '_input_ln_eps', None) is not None: + if getattr(self, "_input_ln_eps", None) is not None: ph = projected_hidden.float() - ph = ph * torch.rsqrt( - ph.pow(2).mean(-1, keepdim=True) + self._input_ln_eps) + ph = ph * torch.rsqrt(ph.pow(2).mean(-1, keepdim=True) + self._input_ln_eps) projected_hidden = ph.to(weight_dtype) elif projected_hidden.dtype != weight_dtype: projected_hidden = projected_hidden.to(weight_dtype) - kv_flat = F.linear(projected_hidden, self._fused_kv_weight, - self._fused_kv_bias) + kv_flat = F.linear(projected_hidden, self._fused_kv_weight, self._fused_kv_bias) # Per-layer layout [L0_K|L0_V|L1_K|L1_V|...] keeps K and V contiguous # after the select() splits — no extra copy required. kv = kv_flat.view(N, L, 2, nkv, hd) @@ -580,7 +579,7 @@ def precompute_context_kv( if self._k_norm_stacked is not None: # Fuse L per-layer RMSNorms into one. k is [N, L, nkv, hd]; # each layer has its own weight ([L, hd]) but shares eps. - k = F.rms_norm(k, (hd, ), eps=self._k_norm_eps) + k = F.rms_norm(k, (hd,), eps=self._k_norm_eps) k = k * self._k_norm_stacked.view(1, L, 1, hd) self._fused_rope_inplace(k.view(N * L, nkv * hd), positions, N, L) @@ -598,8 +597,9 @@ def _get_cos_sin_cache(self) -> torch.Tensor: if not self._rope_initialized: self._init_rope() max_pos = self._rotary_cos_sin.shape[0] - self._cos_sin_cache_fp32 = self._rotary_cos_sin.view(max_pos, -1).to( - torch.float32).contiguous() + self._cos_sin_cache_fp32 = ( + self._rotary_cos_sin.view(max_pos, -1).to(torch.float32).contiguous() + ) return self._cos_sin_cache_fp32 def _fused_rope_inplace( @@ -623,8 +623,7 @@ def _fused_rope_inplace( # scratch so the extra rotate is negligible. need_rows = k_flat.shape[0] dummy_q = self._rope_dummy_q - if (dummy_q is None or dummy_q.dtype != k_flat.dtype - or dummy_q.shape[0] < need_rows): + if dummy_q is None or dummy_q.dtype != k_flat.dtype or dummy_q.shape[0] < need_rows: dummy_q = k_flat.new_empty(need_rows, self._head_dim) self._rope_dummy_q = dummy_q _flashinfer_rope( @@ -638,8 +637,7 @@ def _fused_rope_inplace( return # Pure-PyTorch fallback (older environments without flashinfer). - cos, sin = self._get_rope_cos_sin(positions_int32.view(1, -1), - dtype=k_flat.dtype) + cos, sin = self._get_rope_cos_sin(positions_int32.view(1, -1), dtype=k_flat.dtype) k_roped = RotaryEmbedding.apply_rotary_pos_emb( k_flat.view(k_flat.shape[0], -1, self._head_dim), cos.squeeze(0), @@ -667,41 +665,46 @@ def _build_fused_kv_buffers(self) -> None: # uniformity (the target uses per-layer heads, the drafter does not). for a in layers_attn[1:]: assert ( - a.q_size == q_size and a.kv_size == kv_size - and a.head_dim == head_dim and a.num_heads == num_heads - and a.num_key_value_heads == num_kv_heads), ( - "DFlash fused KV requires all drafter layers to share " - "q_size / kv_size / head_dim / num_heads / num_kv_heads.") + a.q_size == q_size + and a.kv_size == kv_size + and a.head_dim == head_dim + and a.num_heads == num_heads + and a.num_key_value_heads == num_kv_heads + ), ( + "DFlash fused KV requires all drafter layers to share " + "q_size / kv_size / head_dim / num_heads / num_kv_heads." + ) - has_k_norm = [hasattr(a, 'k_norm') for a in layers_attn] + has_k_norm = [hasattr(a, "k_norm") for a in layers_attn] assert all(has_k_norm) or not any(has_k_norm), ( "DFlash fused KV requires either all or no drafter layers to have k_norm." ) - kv_weights = [ - a.qkv_proj.weight[q_size:q_size + 2 * kv_size] for a in layers_attn - ] + kv_weights = [a.qkv_proj.weight[q_size : q_size + 2 * kv_size] for a in layers_attn] # Fold each drafter layer's input_layernorm weight into its KV projection # so context K/V match the query path. vLLM laguna_dflash applies # layer.input_layernorm to context states before KV; RMSNorm gives # (x_hat * w) @ Wkv.T == x_hat @ (Wkv * w).T, and the shared 1/rms(x) is # applied to projected_hidden in precompute_context_kv. dlayers = self.model.layers - if self._context_input_layernorm and all( - hasattr(dl, 'input_layernorm') for dl in dlayers): + if self._context_input_layernorm and all(hasattr(dl, "input_layernorm") for dl in dlayers): eps_set = { - getattr(dl.input_layernorm, 'variance_epsilon', - getattr(self.config, 'rms_norm_eps', 1e-6)) + getattr( + dl.input_layernorm, + "variance_epsilon", + getattr(self.config, "rms_norm_eps", 1e-6), + ) for dl in dlayers } assert len(eps_set) == 1, ( "DFlash fused context input_layernorm needs all drafter layers " - f"to share variance_epsilon; got {sorted(eps_set)}") + f"to share variance_epsilon; got {sorted(eps_set)}" + ) self._input_ln_eps = eps_set.pop() folded = [] for w, dl in zip(kv_weights, dlayers): scale = dl.input_layernorm.weight.data - if getattr(dl.input_layernorm, 'use_gemma', False): + if getattr(dl.input_layernorm, "use_gemma", False): scale = scale + 1 folded.append(w * scale[None, :].to(w.dtype)) kv_weights = folded @@ -709,10 +712,7 @@ def _build_fused_kv_buffers(self) -> None: self._input_ln_eps = None fused_kv_weight = torch.cat(kv_weights, dim=0).contiguous() if attn0.qkv_proj.bias is not None: - kv_biases = [ - a.qkv_proj.bias[q_size:q_size + 2 * kv_size] - for a in layers_attn - ] + kv_biases = [a.qkv_proj.bias[q_size : q_size + 2 * kv_size] for a in layers_attn] self._fused_kv_bias = torch.cat(kv_biases, dim=0).contiguous() else: self._fused_kv_bias = None @@ -723,9 +723,9 @@ def _build_fused_kv_buffers(self) -> None: eps_set = {a.k_norm.variance_epsilon for a in layers_attn} assert len(eps_set) == 1, ( f"DFlash fused k_norm requires all drafter layers to share " - f"variance_epsilon; got {sorted(eps_set)}.") - self._k_norm_stacked = torch.stack( - [a.k_norm.weight.data for a in layers_attn]) + f"variance_epsilon; got {sorted(eps_set)}." + ) + self._k_norm_stacked = torch.stack([a.k_norm.weight.data for a in layers_attn]) self._k_norm_eps = eps else: self._k_norm_stacked = None @@ -739,24 +739,24 @@ def _build_fused_kv_buffers(self) -> None: # fused_qk_norm_rope derives YaRN / partial-rotary frequencies on # the fly, which can disagree with precompute_context_kv's cached # cos/sin. Only enable it when the drafter uses plain RoPE. - self._has_qk_norm = (all(has_k_norm) - and all(hasattr(a, 'q_norm') for a in layers_attn)) - rope_params = getattr(getattr(attn0, 'pos_embd_params', None), 'rope', - None) - scale_type = getattr(rope_params, 'scale_type', None) + self._has_qk_norm = all(has_k_norm) and all(hasattr(a, "q_norm") for a in layers_attn) + rope_params = getattr(getattr(attn0, "pos_embd_params", None), "rope", None) + scale_type = getattr(rope_params, "scale_type", None) partial_rotary_factor = getattr( - getattr(attn0, 'pretrained_config', None), 'partial_rotary_factor', - 1.0) - self._use_fused_qk_norm_rope = (self._has_qk_norm - and hasattr(attn0, 'apply_qk_norm_rope') - and rope_params is not None - and scale_type - in (None, RotaryScalingType.none) - and partial_rotary_factor == 1.0) + getattr(attn0, "pretrained_config", None), "partial_rotary_factor", 1.0 + ) + self._use_fused_qk_norm_rope = ( + self._has_qk_norm + and hasattr(attn0, "apply_qk_norm_rope") + and rope_params is not None + and scale_type in (None, RotaryScalingType.none) + and partial_rotary_factor == 1.0 + ) logger.debug( f"DFlash: fused KV weights built for {self._num_attn_layers} layers " - f"(fused_kv_weight shape={tuple(self._fused_kv_weight.shape)})") + f"(fused_kv_weight shape={tuple(self._fused_kv_weight.shape)})" + ) def _get_rope_cos_sin(self, positions, dtype=None): """Get cos/sin for given positions, suitable for apply_rotary_pos_emb. @@ -782,10 +782,10 @@ def _get_rope_cos_sin(self, positions, dtype=None): def _warn_inferred_attention_windows(self) -> None: """Warn once at initialization when checkpoint metadata enables SWA.""" - if getattr(self.config, 'use_sliding_window', None) is not None: + if getattr(self.config, "use_sliding_window", None) is not None: return - num_hidden_layers = getattr(self.config, 'num_hidden_layers', None) + num_hidden_layers = getattr(self.config, "num_hidden_layers", None) if num_hidden_layers is None: num_hidden_layers = len(self.model.layers) layers_by_window = {} @@ -801,11 +801,12 @@ def _warn_inferred_attention_windows(self) -> None: f"window={window}. Context attention is truncated to {window} " "tokens for these layers; if the drafter expects full context, " "acceptance rate may drop. Set use_sliding_window explicitly " - "to confirm or disable windowing.") + "to confirm or disable windowing." + ) def _get_attention_mask_args(self, layer_idx): """Return FlashAttention causal and local-window arguments for a layer.""" - layer_types = getattr(self.config, 'layer_types', None) + layer_types = getattr(self.config, "layer_types", None) is_sliding_layer = False if layer_types: layer_type = layer_types[layer_idx % len(layer_types)] @@ -852,61 +853,73 @@ def _prepare_dflash_trtllm_gen_buffers( if is_capturing: raise RuntimeError( "DFlash TRTLLM-Gen buffers must be prepared on the current " - "device before CUDA graph capture.") + "device before CUDA graph capture." + ) self._dflash_trtllm_gen_device = device - self._dflash_trtllm_gen_sm_count = ( - torch.cuda.get_device_properties(device).multi_processor_count) + self._dflash_trtllm_gen_sm_count = torch.cuda.get_device_properties( + device + ).multi_processor_count workspace = self._dflash_trtllm_gen_workspace workspace_needs_allocation = ( - workspace is None or workspace.device != device - or workspace.numel() * workspace.element_size() < workspace_bytes) + workspace is None + or workspace.device != device + or workspace.numel() * workspace.element_size() < workspace_bytes + ) if workspace_needs_allocation: if is_capturing: raise RuntimeError( "The DFlash TRTLLM-Gen workspace must be allocated at the " - "required size before CUDA graph capture.") - self._dflash_trtllm_gen_workspace = torch.empty(workspace_bytes, - dtype=torch.uint8, - device=device) + "required size before CUDA graph capture." + ) + self._dflash_trtllm_gen_workspace = torch.empty( + workspace_bytes, dtype=torch.uint8, device=device + ) sm_count = self._dflash_trtllm_gen_sm_count counter_bytes = trtllm_gen_ops.get_multi_ctas_kv_counter_size( - num_heads, max_batch_size, sm_count) + num_heads, max_batch_size, sm_count + ) counters = self._dflash_trtllm_gen_counters - counters_need_allocation = (counters is None - or counters.device != device - or counters.numel() * - counters.element_size() < counter_bytes) + counters_need_allocation = ( + counters is None + or counters.device != device + or counters.numel() * counters.element_size() < counter_bytes + ) if counters_need_allocation: if is_capturing: raise RuntimeError( "The DFlash TRTLLM-Gen counter buffer must be allocated at " - "the required size before CUDA graph capture.") - self._dflash_trtllm_gen_counters = torch.zeros(counter_bytes, - dtype=torch.uint8, - device=device) + "the required size before CUDA graph capture." + ) + self._dflash_trtllm_gen_counters = torch.zeros( + counter_bytes, dtype=torch.uint8, device=device + ) append_batch_indices = self._dflash_batch_indices block_offsets = self._dflash_block_offsets static_indices_need_allocation = ( - append_batch_indices is None or block_offsets is None + append_batch_indices is None + or block_offsets is None or append_batch_indices.device != device or block_offsets.device != device or append_batch_indices.size(0) < max_batch_size or append_batch_indices.size(1) != block_size - or block_offsets.numel() != block_size) + or block_offsets.numel() != block_size + ) if static_indices_need_allocation: if is_capturing: raise RuntimeError( "DFlash TRTLLM-Gen index buffers must be allocated at the " - "required size before CUDA graph capture.") - self._dflash_batch_indices = (torch.arange( - max_batch_size, dtype=torch.int32, - device=device).view(-1, 1).expand(-1, block_size).contiguous()) - self._dflash_block_offsets = torch.arange(block_size, - dtype=torch.int32, - device=device) + "required size before CUDA graph capture." + ) + self._dflash_batch_indices = ( + torch.arange(max_batch_size, dtype=torch.int32, device=device) + .view(-1, 1) + .expand(-1, block_size) + .contiguous() + ) + self._dflash_block_offsets = torch.arange(block_size, dtype=torch.int32, device=device) def dflash_forward( self, @@ -933,18 +946,19 @@ def dflash_forward( Returns: [B * block_size, hidden_size] """ - if self.dflash_attention_backend == 'TRTLLM': + if self.dflash_attention_backend == "TRTLLM": if ctx_kv_cache is None or ctx_page_table is None: raise RuntimeError( "DFlash TRTLLM-Gen requires a paged context cache and page table." ) trtllm_gen_ops = self._dflash_trtllm_gen_ops - elif self.dflash_attention_backend == 'VANILLA': + elif self.dflash_attention_backend == "VANILLA": flash_attention = self._dflash_flash_attention else: raise ValueError( "DFlash attention backend must be VANILLA or TRTLLM, got " - f"{self.dflash_attention_backend!r}.") + f"{self.dflash_attention_backend!r}." + ) if self._fused_kv_weight is None: self._build_fused_kv_buffers() @@ -962,8 +976,9 @@ def dflash_forward( has_qk_norm = self._has_qk_norm is_bf16 = noise_embedding.dtype == torch.bfloat16 use_fused_qk_norm_rope = self._use_fused_qk_norm_rope and is_bf16 - use_fused_rope = (_flashinfer_rope is not None and has_qk_norm - and is_bf16 and not use_fused_qk_norm_rope) + use_fused_rope = ( + _flashinfer_rope is not None and has_qk_norm and is_bf16 and not use_fused_qk_norm_rope + ) B = noise_embedding.shape[0] block_size = noise_embedding.shape[1] @@ -974,8 +989,7 @@ def dflash_forward( # The fused flashinfer path reads self._get_cos_sin_cache() inline. rope_dtype = hidden_states.dtype if not use_fused_rope: - q_rope_cos, q_rope_sin = self._get_rope_cos_sin(query_positions, - dtype=rope_dtype) + q_rope_cos, q_rope_sin = self._get_rope_cos_sin(query_positions, dtype=rope_dtype) _rope = RotaryEmbedding.apply_rotary_pos_emb # cache_seqlens (BEFORE append). flash_attn appends block_size @@ -983,7 +997,7 @@ def dflash_forward( cache_seqlens_i32 = num_ctx_per_req[:B].to(torch.int32) cache_batch_idx_i32 = ctx_cache_batch_idx.to(torch.int32) - if self.dflash_attention_backend == 'TRTLLM': + if self.dflash_attention_backend == "TRTLLM": max_batch_size = ctx_page_table.size(0) self._prepare_dflash_trtllm_gen_buffers( hidden_states.dtype, @@ -994,8 +1008,7 @@ def dflash_forward( num_kv_heads_per_rank, head_dim, ) - block_tables = ctx_page_table.index_select( - 0, cache_batch_idx_i32.long()) + block_tables = ctx_page_table.index_select(0, cache_batch_idx_i32.long()) pages_per_slot = block_tables.size(1) page_size = ctx_kv_cache.size(-2) kv_indices = block_tables.flatten() @@ -1011,8 +1024,10 @@ def dflash_forward( batch_indices = self._dflash_batch_indices append_batch_indices = batch_indices[:B].reshape(-1) append_positions = ( - cache_seqlens_i32.view(-1, 1) + - self._dflash_block_offsets).reshape(-1).contiguous() + (cache_seqlens_i32.view(-1, 1) + self._dflash_block_offsets) + .reshape(-1) + .contiguous() + ) # Flatten query positions once for the fused QK-norm-RoPE kernel. query_positions_flat_i32 = query_positions.reshape(-1).to(torch.int32) @@ -1029,8 +1044,7 @@ def dflash_forward( hs_normed_flat = layer.input_layernorm(hs_flat) else: res_flat = residual.reshape(-1, residual.shape[-1]) - hs_normed_flat, res_flat = layer.input_layernorm( - hs_flat, res_flat) + hs_normed_flat, res_flat = layer.input_layernorm(hs_flat, res_flat) residual = res_flat.reshape(B, block_size, -1) # QKV projection on normed query tokens (2D) @@ -1043,26 +1057,19 @@ def dflash_forward( # shared-cache path below. attn_mod.apply_qk_norm_rope(qkv_query, query_positions_flat_i32) q_all_2d = qkv_query[:, :q_size] - k_noise_2d = qkv_query[:, q_size:q_size + kv_size] - v_noise_2d = qkv_query[:, q_size + kv_size:] - Q_bshd = q_all_2d.reshape(B, block_size, num_heads_per_rank, - head_dim) - k_noise_bshd = k_noise_2d.reshape(B, block_size, - num_kv_heads_per_rank, - head_dim) - v_noise_bshd = v_noise_2d.reshape(B, block_size, - num_kv_heads_per_rank, - head_dim) + k_noise_2d = qkv_query[:, q_size : q_size + kv_size] + v_noise_2d = qkv_query[:, q_size + kv_size :] + Q_bshd = q_all_2d.reshape(B, block_size, num_heads_per_rank, head_dim) + k_noise_bshd = k_noise_2d.reshape(B, block_size, num_kv_heads_per_rank, head_dim) + v_noise_bshd = v_noise_2d.reshape(B, block_size, num_kv_heads_per_rank, head_dim) elif use_fused_rope: # Per-head RMSNorm on q/k (returns new contiguous tensors), # then flashinfer in-place RoPE sharing the same cos/sin cache # as precompute_context_kv. - q = attn_mod.q_norm(qkv_query[:, :q_size].reshape( - -1, head_dim)).view(-1, q_size) - k = attn_mod.k_norm(qkv_query[:, - q_size:q_size + kv_size].reshape( - -1, - head_dim)).view(-1, kv_size) + q = attn_mod.q_norm(qkv_query[:, :q_size].reshape(-1, head_dim)).view(-1, q_size) + k = attn_mod.k_norm( + qkv_query[:, q_size : q_size + kv_size].reshape(-1, head_dim) + ).view(-1, kv_size) _flashinfer_rope( query_positions_flat_i32, q, @@ -1072,56 +1079,63 @@ def dflash_forward( self._is_neox, ) Q_bshd = q.view(B, block_size, num_heads_per_rank, head_dim) - k_noise_bshd = k.view(B, block_size, num_kv_heads_per_rank, - head_dim) - v_noise_bshd = qkv_query[:, q_size + kv_size:].reshape( - B, block_size, num_kv_heads_per_rank, head_dim) + k_noise_bshd = k.view(B, block_size, num_kv_heads_per_rank, head_dim) + v_noise_bshd = qkv_query[:, q_size + kv_size :].reshape( + B, block_size, num_kv_heads_per_rank, head_dim + ) else: qkv_query_3d = qkv_query.reshape(B, block_size, -1) q_all = qkv_query_3d[..., :q_size] - k_noise_all = qkv_query_3d[..., q_size:q_size + kv_size] - v_noise_all = qkv_query_3d[..., q_size + kv_size:] + k_noise_all = qkv_query_3d[..., q_size : q_size + kv_size] + v_noise_all = qkv_query_3d[..., q_size + kv_size :] if has_qk_norm: - q_for_rope = attn_mod.q_norm(q_all.reshape( - -1, head_dim)).reshape(B, block_size, q_size) - k_noise_for_rope = attn_mod.k_norm( - k_noise_all.reshape(-1, head_dim)).reshape( - B, block_size, kv_size) + q_for_rope = attn_mod.q_norm(q_all.reshape(-1, head_dim)).reshape( + B, block_size, q_size + ) + k_noise_for_rope = attn_mod.k_norm(k_noise_all.reshape(-1, head_dim)).reshape( + B, block_size, kv_size + ) else: q_for_rope = q_all k_noise_for_rope = k_noise_all - Q = _rope(q_for_rope.reshape(B, block_size, num_heads_per_rank, - head_dim).transpose(1, 2), - q_rope_cos, - q_rope_sin, - unsqueeze_dim=1, - is_neox=self._is_neox) - k_noise_rope = _rope(k_noise_for_rope.reshape( - B, block_size, num_kv_heads_per_rank, - head_dim).transpose(1, 2), - q_rope_cos, - q_rope_sin, - unsqueeze_dim=1, - is_neox=self._is_neox) + Q = _rope( + q_for_rope.reshape(B, block_size, num_heads_per_rank, head_dim).transpose(1, 2), + q_rope_cos, + q_rope_sin, + unsqueeze_dim=1, + is_neox=self._is_neox, + ) + k_noise_rope = _rope( + k_noise_for_rope.reshape( + B, block_size, num_kv_heads_per_rank, head_dim + ).transpose(1, 2), + q_rope_cos, + q_rope_sin, + unsqueeze_dim=1, + is_neox=self._is_neox, + ) Q_bshd = Q.transpose(1, 2) k_noise_bshd = k_noise_rope.transpose(1, 2) - v_noise_bshd = v_noise_all.reshape(B, block_size, - num_kv_heads_per_rank, - head_dim) + v_noise_bshd = v_noise_all.reshape(B, block_size, num_kv_heads_per_rank, head_dim) # Per-layer view into the pooled ctx cache. causal, window_size = self._get_attention_mask_args(layer_idx) - dspark_window = (self._dspark_layer_windows[layer_idx] if layer_idx - < len(self._dspark_layer_windows) else (-1, -1)) + dspark_window = ( + self._dspark_layer_windows[layer_idx] + if layer_idx < len(self._dspark_layer_windows) + else (-1, -1) + ) if dspark_window != (-1, -1): window_size = dspark_window - if self.dflash_attention_backend == 'TRTLLM': + if self.dflash_attention_backend == "TRTLLM": layer_cache = ctx_kv_cache[layer_idx] trtllm_gen_ops.append_paged_kv_cache( - append_key=k_noise_bshd.reshape(-1, num_kv_heads_per_rank, - head_dim).contiguous(), - append_value=v_noise_bshd.reshape(-1, num_kv_heads_per_rank, - head_dim).contiguous(), + append_key=k_noise_bshd.reshape( + -1, num_kv_heads_per_rank, head_dim + ).contiguous(), + append_value=v_noise_bshd.reshape( + -1, num_kv_heads_per_rank, head_dim + ).contiguous(), batch_indices=append_batch_indices, positions=append_positions, paged_kv_cache=layer_cache, @@ -1156,8 +1170,7 @@ def dflash_forward( kv_cache_sf=None, uses_shared_paged_kv_idx=True, bmm1_scale_log2=None, - multi_ctas_kv_counter_buffer=self. - _dflash_trtllm_gen_counters, + multi_ctas_kv_counter_buffer=self._dflash_trtllm_gen_counters, ) else: cum_seq_lens_q = torch.arange( @@ -1167,12 +1180,12 @@ def dflash_forward( dtype=torch.int32, device=hidden_states.device, ) - cum_seq_lens_kv = torch.cat(( - torch.zeros(1, - dtype=torch.int32, - device=hidden_states.device), - seq_lens_after.cumsum(0, dtype=torch.int32), - )) + cum_seq_lens_kv = torch.cat( + ( + torch.zeros(1, dtype=torch.int32, device=hidden_states.device), + seq_lens_after.cumsum(0, dtype=torch.int32), + ) + ) trtllm_gen_ops.batch_context_with_kv_cache( query=q_flat, kv_cache=(layer_cache[:, 0], layer_cache[:, 1]), @@ -1194,8 +1207,7 @@ def dflash_forward( kv_cache_sf=None, uses_shared_paged_kv_idx=True, causal=False, - multi_ctas_kv_counter_buffer=self. - _dflash_trtllm_gen_counters, + multi_ctas_kv_counter_buffer=self._dflash_trtllm_gen_counters, ) else: # VANILLA, validated before entering the layer loop. layer_k_cache = ctx_k_cache[:, layer_idx] @@ -1208,15 +1220,15 @@ def dflash_forward( # Exact only while every row of the block attends to the same # key set, i.e. non-causal, unwindowed layers. Causal or # windowed layers mask by row, so they stay unpacked. - pack_gqa = (gqa_group_size > 1 and not causal - and window_size == (-1, -1)) + pack_gqa = gqa_group_size > 1 and not causal and window_size == (-1, -1) if pack_gqa: - q_grouped = Q_bshd.reshape(B, block_size, - num_kv_heads_per_rank, - gqa_group_size, head_dim) + q_grouped = Q_bshd.reshape( + B, block_size, num_kv_heads_per_rank, gqa_group_size, head_dim + ) q_packed = q_grouped.permute(0, 3, 1, 2, 4) - q_in = q_packed.reshape(B, gqa_group_size * block_size, - num_kv_heads_per_rank, head_dim) + q_in = q_packed.reshape( + B, gqa_group_size * block_size, num_kv_heads_per_rank, head_dim + ) else: q_in = Q_bshd out = flash_attention( @@ -1232,27 +1244,25 @@ def dflash_forward( ) if pack_gqa: # Undo the packing: [B, group*blk, h_kv, d] -> [B, blk, h_q, d]. - out = out.view(B, gqa_group_size, block_size, - num_kv_heads_per_rank, - head_dim).permute(0, 2, 3, 1, 4) + out = out.view( + B, gqa_group_size, block_size, num_kv_heads_per_rank, head_dim + ).permute(0, 2, 3, 1, 4) attn_output = out.reshape(B * block_size, q_size) # Per-drafter post-attention gate (no-op for generic DFlash; Laguna # applies per-head softplus g_proj gating). gate input is the # input_layernorm output (the attention input). - attn_output = self._post_attention_gate(attn_output, hs_normed_flat, - attn_mod, - num_heads_per_rank, - head_dim) + attn_output = self._post_attention_gate( + attn_output, hs_normed_flat, attn_mod, num_heads_per_rank, head_dim + ) # o_proj (flat 2D, handles all-reduce internally) hidden_out = attn_mod.o_proj(attn_output) # Post-attention layernorm + MLP (flat 2D) res_flat = residual.reshape(-1, residual.shape[-1]) - hidden_out, res_flat = layer.post_attention_layernorm( - hidden_out, res_flat) + hidden_out, res_flat = layer.post_attention_layernorm(hidden_out, res_flat) hidden_out = layer.mlp(hidden_out) hidden_states = hidden_out.reshape(B, block_size, -1) @@ -1261,7 +1271,8 @@ def dflash_forward( # Final norm hidden_states_out, _ = self.model.norm( hidden_states.reshape(-1, hidden_states.shape[-1]), - residual.reshape(-1, residual.shape[-1])) + residual.reshape(-1, residual.shape[-1]), + ) return hidden_states_out def forward( @@ -1309,10 +1320,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 = "VANILLA"): """Pin the Laguna draft-layer class and enable Laguna-specific behaviors (context input_layernorm, causal sliding blocks); reject non-per-head gating.""" @@ -1326,32 +1334,29 @@ def __init__(self, ) self._context_input_layernorm = True self._sliding_layers_causal = True - gating = getattr(self.config, 'gating', True) - if gating not in (True, 'per-head'): + gating = getattr(self.config, "gating", True) + if gating not in (True, "per-head"): raise NotImplementedError( - f"Laguna DFlash drafter supports per-head gating only, " - f"got gating={gating!r}") + f"Laguna DFlash drafter supports per-head gating only, got gating={gating!r}" + ) def load_weights(self, weights, weight_mapper=None, **kwargs): """Build the per-aux ``fc_norm`` from the drafter's ``aux_hidden_norms.*`` weights, then defer the remaining weights to the base loader.""" aux_keys = sorted( - (k for k in weights if k.startswith('aux_hidden_norms.')), - key=lambda k: int(k.split('.')[1])) + (k for k in weights if k.startswith("aux_hidden_norms.")), + key=lambda k: int(k.split(".")[1]), + ) if not aux_keys: - raise ValueError( - "Laguna DFlash checkpoint is missing aux_hidden_norms.* weights" - ) + raise ValueError("Laguna DFlash checkpoint is missing aux_hidden_norms.* weights") weights = dict(weights) - eps = getattr(self.config, 'rms_norm_eps', 1e-6) + eps = getattr(self.config, "rms_norm_eps", 1e-6) norms = [] for k in aux_keys: w = weights.pop(k) - norm = nn.RMSNorm(w.shape[0], - eps=eps, - device='cuda', - elementwise_affine=True, - dtype=w.dtype) + norm = nn.RMSNorm( + w.shape[0], eps=eps, device="cuda", elementwise_affine=True, dtype=w.dtype + ) norm.weight.data.copy_(w) norms.append(norm) self.fc_norm = nn.ModuleList(norms) @@ -1361,23 +1366,20 @@ def project_target_hidden(self, hidden_states): """Project captured target features to the draft width: apply the per-aux ``fc_norm`` to each hidden chunk, then ``fc`` + ``hidden_norm``.""" hidden_states = hidden_states.to(self.fc.weight.dtype) - fc_norm = getattr(self, 'fc_norm', None) + fc_norm = getattr(self, "fc_norm", None) if fc_norm is not None: chunks = hidden_states.chunk(len(fc_norm), dim=-1) - hidden_states = torch.cat( - [norm(chunk) for norm, chunk in zip(fc_norm, chunks)], dim=-1) + hidden_states = torch.cat([norm(chunk) for norm, chunk in zip(fc_norm, chunks)], dim=-1) return self.hidden_norm(self.fc(hidden_states)) - def _post_attention_gate(self, attn_output, gate_input, attn_mod, num_heads, - head_dim): + def _post_attention_gate(self, attn_output, gate_input, attn_mod, num_heads, head_dim): """Apply Laguna's per-head softplus output gate (``g_proj``) to the attention output; a no-op when the layer has no ``g_proj``.""" - g_proj = getattr(attn_mod, 'g_proj', None) + g_proj = getattr(attn_mod, "g_proj", None) if g_proj is None: return attn_output gate = F.softplus(g_proj(gate_input).float()).to(attn_output.dtype) - return (attn_output.unflatten(-1, (num_heads, head_dim)) * - gate.unsqueeze(-1)).flatten(-2) + return (attn_output.unflatten(-1, (num_heads, head_dim)) * gate.unsqueeze(-1)).flatten(-2) @register_draft_model(SpeculativeDecodingMode.DFLASH) @@ -1387,8 +1389,7 @@ def _build_dflash_draft(model_config, draft_config, lm_head, model): Selects the Laguna variant by detecting its architecture in the draft checkpoint's own config. """ - draft_arches = getattr(draft_config.pretrained_config, "architectures", - None) or [] + draft_arches = getattr(draft_config.pretrained_config, "architectures", None) or [] dflash_attention_backend = model_config.spec_config.attention_backend if any("Laguna" in arch for arch in draft_arches): return DFlashLagunaForCausalLM( From 7b3ff14602fc274913e947e6ffed66e42524bf01 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Tue, 18 Aug 2026 23:26:38 -0700 Subject: [PATCH 05/21] [None][refactor] Merge the DSpark draft heads into one implementation The Markov chain loop and its bigram bias existed twice, with no import between them: modeling_dflash.py carried a raw-tensor pair for the DFlash drafter, models/dspark/heads.py an nn.Module tree for the V4-Pro one. Both now share dspark_markov_chain in modeling_speculative.py, the layer both drafters already sit above. Two call-site differences are preserved as parameters rather than folded away: the DFlash drafter needs a TP vocab shard plus a shard-aware argmax, and it casts the bias down to the logits dtype. The V4-Pro drafter builds its heads without a dtype argument, so its Markov weights stay fp32 while its logits are bf16; adding the cast unconditionally would have narrowed its accumulation to bf16 silently. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/models/dspark/draft.py | 4 +- tensorrt_llm/_torch/models/dspark/heads.py | 254 ----------- tensorrt_llm/_torch/models/modeling_dflash.py | 66 +-- tensorrt_llm/_torch/models/modeling_dspark.py | 2 +- .../_torch/models/modeling_speculative.py | 401 ++++++++++++++++++ .../hw_agnostic/test_dspark_draft.py | 2 +- .../hw_agnostic/test_dspark_heads.py | 2 +- .../test_kimi_k3_dspark_semantics.py | 5 +- 8 files changed, 410 insertions(+), 326 deletions(-) delete mode 100644 tensorrt_llm/_torch/models/dspark/heads.py diff --git a/tensorrt_llm/_torch/models/dspark/draft.py b/tensorrt_llm/_torch/models/dspark/draft.py index 1b47f4922965..99a9021ca2c7 100644 --- a/tensorrt_llm/_torch/models/dspark/draft.py +++ b/tensorrt_llm/_torch/models/dspark/draft.py @@ -35,7 +35,7 @@ import torch from torch import nn -from .heads import confident_prefix_length +from ..modeling_speculative import confident_prefix_length def build_draft_input_ids( @@ -94,7 +94,7 @@ def dspark_propose( ) draft_logits = corrected else: - from .heads import greedy_or_sample + from ..modeling_speculative import greedy_or_sample draft_tokens = greedy_or_sample(base_logits, temperature) diff --git a/tensorrt_llm/_torch/models/dspark/heads.py b/tensorrt_llm/_torch/models/dspark/heads.py deleted file mode 100644 index c49e35fbafa0..000000000000 --- a/tensorrt_llm/_torch/models/dspark/heads.py +++ /dev/null @@ -1,254 +0,0 @@ -# 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. -# -# The DSpark Markov/RNN/confidence-head math is ported from DeepSeek's DeepSpec -# reference implementation (https://github.com/deepseek-ai/DeepSpec, MIT License). -"""DSpark draft-network heads (pure-torch, framework-agnostic). - -These modules implement the *sequential refinement* and *acceptance-confidence* -parts of DeepSeek's DSpark speculative-decoding draft network: - - - Markov head: a low-rank token-bigram logit bias ``logits_k += W2(W1[t_{k-1}])`` - applied autoregressively across the ``block_size`` draft positions (the cheap - "sequential" half of DSpark's "semi-parallel" drafting). RNN variant carries - a GRU-style recurrent state across positions. - - Confidence head: predicts a per-position acceptance probability; the cumulative - product over positions estimates prefix-acceptance and is used only to - *truncate* the proposed draft length (NOT to decide acceptance). - -This file deliberately depends on ``torch`` only so it can be unit-tested in -isolation (token-for-token) against the DeepSpec reference. -""" - -from typing import Optional - -import torch -from torch import nn - - -def greedy_or_sample(logits: torch.Tensor, temperature: float) -> torch.Tensor: - """Argmax for temperature<=0, else temperature-scaled multinomial. - - Args: - logits: ``[..., vocab]``. - Returns: - token ids with the trailing vocab dim reduced. - """ - if temperature <= 0.0: - return logits.argmax(dim=-1) - probs = torch.softmax(logits.float() / temperature, dim=-1) - flat = probs.reshape(-1, probs.shape[-1]) - sampled = torch.multinomial(flat, num_samples=1).squeeze(-1) - return sampled.view(probs.shape[:-1]) - - -class VanillaMarkov(nn.Module): - """Low-rank token-bigram logit bias: ``bias = W2(W1[token])``.""" - - markov_head_type = "vanilla" - - def __init__(self, *, vocab_size: int, markov_rank: int): - super().__init__() - self.vocab_size = int(vocab_size) - self.markov_rank = int(markov_rank) - assert self.markov_rank > 0, ( - f"VanillaMarkov requires markov_rank > 0, got {self.markov_rank}." - ) - self.markov_w1 = nn.Embedding(self.vocab_size, self.markov_rank) - self.markov_w2 = nn.Linear(self.markov_rank, self.vocab_size, bias=False) - - def get_prev_embeddings(self, token_ids: torch.Tensor) -> torch.Tensor: - return self.markov_w1(token_ids.long()) - - def project_bias(self, latent_states: torch.Tensor) -> torch.Tensor: - return self.markov_w2(latent_states) - - def compute_step_bias( - self, token_ids: torch.Tensor, hidden_states: Optional[torch.Tensor] - ) -> torch.Tensor: - del hidden_states - return self.project_bias(self.get_prev_embeddings(token_ids)) - - def apply_step_logits( - self, - logits: torch.Tensor, - *, - token_ids: torch.Tensor, - hidden_states: Optional[torch.Tensor], - ) -> torch.Tensor: - return logits + self.compute_step_bias(token_ids, hidden_states) - - def sample_block_tokens( - self, - base_logits: torch.Tensor, - *, - first_prev_token_ids: torch.Tensor, - hidden_states: Optional[torch.Tensor], - temperature: float = 0.0, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Autoregressive block sampling with the (memoryless) Markov bias. - - Args: - base_logits: ``[batch, block_size, vocab]`` from the backbone+lm_head. - first_prev_token_ids: ``[batch]`` token preceding the first position. - hidden_states: ``[batch, block_size, d]`` (unused by vanilla/gated). - Returns: - sampled_tokens ``[batch, block_size]``, corrected_logits ``[batch, block_size, vocab]``. - """ - batch_size, block_size = base_logits.shape[:2] - if block_size == 0: - empty = torch.empty(batch_size, 0, dtype=torch.long, device=base_logits.device) - return empty, base_logits - sampled, corrected = [], [] - prev = first_prev_token_ids.long() - for k in range(block_size): - step_hidden = None if hidden_states is None else hidden_states[:, k] - step_logits = self.apply_step_logits( - base_logits[:, k], token_ids=prev, hidden_states=step_hidden - ) - corrected.append(step_logits.unsqueeze(1)) - prev = greedy_or_sample(step_logits, temperature) - sampled.append(prev) - return torch.stack(sampled, dim=1), torch.cat(corrected, dim=1) - - -class GatedMarkovHead(VanillaMarkov): - """Markov bias gated by a sigmoid of [hidden, prev_embedding].""" - - markov_head_type = "gated" - - def __init__(self, *, vocab_size: int, markov_rank: int, hidden_size: int): - super().__init__(vocab_size=vocab_size, markov_rank=markov_rank) - self.gate_proj = nn.Linear(hidden_size + markov_rank, markov_rank) - - def compute_step_bias( - self, token_ids: torch.Tensor, hidden_states: Optional[torch.Tensor] - ) -> torch.Tensor: - assert hidden_states is not None - prev_emb = self.get_prev_embeddings(token_ids) - gate = torch.sigmoid(self.gate_proj(torch.cat([hidden_states, prev_emb], dim=-1))).to( - dtype=prev_emb.dtype - ) - return self.project_bias(gate * prev_emb) - - -class RNNHead(VanillaMarkov): - """GRU-style head carrying recurrent state across block positions.""" - - markov_head_type = "rnn" - - def __init__(self, *, vocab_size: int, markov_rank: int, hidden_size: int): - super().__init__(vocab_size=vocab_size, markov_rank=markov_rank) - self.hidden_size = int(hidden_size) - # [s_{k-1}; W1[x_{k-1}]; h_k] -> [gate; candidate; output] - self.joint_proj = nn.Linear(2 * markov_rank + hidden_size, 3 * markov_rank) - - def _rnn_step(self, state, prev_embeddings, hidden_states): - z = torch.cat([state, prev_embeddings, hidden_states], dim=-1) - gate_raw, cand_raw, out_raw = self.joint_proj(z).chunk(3, dim=-1) - gate = torch.sigmoid(gate_raw) - candidate = torch.tanh(cand_raw) - new_state = gate * state + (1.0 - gate) * candidate - bias = self.project_bias(torch.tanh(out_raw)) - return new_state, bias - - def sample_block_tokens( - self, - base_logits: torch.Tensor, - *, - first_prev_token_ids: torch.Tensor, - hidden_states: Optional[torch.Tensor], - temperature: float = 0.0, - ) -> tuple[torch.Tensor, torch.Tensor]: - assert hidden_states is not None - batch_size, block_size = base_logits.shape[:2] - if block_size == 0: - empty = torch.empty(batch_size, 0, dtype=torch.long, device=base_logits.device) - return empty, base_logits - state = torch.zeros( - batch_size, self.markov_rank, device=base_logits.device, dtype=hidden_states.dtype - ) - sampled, corrected = [], [] - prev = first_prev_token_ids.long() - for k in range(block_size): - prev_emb = self.get_prev_embeddings(prev) - state, bias = self._rnn_step(state, prev_emb, hidden_states[:, k]) - step_logits = base_logits[:, k] + bias - corrected.append(step_logits.unsqueeze(1)) - prev = greedy_or_sample(step_logits, temperature) - sampled.append(prev) - return torch.stack(sampled, dim=1), torch.cat(corrected, dim=1) - - -def build_markov_head( - *, markov_head_type: str, vocab_size: int, markov_rank: int, hidden_size: int -) -> Optional[nn.Module]: - """Factory mirroring DeepSpec ``build_markov_head``; returns None if rank==0.""" - if int(markov_rank) <= 0: - return None - kind = str(markov_head_type).lower() - if kind == "vanilla": - return VanillaMarkov(vocab_size=vocab_size, markov_rank=markov_rank) - if kind == "gated": - return GatedMarkovHead( - vocab_size=vocab_size, markov_rank=markov_rank, hidden_size=hidden_size - ) - if kind == "rnn": - return RNNHead(vocab_size=vocab_size, markov_rank=markov_rank, hidden_size=hidden_size) - raise ValueError(f"Unsupported markov_head_type: {markov_head_type!r}") - - -class DSparkConfidenceHead(nn.Module): - """Per-position acceptance-confidence predictor (DeepSpec AcceptRatePredictor). - - Input features are the backbone hidden state, optionally concatenated with the - Markov head's previous-token embedding. Output is a single logit per position. - """ - - def __init__(self, *, hidden_size: int, markov_rank: int = 0, with_markov: bool = False): - super().__init__() - self.with_markov = bool(with_markov) - input_dim = int(hidden_size) + (int(markov_rank) if with_markov else 0) - # The checkpoint stores ``proj`` as a bias-free bf16 weight, but the - # confidence score is computed in fp32 (mirrors the DeepSpec reference - # ``Linear(input_dim, 1, dtype=torch.float32)`` with the fp32 matmul). - self.proj = nn.Linear(input_dim, 1, bias=False, dtype=torch.float32) - - def forward( - self, hidden_states: torch.Tensor, prev_embeddings: Optional[torch.Tensor] = None - ) -> torch.Tensor: - if self.with_markov: - assert prev_embeddings is not None - features = torch.cat([hidden_states, prev_embeddings.to(hidden_states.dtype)], dim=-1) - else: - features = hidden_states - # fp32 matmul for a stable confidence score (mirrors the reference). - return self.proj(features.float()).squeeze(-1) - - -def confident_prefix_length( - confidence_logits: torch.Tensor, *, block_size: int, threshold: float -) -> int: - """First position k where ``sigmoid(confidence_k) < threshold``. - - Returns ``block_size`` when threshold<=0 (no truncation) or all positions - are confident. Assumes batch size 1 (functional-first scope). - """ - if threshold <= 0.0: - return int(block_size) - below = confidence_logits.sigmoid() < threshold - if not bool(below[0].any().item()): - return int(block_size) - return int(torch.nonzero(below[0], as_tuple=False)[0].item()) diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index c3e4152904dc..79f65a679647 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -21,6 +21,7 @@ from ..pyexecutor.config_utils import _is_sliding_attention_layer, get_layer_attention_window from ..speculative.dflash_attention import get_dflash_flash_attention, get_dflash_trtllm_gen_ops from ..speculative.interface import SpeculativeDecodingMode +from .modeling_speculative import dspark_markov_chain_logits from .modeling_utils import get_model_architecture, register_draft_model @@ -51,69 +52,6 @@ def dspark_layer_window_size( return (swa_window - 1, swa_window - 1) -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. - - Reference: DeepSpec ``VanillaMarkov`` (deepspec/modeling/dspark/ - markov_head.py): ``bias = markov_w2(markov_w1(prev_token))`` where - markov_w1 is nn.Embedding(vocab, rank) and markov_w2 is - nn.Linear(rank, vocab, bias=False). With both weights stored - [vocab, rank] this is ``markov_w1[prev] @ markov_w2.T``. - - Args: - prev_tokens: [B] long, previous token per request (draft vocab). - markov_w1: [vocab, rank]. - markov_w2: [vocab_or_shard, rank] (rows may be a TP vocab shard). - Returns: - [B, vocab_or_shard] bias in the markov weights' dtype. - """ - return F.linear(F.embedding(prev_tokens, markov_w1), markov_w2) - - -def dspark_markov_chain_logits( - base_logits: torch.Tensor, - first_prev_tokens: torch.Tensor, - markov_w1: torch.Tensor, - markov_w2: torch.Tensor, - argmax_fn=None, -) -> torch.Tensor: - """Apply the vanilla Markov intra-block bias across a drafted block. - - Reference: DeepSpec ``VanillaMarkov.sample_block_tokens`` at - temperature 0: for step i, ``logits_i += bias(prev_i)`` with - ``prev_0`` = the anchor token (last accepted token, block slot 0) and - ``prev_{i>0}`` = the greedy token from step i-1's *biased* logits. - llama.cpp PR #25173 implements the same greedy chain. - - Args: - base_logits: [B, K, vocab_or_shard] shared-lm_head logits. - first_prev_tokens: [B] long, anchor token ids (draft vocab). - markov_w1 / markov_w2: see :func:`dspark_markov_step_bias`. - argmax_fn: callable([B, vocab_or_shard]) -> [B] token ids in the - full draft vocab; defaults to plain argmax. Workers pass a - TP-aware argmax when the draft logits are vocab-sharded. - Returns: - [B, K, vocab_or_shard] biased logits. Greedy per-position argmax of - the result reproduces the reference sampled chain exactly. - """ - K = base_logits.shape[1] - if K == 0: - return base_logits - prev = first_prev_tokens.long() - steps = [] - for i in range(K): - bias = dspark_markov_step_bias(prev, markov_w1, markov_w2) - step_logits = base_logits[:, i] + bias.to(base_logits.dtype) - steps.append(step_logits) - if argmax_fn is not None: - prev = argmax_fn(step_logits).long() - else: - prev = torch.argmax(step_logits, dim=-1) - return torch.stack(steps, dim=1) - - class DFlashForCausalLM(nn.Module): """Draft model wrapper for DFlash speculative decoding. @@ -395,7 +333,7 @@ def apply_markov_chain_logits( """Apply the dspark vanilla-Markov intra-block bias to block logits. No-op (returns ``base_logits`` unchanged) for non-dspark drafters. - See :func:`dspark_markov_chain_logits` for the semantics; when + See :func:`dspark_markov_chain` for the semantics; when ``base_logits`` is a TP vocab shard, the caller must pass this rank's ``vocab_slice`` (to shard the markov_w2 rows identically) and an ``argmax_fn`` returning full-vocab token ids — DFlashWorker diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 3281336e1de6..8bbd0e2f2667 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -62,7 +62,6 @@ precompute_dspark_freqs_cis, ) from .dspark.draft import build_draft_input_ids, dspark_propose -from .dspark.heads import DSparkConfidenceHead, build_markov_head from .modeling_deepseekv4 import ( DeepseekV4DecoderLayer, DeepseekV4WeightLoader, @@ -72,6 +71,7 @@ _rename_deepseek_v4_attn_subkey, _rename_deepseek_v4_ffn_subkey, ) +from .modeling_speculative import DSparkConfidenceHead, build_markov_head from .modeling_utils import register_draft_model # Matches the draft namespace ``mtp..`` in the V4-Pro-DSpark diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index eaf6605279e4..33f6ded534b9 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -6,6 +6,7 @@ from typing import Dict, Generic, List, Optional, Tuple import torch +import torch.nn.functional as F from torch import nn from transformers import LlamaConfig, PretrainedConfig @@ -59,6 +60,406 @@ def _slice_spec_position_ids(position_ids: Optional[torch.Tensor], return position_ids[..., :num_tokens] +# --------------------------------------------------------------------------- +# DSpark draft-network heads, shared by both drafters that implement DSpark: +# the DFlash path (modeling_dflash.py, standalone Kimi K3 style drafters) and +# the DeepSeek-V4-Pro path (modeling_dspark.py, mtp.* stages inside the target +# checkpoint). DFlash is the degenerate case -- DSpark with the Markov and +# confidence heads switched off -- so the math below is the *only* copy; the +# two drafters differ solely in how they store the weights and whether their +# draft lm_head is TP vocab-sharded. +# +# Ported from DeepSeek's DeepSpec reference implementation +# (https://github.com/deepseek-ai/DeepSpec, MIT License). +# --------------------------------------------------------------------------- + + +def greedy_or_sample(logits: torch.Tensor, temperature: float) -> torch.Tensor: + """Argmax for temperature<=0, else temperature-scaled multinomial. + + Args: + logits: ``[..., vocab]``. + Returns: + token ids with the trailing vocab dim reduced. + """ + if temperature <= 0.0: + return logits.argmax(dim=-1) + probs = torch.softmax(logits.float() / temperature, dim=-1) + flat = probs.reshape(-1, probs.shape[-1]) + sampled = torch.multinomial(flat, num_samples=1).squeeze(-1) + return sampled.view(probs.shape[:-1]) + + +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. + + Reference: DeepSpec ``VanillaMarkov`` (deepspec/modeling/dspark/ + markov_head.py): ``bias = markov_w2(markov_w1(prev_token))`` where + markov_w1 is nn.Embedding(vocab, rank) and markov_w2 is + nn.Linear(rank, vocab, bias=False). With both weights stored + [vocab, rank] this is ``markov_w1[prev] @ markov_w2.T``. + + Args: + prev_tokens: [B] long, previous token per request (draft vocab). + markov_w1: [vocab, rank]. + markov_w2: [vocab_or_shard, rank] (rows may be a TP vocab shard). + Returns: + [B, vocab_or_shard] bias in the markov weights' dtype. + """ + return F.linear(F.embedding(prev_tokens, markov_w1), markov_w2) + + +def dspark_markov_chain( + base_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + step_bias_fn, + *, + hidden_states: Optional[torch.Tensor] = None, + next_token_fn=None, + cast_bias_to_logits: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """The intra-block Markov refinement loop, shared by every DSpark head. + + Reference: DeepSpec ``VanillaMarkov.sample_block_tokens``. For step i, + ``logits_i += bias(prev_i)`` with ``prev_0`` = the anchor token (the last + accepted token, block slot 0) and ``prev_{i>0}`` = the token drawn from + step i-1's *biased* logits. llama.cpp PR #25173 implements the same chain. + + The chain being sequential is load-bearing and is NOT derivable from a + checkpoint: published DSpark drafters ship only the training-time + ``apply_block_logits``, which biases the whole block in one teacher-forced + pass. SGLang ``srt/models/dspark.py:34-64 run_markov_block`` settles it -- + it steps the same way, feeding back the token drawn from the biased logits. + + Two callers drive this, and a change here has to satisfy both: the embedded + flavour chains it inside the model, which owns its sampler, while the + standalone flavour is driven from the worker through ``next_token_fn`` so + the chain can advance on a TP-gathered global argmax the model cannot + compute on its own. + + Args: + base_logits: [B, K, vocab_or_shard] shared-lm_head logits. + first_prev_tokens: [B] long, anchor token ids (draft vocab). + step_bias_fn: ``(prev_tokens [B], step_hidden or None) -> bias``. A + closure, so a stateful head (RNN) can carry its recurrent state + across positions without a second loop. + hidden_states: [B, K, d] fed to ``step_bias_fn`` one position at a + time; None for the memoryless heads. + next_token_fn: ``([B, vocab_or_shard]) -> [B]`` token ids in the FULL + draft vocab; defaults to a plain argmax. Drafters whose draft + logits are TP vocab-sharded pass a shard-aware argmax here. + cast_bias_to_logits: cast the bias down to ``base_logits.dtype`` + before adding. The DFlash drafter does; the V4-Pro drafter does + not (its Markov weights and its logits already agree), and adding + the cast there would silently narrow its accumulation dtype. + Returns: + sampled_tokens [B, K], corrected_logits [B, K, vocab_or_shard]. + Greedy per-position argmax of the corrected logits reproduces the + reference sampled chain exactly. + """ + batch_size, block_size = base_logits.shape[:2] + if block_size == 0: + empty = torch.empty(batch_size, + 0, + dtype=torch.long, + device=base_logits.device) + return empty, base_logits + sampled, corrected = [], [] + prev = first_prev_tokens.long() + for k in range(block_size): + step_hidden = None if hidden_states is None else hidden_states[:, k] + bias = step_bias_fn(prev, step_hidden) + if cast_bias_to_logits: + bias = bias.to(base_logits.dtype) + step_logits = base_logits[:, k] + bias + corrected.append(step_logits.unsqueeze(1)) + if next_token_fn is None: + prev = torch.argmax(step_logits, dim=-1) + else: + prev = next_token_fn(step_logits).long() + sampled.append(prev) + return torch.stack(sampled, dim=1), torch.cat(corrected, dim=1) + + +def dspark_markov_chain_logits(base_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + markov_w1: torch.Tensor, + markov_w2: torch.Tensor, + argmax_fn=None) -> torch.Tensor: + """Raw-tensor entry to :func:`dspark_markov_chain`, corrected logits only. + + For drafters that keep the Markov head as plain checkpoint tensors rather + than a :class:`VanillaMarkov` module (the DFlash path). ``markov_w2`` may + already be sliced down to this rank's TP vocab shard, in which case + ``argmax_fn`` must map a shard-local row back to a full-vocab token id. + """ + + def _step_bias(prev_tokens, step_hidden): + del step_hidden + return dspark_markov_step_bias(prev_tokens, markov_w1, markov_w2) + + _, corrected = dspark_markov_chain(base_logits, + first_prev_tokens, + _step_bias, + next_token_fn=argmax_fn, + cast_bias_to_logits=True) + return corrected + + +class VanillaMarkov(nn.Module): + """Low-rank token-bigram logit bias: ``bias = W2(W1[token])``.""" + + markov_head_type = "vanilla" + + def __init__(self, *, vocab_size: int, markov_rank: int): + super().__init__() + self.vocab_size = int(vocab_size) + self.markov_rank = int(markov_rank) + assert self.markov_rank > 0, ( + f"VanillaMarkov requires markov_rank > 0, got {self.markov_rank}.") + self.markov_w1 = nn.Embedding(self.vocab_size, self.markov_rank) + self.markov_w2 = nn.Linear(self.markov_rank, + self.vocab_size, + bias=False) + + def get_prev_embeddings(self, token_ids: torch.Tensor) -> torch.Tensor: + return F.embedding(token_ids.long(), self.markov_w1.weight) + + def project_bias(self, + latent_states: torch.Tensor, + *, + vocab_slice: Optional[slice] = None) -> torch.Tensor: + w2 = self.markov_w2.weight + if vocab_slice is not None: + w2 = w2[vocab_slice] + return F.linear(latent_states, w2) + + def compute_step_bias(self, + token_ids: torch.Tensor, + hidden_states: Optional[torch.Tensor], + *, + vocab_slice: Optional[slice] = None) -> torch.Tensor: + del hidden_states + w2 = self.markov_w2.weight + if vocab_slice is not None: + w2 = w2[vocab_slice] + return dspark_markov_step_bias(token_ids.long(), self.markov_w1.weight, + w2) + + def apply_step_logits( + self, + logits: torch.Tensor, + *, + token_ids: torch.Tensor, + hidden_states: Optional[torch.Tensor], + ) -> torch.Tensor: + return logits + self.compute_step_bias(token_ids, hidden_states) + + def sample_block_tokens( + self, + base_logits: torch.Tensor, + *, + first_prev_token_ids: torch.Tensor, + hidden_states: Optional[torch.Tensor], + temperature: float = 0.0, + vocab_slice: Optional[slice] = None, + next_token_fn=None, + cast_bias_to_logits: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Autoregressive block sampling with the (memoryless) Markov bias. + + Args: + base_logits: ``[batch, block_size, vocab]`` from backbone+lm_head. + first_prev_token_ids: ``[batch]`` token preceding the 1st position. + hidden_states: ``[batch, block_size, d]`` (unused by vanilla). + vocab_slice / next_token_fn / cast_bias_to_logits: see + :func:`dspark_markov_chain`; only the TP vocab-sharded DFlash + drafter sets them. + Returns: + sampled_tokens ``[batch, block_size]``, + corrected_logits ``[batch, block_size, vocab]``. + """ + + def _step_bias(prev_tokens, step_hidden): + return self.compute_step_bias(prev_tokens, + step_hidden, + vocab_slice=vocab_slice) + + def _sample_step(step_logits): + return greedy_or_sample(step_logits, temperature) + + return dspark_markov_chain( + base_logits, + first_prev_token_ids, + _step_bias, + hidden_states=hidden_states, + next_token_fn=_sample_step + if next_token_fn is None else next_token_fn, + cast_bias_to_logits=cast_bias_to_logits, + ) + + +class GatedMarkovHead(VanillaMarkov): + """Markov bias gated by a sigmoid of [hidden, prev_embedding].""" + + markov_head_type = "gated" + + def __init__(self, *, vocab_size: int, markov_rank: int, hidden_size: int): + super().__init__(vocab_size=vocab_size, markov_rank=markov_rank) + self.gate_proj = nn.Linear(hidden_size + markov_rank, markov_rank) + + def compute_step_bias(self, + token_ids: torch.Tensor, + hidden_states: Optional[torch.Tensor], + *, + vocab_slice: Optional[slice] = None) -> torch.Tensor: + assert hidden_states is not None + prev_emb = self.get_prev_embeddings(token_ids) + gate = torch.sigmoid( + self.gate_proj(torch.cat([hidden_states, prev_emb], + dim=-1))).to(dtype=prev_emb.dtype) + return self.project_bias(gate * prev_emb, vocab_slice=vocab_slice) + + +class RNNHead(VanillaMarkov): + """GRU-style head carrying recurrent state across block positions.""" + + markov_head_type = "rnn" + + def __init__(self, *, vocab_size: int, markov_rank: int, hidden_size: int): + super().__init__(vocab_size=vocab_size, markov_rank=markov_rank) + self.hidden_size = int(hidden_size) + # [s_{k-1}; W1[x_{k-1}]; h_k] -> [gate; candidate; output] + self.joint_proj = nn.Linear(2 * markov_rank + hidden_size, + 3 * markov_rank) + + def _rnn_step(self, + state, + prev_embeddings, + hidden_states, + *, + vocab_slice: Optional[slice] = None): + z = torch.cat([state, prev_embeddings, hidden_states], dim=-1) + gate_raw, cand_raw, out_raw = self.joint_proj(z).chunk(3, dim=-1) + gate = torch.sigmoid(gate_raw) + candidate = torch.tanh(cand_raw) + new_state = gate * state + (1.0 - gate) * candidate + bias = self.project_bias(torch.tanh(out_raw), vocab_slice=vocab_slice) + return new_state, bias + + def sample_block_tokens( + self, + base_logits: torch.Tensor, + *, + first_prev_token_ids: torch.Tensor, + hidden_states: Optional[torch.Tensor], + temperature: float = 0.0, + vocab_slice: Optional[slice] = None, + next_token_fn=None, + cast_bias_to_logits: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + assert hidden_states is not None + state = torch.zeros(base_logits.shape[0], + self.markov_rank, + device=base_logits.device, + dtype=hidden_states.dtype) + + def _step_bias(prev_tokens, step_hidden): + nonlocal state + prev_emb = self.get_prev_embeddings(prev_tokens) + state, bias = self._rnn_step(state, + prev_emb, + step_hidden, + vocab_slice=vocab_slice) + return bias + + def _sample_step(step_logits): + return greedy_or_sample(step_logits, temperature) + + return dspark_markov_chain( + base_logits, + first_prev_token_ids, + _step_bias, + hidden_states=hidden_states, + next_token_fn=_sample_step + if next_token_fn is None else next_token_fn, + cast_bias_to_logits=cast_bias_to_logits, + ) + + +def build_markov_head(*, markov_head_type: str, vocab_size: int, + markov_rank: int, + hidden_size: int) -> Optional[nn.Module]: + """Factory mirroring DeepSpec ``build_markov_head``; None if rank==0.""" + if int(markov_rank) <= 0: + return None + kind = str(markov_head_type).lower() + if kind == "vanilla": + return VanillaMarkov(vocab_size=vocab_size, markov_rank=markov_rank) + if kind == "gated": + return GatedMarkovHead(vocab_size=vocab_size, + markov_rank=markov_rank, + hidden_size=hidden_size) + if kind == "rnn": + return RNNHead(vocab_size=vocab_size, + markov_rank=markov_rank, + hidden_size=hidden_size) + raise ValueError(f"Unsupported markov_head_type: {markov_head_type!r}") + + +class DSparkConfidenceHead(nn.Module): + """Per-position acceptance-confidence predictor (DeepSpec + AcceptRatePredictor). + + Input features are the backbone hidden state, optionally concatenated with + the Markov head's previous-token embedding. Output is a single logit per + position. + """ + + def __init__(self, + *, + hidden_size: int, + markov_rank: int = 0, + with_markov: bool = False): + super().__init__() + self.with_markov = bool(with_markov) + input_dim = int(hidden_size) + (int(markov_rank) if with_markov else 0) + # The checkpoint stores ``proj`` as a bias-free bf16 weight, but the + # confidence score is computed in fp32 (mirrors the DeepSpec reference + # ``Linear(input_dim, 1, dtype=torch.float32)`` with the fp32 matmul). + self.proj = nn.Linear(input_dim, 1, bias=False, dtype=torch.float32) + + def forward(self, + hidden_states: torch.Tensor, + prev_embeddings: Optional[torch.Tensor] = None) -> torch.Tensor: + if self.with_markov: + assert prev_embeddings is not None + features = torch.cat( + [hidden_states, + prev_embeddings.to(hidden_states.dtype)], + dim=-1) + else: + features = hidden_states + # fp32 matmul for a stable confidence score (mirrors the reference). + return self.proj(features.float()).squeeze(-1) + + +def confident_prefix_length(confidence_logits: torch.Tensor, *, block_size: int, + threshold: float) -> int: + """First position k where ``sigmoid(confidence_k) < threshold``. + + Returns ``block_size`` when threshold<=0 (no truncation) or all positions + are confident. Assumes batch size 1 (functional-first scope). + """ + if threshold <= 0.0: + return int(block_size) + below = confidence_logits.sigmoid() < threshold + if not bool(below[0].any().item()): + return int(block_size) + return int(torch.nonzero(below[0], as_tuple=False)[0].item()) + + class Eagle3Attention(Attention): def __init__( diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py index fed0da3a2803..f31dab6cb12a 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py @@ -17,7 +17,7 @@ import torch from tensorrt_llm._torch.models.dspark.draft import build_draft_input_ids, dspark_propose -from tensorrt_llm._torch.models.dspark.heads import DSparkConfidenceHead, build_markov_head +from tensorrt_llm._torch.models.modeling_speculative import DSparkConfidenceHead, build_markov_head VOCAB, HID, RANK, B, BLK = 257, 32, 16, 2, 5 NOISE_ID = 199 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 16b6a01acd17..5546c1b03f85 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py @@ -17,7 +17,7 @@ import pytest import torch -from tensorrt_llm._torch.models.dspark.heads import ( +from tensorrt_llm._torch.models.modeling_speculative import ( DSparkConfidenceHead, RNNHead, VanillaMarkov, 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 61c1fda9432a..3768d4198efe 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 @@ -23,9 +23,8 @@ import torch import torch.nn.functional as F -from tensorrt_llm._torch.models.modeling_dflash import ( - DFlashForCausalLM, - dspark_layer_window_size, +from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM, dspark_layer_window_size +from tensorrt_llm._torch.models.modeling_speculative import ( dspark_markov_chain_logits, dspark_markov_step_bias, ) From 3846595da899905c6bd2c61ad773729dd9d5a149 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Wed, 19 Aug 2026 01:35:09 -0700 Subject: [PATCH 06/21] [None][refactor] Fold the DSpark package into modeling_dspark.py The models/dspark/ package held the captured-context attention primitives and the block draft I/O for one consumer, modeling_dspark.py. Fold both in and drop the package; the heads left in cut 2. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/models/dspark/__init__.py | 16 - .../_torch/models/dspark/attention.py | 567 -------------- tensorrt_llm/_torch/models/dspark/draft.py | 130 ---- tensorrt_llm/_torch/models/modeling_dspark.py | 711 +++++++++++++++++- .../hw_agnostic/test_dspark_attention.py | 10 +- .../hw_agnostic/test_dspark_cuda_graph.py | 2 +- .../hw_agnostic/test_dspark_draft.py | 2 +- .../test_dspark_cute_dsl_attention.py | 4 +- 8 files changed, 695 insertions(+), 747 deletions(-) delete mode 100644 tensorrt_llm/_torch/models/dspark/__init__.py delete mode 100644 tensorrt_llm/_torch/models/dspark/attention.py delete mode 100644 tensorrt_llm/_torch/models/dspark/draft.py diff --git a/tensorrt_llm/_torch/models/dspark/__init__.py b/tensorrt_llm/_torch/models/dspark/__init__.py deleted file mode 100644 index b33d7553d877..000000000000 --- a/tensorrt_llm/_torch/models/dspark/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -# 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. - -"""DSpark draft-model components.""" diff --git a/tensorrt_llm/_torch/models/dspark/attention.py b/tensorrt_llm/_torch/models/dspark/attention.py deleted file mode 100644 index f3cfcedfbef4..000000000000 --- a/tensorrt_llm/_torch/models/dspark/attention.py +++ /dev/null @@ -1,567 +0,0 @@ -# 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. -# -# The DSpark captured-context attention primitives are ported from DeepSeek's -# DeepSpec reference ``inference/kernel.py`` (``sparse_attn``) and -# ``inference/model.py`` (``get_dspark_topk_idxs``). The reference computes these -# with a TileLang kernel; this is a functional-first pure-PyTorch port with the -# same math (index-gather + online softmax + a learnable attention sink that -# contributes only to the softmax denominator). -"""DSpark draft captured-context attention primitives (hardware-agnostic). - -The DSpark draft uses *dense* sliding-window MLA (``compress_ratio == 0``): the -query comes from the block's draft tokens, while the keys/values are gathered -from a small per-request set of positions (a sliding window of the projected -captured context plus the current block's own positions). Two primitives capture -the parts that differ from the standard MLA path: - -* :func:`get_dspark_topk_idxs` — the (window-context + block) position list. -* :func:`dspark_sparse_attn` — index-gathered attention with an attention sink. -""" - -from functools import lru_cache - -import torch -import torch.nn.functional as F - -from ...._utils import is_sm_100f -from ...cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE - -if IS_CUTLASS_DSL_AVAILABLE: - from ...custom_ops.dspark_attention_custom_op import ( - cute_dsl_dspark_attention, - is_fused_dspark_attention_supported, - ) - from ...custom_ops.dspark_rmsnorm_rope_custom_op import ( - cute_dsl_dspark_rmsnorm_rope, - is_fused_dspark_rmsnorm_rope_supported, - ) - -__all__ = [ - "get_dspark_topk_idxs", - "get_dspark_topk_idxs_batched", - "dspark_sparse_attn", - "precompute_dspark_freqs_cis", - "apply_dspark_rotary", - "apply_dspark_rotary_batched", - "dspark_attention_forward", - "dspark_attention_forward_batched", -] - - -def precompute_dspark_freqs_cis( - rope_head_dim: int, - seqlen: int, - rope_theta: float = 10000.0, - device: torch.device | str = "cpu", -) -> torch.Tensor: - """Plain (non-YaRN) RoPE complex exponentials for the DSpark draft. - - The dense draft attention (``compress_ratio == 0``) disables YaRN and uses the - base ``rope_theta`` (DeepSpec ``precompute_freqs_cis`` with - ``original_seq_len == 0``). - - Returns: - complex64 tensor ``[seqlen, rope_head_dim // 2]``. - """ - freqs = 1.0 / ( - rope_theta - ** (torch.arange(0, rope_head_dim, 2, dtype=torch.float32, device=device) / rope_head_dim) - ) - t = torch.arange(seqlen, dtype=torch.float32, device=device) - freqs = torch.outer(t, freqs) - return torch.polar(torch.ones_like(freqs), freqs) - - -def apply_dspark_rotary( - x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False -) -> torch.Tensor: - """Apply (or, with ``inverse``, de-apply) rotary embeddings, DeepSpec-style. - - Functional (non-in-place) port of DeepSpec ``apply_rotary_emb``: treats the - last dim as adjacent (re, im) pairs, rotates by ``freqs_cis`` indexed along the - sequence axis, and conjugates for the inverse (de-rotation applied to the - attention output). ``x`` is the rope-dim slice only: ``[b, s, rd]`` (3D) or - ``[b, s, h, rd]`` (4D), with ``freqs_cis`` of shape ``[s, rd // 2]``. - """ - orig_dtype = x.dtype - xc = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) - if inverse: - freqs_cis = freqs_cis.conj() - if xc.ndim == 3: - fc = freqs_cis.view(1, xc.size(1), xc.size(-1)) - else: - fc = freqs_cis.view(1, xc.size(1), 1, xc.size(-1)) - out = torch.view_as_real(xc * fc).flatten(-2) - return out.to(orig_dtype) - - -def apply_dspark_rotary_batched( - x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False -) -> torch.Tensor: - """Per-row (batched) variant of :func:`apply_dspark_rotary`. - - Identical math, but ``freqs_cis`` carries a leading batch axis so each row of - ``x`` is rotated by its own per-request phases (the generation draft runs each - request at a different absolute ``start_pos``). ``x`` is the rope-dim slice - only: ``[G, s, rd]`` (3D) or ``[G, s, h, rd]`` (4D), with ``freqs_cis`` of shape - ``[G, s, rd // 2]``. - """ - orig_dtype = x.dtype - xc = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) - if inverse: - freqs_cis = freqs_cis.conj() - g, s, half = freqs_cis.shape - if xc.ndim == 3: - fc = freqs_cis.view(g, s, half) - else: - fc = freqs_cis.view(g, s, 1, half) - out = torch.view_as_real(xc * fc).flatten(-2) - return out.to(orig_dtype) - - -@lru_cache(maxsize=64) -def _topk_matrix(window_size: int, block_size: int, start_pos: int) -> torch.Tensor: - # [min(window, start_pos+1)] context positions in the rolling KV window, - # followed by [block_size] positions for the current block's own K/V (which - # the caller appends to the window at offset ``window_size``). - ctx = torch.arange(min(window_size, start_pos + 1)) - blk = window_size + torch.arange(block_size) - return torch.cat([ctx, blk]).int() - - -def get_dspark_topk_idxs( - window_size: int, - bsz: int, - block_size: int, - start_pos: int, - device: torch.device | str = "cpu", -) -> torch.Tensor: - """Per-query attended-position indices for the DSpark draft block. - - Mirrors DeepSpec ``get_dspark_topk_idxs``: every one of the ``block_size`` - query positions attends to the same set — the ``min(window_size, start_pos+1)`` - most-recent context positions in the rolling KV window, then the - ``block_size`` positions of the current block (stored at offset - ``window_size`` in the concatenated KV). Note this is *non-causal* within the - block (every position sees every block position), matching the reference. - - Args: - window_size: sliding-window length of the captured-context KV cache. - bsz: batch size. - block_size: number of draft positions per request. - start_pos: absolute decode position (must be > 0); bounds the context. - device: device for the returned index tensor. - - Returns: - int32 tensor ``[bsz, block_size, topk]`` with - ``topk = min(window_size, start_pos+1) + block_size``. - """ - assert start_pos > 0, "DSpark draft attention runs at generation (start_pos > 0)" - matrix = _topk_matrix(int(window_size), int(block_size), int(start_pos)).to(device) - return matrix.view(1, 1, -1).expand(bsz, block_size, -1).contiguous() - - -def get_dspark_topk_idxs_batched( - window_size: int, - block_size: int, - start_pos: torch.Tensor, - valid_len: torch.Tensor | None = None, -) -> torch.Tensor: - """Sync-free, fixed-size (CUDA-graph-safe) batched ``get_dspark_topk_idxs``. - - Unlike the scalar :func:`get_dspark_topk_idxs` (whose ``topk`` width - ``min(window_size, start_pos+1) + block_size`` depends on the host int - ``start_pos``), this always returns the **fixed** width ``window_size + - block_size`` and masks the unfilled context slots with ``-1``. The masked - slots are excluded by :func:`dspark_sparse_attn` exactly as if they were - absent while the shape remains CUDA-graph safe. - - Every query attends to the actually written circular-window suffix, followed - by the current-block positions. Without ``valid_len`` this preserves the - legacy ``start_pos``-only behavior. - - Args: - window_size: sliding-window length of the captured-context KV cache. - block_size: number of draft positions per request. - start_pos: ``[G]`` int tensor of per-request absolute decode positions. - valid_len: optional ``[G]`` count of actually written rolling-window - entries. When omitted, preserve the legacy ``start_pos`` mask. - - Returns: - int32 tensor ``[G, block_size, window_size + block_size]``. - """ - device = start_pos.device - g = start_pos.shape[0] - ctx_cols = torch.arange(window_size, device=device) # [win] - if valid_len is None: - valid = ctx_cols.unsqueeze(0) <= start_pos.unsqueeze(1) # [G, win] - else: - # The valid entries are the contiguous logical suffix ending at - # start_pos, but their physical slots wrap modulo window_size. - valid_len = valid_len.clamp(min=0, max=window_size) - age = torch.remainder(start_pos.unsqueeze(1) - ctx_cols.unsqueeze(0), window_size) - valid = age < valid_len.unsqueeze(1) - ctx_idx = torch.where( - valid, ctx_cols.unsqueeze(0).expand(g, -1), torch.full_like(valid, -1, dtype=torch.long) - ) - blk_idx = window_size + torch.arange(block_size, device=device) # [block] - blk_idx = blk_idx.unsqueeze(0).expand(g, -1) # [G, block] - row = torch.cat([ctx_idx, blk_idx], dim=1).to(torch.int32) # [G, win+block] - return row.unsqueeze(1).expand(g, block_size, -1).contiguous() - - -def dspark_sparse_attn( - q: torch.Tensor, - kv: torch.Tensor, - attn_sink: torch.Tensor, - topk_idxs: torch.Tensor, - softmax_scale: float, -) -> torch.Tensor: - """Index-gathered multi-query attention with an attention sink. - - Functional-first port of the DeepSpec ``sparse_attn`` TileLang kernel. For - each ``(batch, query, head)`` it gathers the ``topk`` KV rows named by - ``topk_idxs`` (an index of ``-1`` masks that slot), computes a scaled - dot-product softmax over them, and adds a per-head learnable *sink* logit that - participates only in the softmax denominator (i.e. an "attend-to-nothing" - option with a zero value vector). KV is shared across query heads (MQA). - - Args: - q: ``[b, m, h, d]`` query (``m`` = block_size, ``h`` = heads). - kv: ``[b, n, d]`` keys/values (shared across heads). - attn_sink: ``[h]`` per-head sink logits (fp32). - topk_idxs: ``[b, m, topk]`` int gather indices into ``kv`` (``-1`` masks). - softmax_scale: scalar applied to the q·k scores (``head_dim ** -0.5``). - - Returns: - ``[b, m, h, d]`` attention output, in ``q.dtype``. - """ - b, m, h, d = q.shape - idx = topk_idxs.long() # [b, m, topk] - valid = idx >= 0 - safe = idx.clamp(min=0) - - # Invalid slots read kv[0, :] (via safe.clamp), but masked_fill below - # zeros their softmax probs, so the einsum nullifies them. - kv_exp = kv.unsqueeze(1).expand(b, m, kv.shape[1], d) - gathered = torch.gather(kv_exp, 2, safe.unsqueeze(-1).expand(b, m, safe.shape[-1], d)).float() - - # Scores [b, m, h, topk]; mask invalid slots to -inf before the softmax. - scores = torch.einsum("bmhd,bmkd->bmhk", q.float(), gathered) * softmax_scale - scores = scores.masked_fill(~valid.unsqueeze(2), float("-inf")) - - # Online-softmax max is taken over gathered positions only (the sink is added - # to the denominator afterwards), matching the kernel's reduce order. - smax = scores.max(dim=-1, keepdim=True).values # [b, m, h, 1] - smax = torch.where(torch.isinf(smax), torch.zeros_like(smax), smax) - probs = torch.exp(scores - smax) # masked slots -> exp(-inf) = 0 - sink = torch.exp(attn_sink.to(torch.float32).view(1, 1, h) - smax.squeeze(-1)) - denom = probs.sum(dim=-1) + sink # [b, m, h] - out = torch.einsum("bmhk,bmkd->bmhd", probs, gathered) / denom.unsqueeze(-1) - return out.to(q.dtype) - - -def _rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: - """RMSNorm matching the DeepSpec reference (fp32 reduce, then * weight).""" - dtype = x.dtype - xf = x.float() - xf = xf * torch.rsqrt(xf.square().mean(-1, keepdim=True) + eps) - return (weight.float() * xf).to(dtype) - - -def _rope_last_dims( - t: torch.Tensor, rope_head_dim: int, freqs_cis: torch.Tensor, inverse: bool = False -) -> torch.Tensor: - """Apply RoPE to the last ``rope_head_dim`` dims; pass the rest through.""" - nope = t[..., :-rope_head_dim] - rope = apply_dspark_rotary(t[..., -rope_head_dim:], freqs_cis, inverse=inverse) - return torch.cat([nope, rope], dim=-1) - - -def _rope_last_dims_batched( - t: torch.Tensor, rope_head_dim: int, freqs_cis: torch.Tensor, inverse: bool = False -) -> torch.Tensor: - """Per-row variant of :func:`_rope_last_dims` (``freqs_cis`` has a batch axis).""" - nope = t[..., :-rope_head_dim] - rope = apply_dspark_rotary_batched(t[..., -rope_head_dim:], freqs_cis, inverse=inverse) - return torch.cat([nope, rope], dim=-1) - - -def _rmsnorm_rope_batched( - t: torch.Tensor, - weight: torch.Tensor, - eps: float, - rope_head_dim: int, - freqs_cis: torch.Tensor, - *, - num_heads: int = 1, - apply_weight: bool = True, - apply_rmsnorm: bool = True, - inverse_rope: bool = False, -) -> torch.Tensor: - """Fuse DSpark RMSNorm and last-dimension RoPE when supported.""" - 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): - return cute_dsl_dspark_rmsnorm_rope( - t, - weight, - freqs_real, - num_heads, - rope_head_dim, - eps, - apply_weight, - apply_rmsnorm, - inverse_rope, - ) - - if apply_rmsnorm: - if apply_weight: - t = _rmsnorm(t, weight, eps) - else: - t = t * torch.rsqrt(t.square().mean(-1, keepdim=True) + eps) - elif apply_weight: - t = (t.float() * weight.float()).to(t.dtype) - if rope_head_dim > 0: - t = _rope_last_dims_batched(t, rope_head_dim, freqs_cis, inverse=inverse_rope) - return t - - -def dspark_attention_forward( - x: torch.Tensor, - main_x: torch.Tensor, - start_pos: int, - kv_cache: torch.Tensor, - *, - wq_a: torch.Tensor, - q_norm_w: torch.Tensor, - wq_b: torch.Tensor, - wkv: torch.Tensor, - kv_norm_w: torch.Tensor, - wo_a: torch.Tensor, - wo_b: torch.Tensor, - attn_sink: torch.Tensor, - n_heads: int, - head_dim: int, - rope_head_dim: int, - n_groups: int, - o_lora_rank: int, - window_size: int, - eps: float, - softmax_scale: float, - freqs_cis: torch.Tensor, - persist: bool = False, -) -> torch.Tensor: - """Captured-context DSpark draft attention (generation path, ``start_pos > 0``). - - Functional port of DeepSpec ``DSparkAttention.forward`` for the dense - (``compress_ratio == 0``) draft: low-rank Q (``wq_a`` -> ``q_norm`` -> ``wq_b``) - with a per-head RMS + RoPE, MQA K/V from ``wkv`` (shared across heads), keys - gathered from a rolling captured-context window (``kv_cache``, into which the - projected ``main_x`` context is written at ``start_pos % window_size``) plus the - block's own positions, attention-sink softmax, inverse-RoPE on the output, and a - grouped low-rank O projection (``wo_a`` einsum + ``wo_b``). - - Weights are plain tensors for ``F.linear`` (the caller supplies the loaded / - dequantized projection weights); ``wo_a`` is the raw grouped weight matrix - ``[n_groups * o_lora_rank, n_heads * head_dim // n_groups]``. ``kv_cache`` is - ``[b, window_size, head_dim]`` and is updated functionally (cloned). - - Returns: - ``[b, block_size, dim]`` attention output (residual stream contribution). - """ - assert start_pos > 0, "DSpark draft attention runs at generation (start_pos > 0)" - b, block, _ = x.shape - rd = rope_head_dim - main_freqs = freqs_cis[start_pos : start_pos + 1] - blk_freqs = freqs_cis[start_pos + 1 : start_pos + 1 + block] - - # Captured-context K/V from main_x (MQA, shared across heads). - main_kv = _rmsnorm(F.linear(main_x, wkv), kv_norm_w, eps) # [b, 1, head_dim] - main_kv = _rope_last_dims(main_kv, rd, main_freqs) - - # Query: low-rank + per-head RMS + RoPE. - q = _rmsnorm(F.linear(x, wq_a), q_norm_w, eps) - q = F.linear(q, wq_b).unflatten(-1, (n_heads, head_dim)) # [b, block, h, head_dim] - # Per-head RMS in the query dtype (matches the reference inline normalization, - # which is NOT the fp32 RMSNorm path). - q = q * torch.rsqrt(q.square().mean(-1, keepdim=True) + eps) - q = _rope_last_dims(q, rd, blk_freqs) - - # Block K/V. - kv = _rmsnorm(F.linear(x, wkv), kv_norm_w, eps) # [b, block, head_dim] - kv = _rope_last_dims(kv, rd, blk_freqs) - - # Write the context K/V into the rolling window, then attend over - # [window context | block] with the sink. ``persist=True`` writes through - # to the caller's buffer (cross-step decode, worker-owned window); the - # default clones so single-shot callers (golden / unit tests) stay pure. - cache = kv_cache if persist else kv_cache.clone() - cache[:, start_pos % window_size] = main_kv.squeeze(1) - kv_full = torch.cat([cache, kv], dim=1) # [b, window + block, head_dim] - topk = get_dspark_topk_idxs(window_size, b, block, start_pos, device=x.device) - o = dspark_sparse_attn(q, kv_full, attn_sink, topk, softmax_scale) # [b, block, h, head_dim] - o = _rope_last_dims(o, rd, blk_freqs, inverse=True) - - # Grouped low-rank O projection. - o = o.reshape(b, block, n_groups, -1) - wo_a_v = wo_a.view(n_groups, o_lora_rank, -1) - o = torch.einsum("bsgd,grd->bsgr", o, wo_a_v) - return F.linear(o.flatten(2), wo_b) - - -def dspark_attention_forward_batched( - x: torch.Tensor, - main_x: torch.Tensor, - start_pos: torch.Tensor, - kv_cache: torch.Tensor, - slots: torch.Tensor, - valid_len: torch.Tensor | None = None, - *, - wq_a: torch.Tensor, - q_norm_w: torch.Tensor, - wq_b: torch.Tensor, - wkv: torch.Tensor, - kv_norm_w: torch.Tensor, - wo_a: torch.Tensor, - wo_b: torch.Tensor, - attn_sink: torch.Tensor, - n_heads: int, - head_dim: int, - rope_head_dim: int, - n_groups: int, - o_lora_rank: int, - window_size: int, - eps: float, - softmax_scale: float, - freqs_cis: torch.Tensor, - persist: bool = False, -) -> torch.Tensor: - """Batched, CUDA-graph-safe captured-context DSpark draft attention. - - Numerically identical, per request, to :func:`dspark_attention_forward`, but - free of host syncs and data-dependent shapes so it can be captured into a CUDA - graph (the one-engine drafter runs inside the target's graph). The differences - from the scalar path are purely mechanical: - - * ``start_pos`` is a ``[G]`` int tensor (one absolute decode position per gen - request) instead of a python int; RoPE phases are *gathered* per request from - the fixed ``freqs_cis`` table rather than sliced. - * the rolling-window context K/V is written/read through the ``slots`` index - into a shared ``kv_cache`` (``persist=True`` writes through to the caller's - worker-owned buffer; otherwise a clone is used), instead of mutating a - per-request cache in place. - * the attended-position list has the fixed width ``window_size + block_size`` - with ``-1`` masking (see :func:`get_dspark_topk_idxs_batched`). - - Args: - x: ``[G, block, dim]`` block layer input (per gen request). - main_x: ``[G, 1, hidden]`` projected captured context. - start_pos: ``[G]`` int tensor of absolute decode positions (> 0). - kv_cache: ``[N, window_size, head_dim]`` rolling captured-context windows - (``N`` rows indexed by ``slots``; ``N == G`` for single-shot callers). - slots: ``[G]`` int tensor mapping each request to its ``kv_cache`` row. - valid_len: optional ``[G]`` count of actually written context entries; - masks holes left when absolute positions are bootstrapped without - receiving the corresponding DSpark rolling-window state. - freqs_cis: ``[maxlen, rope_head_dim // 2]`` precomputed plain-RoPE table; - must satisfy ``maxlen > start_pos.max() + block_size``. - - Returns: - ``[G, block, dim]`` attention output (residual stream contribution). - """ - g, block, _ = x.shape - if kv_cache.shape[1] != window_size: - raise ValueError( - f"kv_cache window extent {kv_cache.shape[1]} does not match window_size {window_size}" - ) - rd = rope_head_dim - # Per-request RoPE phases gathered from the fixed table (no host-int slicing). - main_freqs = freqs_cis[start_pos].unsqueeze(1) # [G, 1, rd//2] - blk_pos = start_pos.unsqueeze(1) + 1 + torch.arange(block, device=x.device) # [G, block] - blk_freqs = freqs_cis[blk_pos] # [G, block, rd//2] - - # Captured-context K/V from main_x (MQA, shared across heads). - main_kv = _rmsnorm_rope_batched(F.linear(main_x, wkv), kv_norm_w, eps, rd, main_freqs) - - # Query: low-rank + per-head RMS + RoPE. - q = _rmsnorm_rope_batched(F.linear(x, wq_a), q_norm_w, eps, 0, blk_freqs) - 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, - eps, - rd, - blk_freqs, - num_heads=n_heads, - apply_weight=False, - ) - - # Block K/V. - kv = _rmsnorm_rope_batched(F.linear(x, wkv), kv_norm_w, eps, rd, blk_freqs) - - # Write the context K/V into the rolling window at slot start_pos%window_size, - # then attend over [window context | block]. ``persist=True`` writes through to - # the worker-owned buffer (cross-step decode); otherwise clone so single-shot - # callers stay pure. - write_target = kv_cache if persist else kv_cache.clone() - main_kv_flat = main_kv.squeeze(1).to(write_target.dtype) - if ( - valid_len is None - and IS_CUTLASS_DSL_AVAILABLE - and is_fused_dspark_attention_supported( - q, main_kv_flat, kv, write_target, slots, start_pos, attn_sink - ) - ): - # One custom op performs the rolling-cache write/read, validity handling, - # QK, attention-sink online softmax, and PV. In particular it creates no - # topk index, gathered KV, score, or probability tensors. - o = cute_dsl_dspark_attention( - q, - main_kv_flat, - kv, - write_target, - slots, - start_pos, - attn_sink, - softmax_scale, - ) - else: - slot_pos = start_pos % window_size # [G] - write_target[slots, slot_pos] = main_kv_flat - cache_rows = write_target[slots] # [G, window, head_dim] - kv_full = torch.cat([cache_rows, kv], dim=1) # [G, window + block, head_dim] - topk = get_dspark_topk_idxs_batched(window_size, block, start_pos, valid_len) - o = dspark_sparse_attn( - q, kv_full, attn_sink, topk, softmax_scale - ) # [G, block, h, head_dim] - o = _rmsnorm_rope_batched( - o, - kv_norm_w, - eps, - rd, - blk_freqs, - num_heads=n_heads, - apply_weight=False, - apply_rmsnorm=False, - inverse_rope=True, - ) - - # Grouped low-rank O projection. - o = o.reshape(g, block, n_groups, -1) - wo_a_v = wo_a.view(n_groups, o_lora_rank, -1) - o = torch.einsum("bsgd,grd->bsgr", o, wo_a_v) - return F.linear(o.flatten(2), wo_b) diff --git a/tensorrt_llm/_torch/models/dspark/draft.py b/tensorrt_llm/_torch/models/dspark/draft.py deleted file mode 100644 index 99a9021ca2c7..000000000000 --- a/tensorrt_llm/_torch/models/dspark/draft.py +++ /dev/null @@ -1,130 +0,0 @@ -# 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. -# -# DSpark draft I/O logic is ported from DeepSeek's DeepSeek-V4-Pro-DSpark -# reference (`inference/model.py`, DSparkBlock.forward_embed / forward_head). -"""DSpark draft I/O: block input and proposal stages. - -This module holds the *framework-agnostic* (pure-torch) input/output stages of -the DSpark draft block, separated from the heavy V4 backbone (MLA + MoE + mHC) so -they can be unit-tested in isolation: - - - ``build_draft_input_ids``: ``[bonus_token, noise, noise, ...]`` block input. - - ``dspark_propose``: given the per-position backbone ``base_logits`` and the - Markov / confidence heads, run the autoregressive Markov refinement to sample - the block tokens and apply the static confidence-threshold truncation. - -The backbone (3 V4 blocks producing ``block_hidden``) lives in the model module; -this file is the part fully specified by the reference and validated against it. -""" - -from typing import Optional - -import torch -from torch import nn - -from ..modeling_speculative import confident_prefix_length - - -def build_draft_input_ids( - bonus_token_ids: torch.Tensor, *, block_size: int, noise_token_id: int -) -> torch.Tensor: - """``[batch] -> [batch, block_size]`` = ``[bonus, noise, noise, ...]``. - - The first position is the verified bonus token (the target's last accepted - token); the rest are the DSpark noise/mask token (id 128799 for V4-Pro). - """ - batch = bonus_token_ids.shape[0] - out = bonus_token_ids.new_full((batch, block_size), int(noise_token_id)) - out[:, 0] = bonus_token_ids - return out - - -def dspark_propose( - base_logits: torch.Tensor, - *, - bonus_token_ids: torch.Tensor, - block_hidden: torch.Tensor, - markov_head: Optional[nn.Module], - confidence_head: Optional[nn.Module], - block_size: int, - temperature: float = 0.0, - confidence_threshold: float = 0.0, - return_logits: bool = False, -) -> tuple: - """Produce DSpark draft tokens for one block (functional-first, static length). - - Args: - base_logits: ``[batch, block_size, vocab]`` from the backbone + lm_head. - bonus_token_ids: ``[batch]`` the token preceding the first draft position. - block_hidden: ``[batch, block_size, hidden]`` backbone hidden (feeds the - confidence head, and the RNN-head variant). - markov_head / confidence_head: the validated DSpark heads (may be None). - Returns: - draft_tokens: ``[batch, block_size]`` sampled tokens (full block; callers - keep the tensor fixed-width for CUDA-graph safety). - num_proposed: ``[batch]`` int32 — how many leading tokens survive the - static confidence-threshold truncation (== block_size when no head / - threshold<=0). - """ - batch = base_logits.shape[0] - # ``draft_logits`` are the per-position distributions the draft token is drawn - # from (markov-corrected when a head is present, else the raw base logits). - # Surfaced under ``return_logits`` for the §7.9 probabilistic-acceptance - # (1-TV) measurement; the normal path ignores them. - draft_logits = base_logits - if markov_head is not None: - draft_tokens, corrected = markov_head.sample_block_tokens( - base_logits, - first_prev_token_ids=bonus_token_ids, - hidden_states=block_hidden, - temperature=temperature, - ) - draft_logits = corrected - else: - from ..modeling_speculative import greedy_or_sample - - draft_tokens = greedy_or_sample(base_logits, temperature) - - # Scaffolding: confidence-based dynamic drafting is NOT enabled in this PR. - # The worker always calls with confidence_threshold=0.0, so the block below is - # inert and num_proposed stays == block_size (the full block is proposed). The - # returned num_proposed is intentionally not yet consumed by the speculative - # scheduler/verifier; wiring it through is a follow-up (see PR description). - num_proposed = torch.full( - (batch,), int(block_size), dtype=torch.int32, device=base_logits.device - ) - if confidence_head is not None and confidence_threshold > 0.0: - # prev token at position k is [bonus, draft_0, ..., draft_{k-1}] - prev_ids = torch.cat([bonus_token_ids.unsqueeze(1), draft_tokens[:, :-1]], dim=1) - prev_emb = ( - markov_head.get_prev_embeddings(prev_ids) - if (markov_head is not None and getattr(confidence_head, "with_markov", False)) - else None - ) - conf_logits = ( - confidence_head(block_hidden, prev_embeddings=prev_emb) - if prev_emb is not None - else confidence_head(block_hidden) - ) - # Per-request prefix truncation (batch handled row-wise to stay simple; - # functional-first scope typically runs batch=1 for the draft). - for b in range(batch): - num_proposed[b] = confident_prefix_length( - conf_logits[b : b + 1], block_size=block_size, threshold=confidence_threshold - ) - if return_logits: - return draft_tokens, num_proposed, draft_logits - return draft_tokens, num_proposed diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 8bbd0e2f2667..aebafa552838 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -15,29 +15,59 @@ # # DSpark backbone ported from the DeepSeek-V4-Pro-DSpark reference # (`inference/model.py`: DSparkBlock / Transformer.forward_spec). -"""DeepSeek-V4-Pro DSpark speculative-decoding draft backbone. - -The DSpark draft is ``n_mtp_layers`` (3 for V4-Pro) **full DeepSeek-V4 blocks** -stored under the ``mtp.*`` checkpoint namespace — it reuses the V4 decoder block -(MLA attention + MoE + manifold Hyper-Connections) and adds: - - - **stage 0**: ``main_proj`` (Linear, fp8) + ``main_norm`` (RMSNorm) — projects - the concatenation of captured target-layer hidden states ([58,59,60]) into the - draft's cross-attention context (``main_x``); replaces vanilla-MTP's - enorm/hnorm + e_proj/h_proj single-hidden mixing. - - **last stage**: ``norm`` + ``markov_head`` + ``confidence_head`` + - flat ``hc_head`` — the block-draft output head (see dspark_heads/dspark_draft). - -The per-stage *backbone* forward (block attention whose K/V derive from ``main_x``, -+ MoE + mHC) is brought up and numerically validated against the real fp8 weights -separately; ``forward_embed`` (capture) and ``forward_head`` (block draft) below -are the reference-faithful, unit-validated I/O stages. +# +# The captured-context attention primitives are ported from DeepSeek's DeepSpec +# reference ``inference/kernel.py`` (``sparse_attn``) and ``inference/model.py`` +# (``get_dspark_topk_idxs``). The reference computes these with a TileLang +# kernel; this is a functional-first pure-PyTorch port with the same math +# (index-gather + online softmax + a learnable attention sink that contributes +# only to the softmax denominator). +# +# The draft I/O stages are ported from the same reference +# (`inference/model.py`: DSparkBlock.forward_embed / forward_head). +"""DeepSeek-V4-Pro DSpark speculative-decoding draft. + +``DSpark`` names the speculative-decoding *algorithm*: a parallel block draft +over captured target hidden states, refined by a low-rank Markov head and +scheduled by a confidence head. This module holds its **in-checkpoint** flavour, +where the draft weights ship inside the DeepSeek-V4-Pro *target* checkpoint under +the ``mtp.*`` namespace and reuse the V4 decoder block, so the draft inherits the +target's EPLB layer namespace and fp8/NVFP4 quantization. Standalone DSpark +drafters — shipped as their own checkpoint, backbone resolved through the model +registry — are built by :mod:`modeling_dflash` instead. + +Three parts live here: + +1. **Draft backbone** — ``n_mtp_layers`` (3 for V4-Pro) full DeepSeek-V4 blocks + (MLA attention + MoE + manifold Hyper-Connections), plus: + + - **stage 0**: ``main_proj`` (Linear, fp8) + ``main_norm`` (RMSNorm) — + projects the concatenation of captured target-layer hidden states + ([58,59,60]) into the draft's cross-attention context (``main_x``); + replaces vanilla-MTP's enorm/hnorm + e_proj/h_proj single-hidden mixing. + - **last stage**: ``norm`` + ``markov_head`` + ``confidence_head`` + flat + ``hc_head`` — the block-draft output head. The heads themselves are shared + with the standalone path and live in :mod:`modeling_speculative`. + +2. **Captured-context attention primitives** — the dense sliding-window MLA the + draft block attends with (``get_dspark_topk_idxs`` / ``dspark_sparse_attn`` + and the rotary helpers). Hardware-agnostic and unit-testable in isolation. + +3. **Block draft I/O** — ``build_draft_input_ids`` (the + ``[bonus_token, noise, ...]`` block input) and ``dspark_propose`` (Markov + refinement + static confidence truncation). + +The per-stage *backbone* forward (block attention whose K/V derive from +``main_x``, + MoE + mHC) is brought up and numerically validated against the real +fp8 weights separately; ``forward_embed`` (capture) and ``forward_head`` (block +draft) are the reference-faithful, unit-validated I/O stages. """ import copy import json import os import re +from functools import lru_cache from typing import Dict, List, Optional import torch @@ -47,21 +77,14 @@ from tensorrt_llm.logger import logger 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 ..modules.linear import Linear from ..modules.mhc.hyper_connection import HCHead from ..modules.rms_norm import RMSNorm from ..speculative.interface import SpeculativeDecodingMode from ..utils import AuxStreamType -from .dspark.attention import ( - _rmsnorm, - _rope_last_dims, - _rope_last_dims_batched, - dspark_attention_forward, - dspark_attention_forward_batched, - precompute_dspark_freqs_cis, -) -from .dspark.draft import build_draft_input_ids, dspark_propose from .modeling_deepseekv4 import ( DeepseekV4DecoderLayer, DeepseekV4WeightLoader, @@ -71,9 +94,631 @@ _rename_deepseek_v4_attn_subkey, _rename_deepseek_v4_ffn_subkey, ) -from .modeling_speculative import DSparkConfidenceHead, build_markov_head +from .modeling_speculative import DSparkConfidenceHead, build_markov_head, confident_prefix_length from .modeling_utils import register_draft_model +if IS_CUTLASS_DSL_AVAILABLE: + from ..custom_ops.dspark_attention_custom_op import ( + cute_dsl_dspark_attention, + is_fused_dspark_attention_supported, + ) + from ..custom_ops.dspark_rmsnorm_rope_custom_op import ( + cute_dsl_dspark_rmsnorm_rope, + is_fused_dspark_rmsnorm_rope_supported, + ) + +# ---------------------------------------------------------------------------- +# Captured-context attention primitives. +# ---------------------------------------------------------------------------- + + +def precompute_dspark_freqs_cis( + rope_head_dim: int, + seqlen: int, + rope_theta: float = 10000.0, + device: torch.device | str = "cpu", +) -> torch.Tensor: + """Plain (non-YaRN) RoPE complex exponentials for the DSpark draft. + + The dense draft attention (``compress_ratio == 0``) disables YaRN and uses the + base ``rope_theta`` (DeepSpec ``precompute_freqs_cis`` with + ``original_seq_len == 0``). + + Returns: + complex64 tensor ``[seqlen, rope_head_dim // 2]``. + """ + freqs = 1.0 / ( + rope_theta + ** (torch.arange(0, rope_head_dim, 2, dtype=torch.float32, device=device) / rope_head_dim) + ) + t = torch.arange(seqlen, dtype=torch.float32, device=device) + freqs = torch.outer(t, freqs) + return torch.polar(torch.ones_like(freqs), freqs) + + +def apply_dspark_rotary( + x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False +) -> torch.Tensor: + """Apply (or, with ``inverse``, de-apply) rotary embeddings, DeepSpec-style. + + Functional (non-in-place) port of DeepSpec ``apply_rotary_emb``: treats the + last dim as adjacent (re, im) pairs, rotates by ``freqs_cis`` indexed along the + sequence axis, and conjugates for the inverse (de-rotation applied to the + attention output). ``x`` is the rope-dim slice only: ``[b, s, rd]`` (3D) or + ``[b, s, h, rd]`` (4D), with ``freqs_cis`` of shape ``[s, rd // 2]``. + """ + orig_dtype = x.dtype + xc = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) + if inverse: + freqs_cis = freqs_cis.conj() + if xc.ndim == 3: + fc = freqs_cis.view(1, xc.size(1), xc.size(-1)) + else: + fc = freqs_cis.view(1, xc.size(1), 1, xc.size(-1)) + out = torch.view_as_real(xc * fc).flatten(-2) + return out.to(orig_dtype) + + +def apply_dspark_rotary_batched( + x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False +) -> torch.Tensor: + """Per-row (batched) variant of :func:`apply_dspark_rotary`. + + Identical math, but ``freqs_cis`` carries a leading batch axis so each row of + ``x`` is rotated by its own per-request phases (the generation draft runs each + request at a different absolute ``start_pos``). ``x`` is the rope-dim slice + only: ``[G, s, rd]`` (3D) or ``[G, s, h, rd]`` (4D), with ``freqs_cis`` of shape + ``[G, s, rd // 2]``. + """ + orig_dtype = x.dtype + xc = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) + if inverse: + freqs_cis = freqs_cis.conj() + g, s, half = freqs_cis.shape + if xc.ndim == 3: + fc = freqs_cis.view(g, s, half) + else: + fc = freqs_cis.view(g, s, 1, half) + out = torch.view_as_real(xc * fc).flatten(-2) + return out.to(orig_dtype) + + +@lru_cache(maxsize=64) +def _topk_matrix(window_size: int, block_size: int, start_pos: int) -> torch.Tensor: + # [min(window, start_pos+1)] context positions in the rolling KV window, + # followed by [block_size] positions for the current block's own K/V (which + # the caller appends to the window at offset ``window_size``). + ctx = torch.arange(min(window_size, start_pos + 1)) + blk = window_size + torch.arange(block_size) + return torch.cat([ctx, blk]).int() + + +def get_dspark_topk_idxs( + window_size: int, + bsz: int, + block_size: int, + start_pos: int, + device: torch.device | str = "cpu", +) -> torch.Tensor: + """Per-query attended-position indices for the DSpark draft block. + + Mirrors DeepSpec ``get_dspark_topk_idxs``: every one of the ``block_size`` + query positions attends to the same set — the ``min(window_size, start_pos+1)`` + most-recent context positions in the rolling KV window, then the + ``block_size`` positions of the current block (stored at offset + ``window_size`` in the concatenated KV). Note this is *non-causal* within the + block (every position sees every block position), matching the reference. + + Args: + window_size: sliding-window length of the captured-context KV cache. + bsz: batch size. + block_size: number of draft positions per request. + start_pos: absolute decode position (must be > 0); bounds the context. + device: device for the returned index tensor. + + Returns: + int32 tensor ``[bsz, block_size, topk]`` with + ``topk = min(window_size, start_pos+1) + block_size``. + """ + assert start_pos > 0, "DSpark draft attention runs at generation (start_pos > 0)" + matrix = _topk_matrix(int(window_size), int(block_size), int(start_pos)).to(device) + return matrix.view(1, 1, -1).expand(bsz, block_size, -1).contiguous() + + +def get_dspark_topk_idxs_batched( + window_size: int, + block_size: int, + start_pos: torch.Tensor, + valid_len: torch.Tensor | None = None, +) -> torch.Tensor: + """Sync-free, fixed-size (CUDA-graph-safe) batched ``get_dspark_topk_idxs``. + + Unlike the scalar :func:`get_dspark_topk_idxs` (whose ``topk`` width + ``min(window_size, start_pos+1) + block_size`` depends on the host int + ``start_pos``), this always returns the **fixed** width ``window_size + + block_size`` and masks the unfilled context slots with ``-1``. The masked + slots are excluded by :func:`dspark_sparse_attn` exactly as if they were + absent while the shape remains CUDA-graph safe. + + Every query attends to the actually written circular-window suffix, followed + by the current-block positions. Without ``valid_len`` this preserves the + legacy ``start_pos``-only behavior. + + Args: + window_size: sliding-window length of the captured-context KV cache. + block_size: number of draft positions per request. + start_pos: ``[G]`` int tensor of per-request absolute decode positions. + valid_len: optional ``[G]`` count of actually written rolling-window + entries. When omitted, preserve the legacy ``start_pos`` mask. + + Returns: + int32 tensor ``[G, block_size, window_size + block_size]``. + """ + device = start_pos.device + g = start_pos.shape[0] + ctx_cols = torch.arange(window_size, device=device) # [win] + if valid_len is None: + valid = ctx_cols.unsqueeze(0) <= start_pos.unsqueeze(1) # [G, win] + else: + # The valid entries are the contiguous logical suffix ending at + # start_pos, but their physical slots wrap modulo window_size. + valid_len = valid_len.clamp(min=0, max=window_size) + age = torch.remainder(start_pos.unsqueeze(1) - ctx_cols.unsqueeze(0), window_size) + valid = age < valid_len.unsqueeze(1) + ctx_idx = torch.where( + valid, ctx_cols.unsqueeze(0).expand(g, -1), torch.full_like(valid, -1, dtype=torch.long) + ) + blk_idx = window_size + torch.arange(block_size, device=device) # [block] + blk_idx = blk_idx.unsqueeze(0).expand(g, -1) # [G, block] + row = torch.cat([ctx_idx, blk_idx], dim=1).to(torch.int32) # [G, win+block] + return row.unsqueeze(1).expand(g, block_size, -1).contiguous() + + +def dspark_sparse_attn( + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + """Index-gathered multi-query attention with an attention sink. + + Functional-first port of the DeepSpec ``sparse_attn`` TileLang kernel. For + each ``(batch, query, head)`` it gathers the ``topk`` KV rows named by + ``topk_idxs`` (an index of ``-1`` masks that slot), computes a scaled + dot-product softmax over them, and adds a per-head learnable *sink* logit that + participates only in the softmax denominator (i.e. an "attend-to-nothing" + option with a zero value vector). KV is shared across query heads (MQA). + + Args: + q: ``[b, m, h, d]`` query (``m`` = block_size, ``h`` = heads). + kv: ``[b, n, d]`` keys/values (shared across heads). + attn_sink: ``[h]`` per-head sink logits (fp32). + topk_idxs: ``[b, m, topk]`` int gather indices into ``kv`` (``-1`` masks). + softmax_scale: scalar applied to the q·k scores (``head_dim ** -0.5``). + + Returns: + ``[b, m, h, d]`` attention output, in ``q.dtype``. + """ + b, m, h, d = q.shape + idx = topk_idxs.long() # [b, m, topk] + valid = idx >= 0 + safe = idx.clamp(min=0) + + # Invalid slots read kv[0, :] (via safe.clamp), but masked_fill below + # zeros their softmax probs, so the einsum nullifies them. + kv_exp = kv.unsqueeze(1).expand(b, m, kv.shape[1], d) + gathered = torch.gather(kv_exp, 2, safe.unsqueeze(-1).expand(b, m, safe.shape[-1], d)).float() + + # Scores [b, m, h, topk]; mask invalid slots to -inf before the softmax. + scores = torch.einsum("bmhd,bmkd->bmhk", q.float(), gathered) * softmax_scale + scores = scores.masked_fill(~valid.unsqueeze(2), float("-inf")) + + # Online-softmax max is taken over gathered positions only (the sink is added + # to the denominator afterwards), matching the kernel's reduce order. + smax = scores.max(dim=-1, keepdim=True).values # [b, m, h, 1] + smax = torch.where(torch.isinf(smax), torch.zeros_like(smax), smax) + probs = torch.exp(scores - smax) # masked slots -> exp(-inf) = 0 + sink = torch.exp(attn_sink.to(torch.float32).view(1, 1, h) - smax.squeeze(-1)) + denom = probs.sum(dim=-1) + sink # [b, m, h] + out = torch.einsum("bmhk,bmkd->bmhd", probs, gathered) / denom.unsqueeze(-1) + return out.to(q.dtype) + + +def _rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """RMSNorm matching the DeepSpec reference (fp32 reduce, then * weight).""" + dtype = x.dtype + xf = x.float() + xf = xf * torch.rsqrt(xf.square().mean(-1, keepdim=True) + eps) + return (weight.float() * xf).to(dtype) + + +def _rope_last_dims( + t: torch.Tensor, rope_head_dim: int, freqs_cis: torch.Tensor, inverse: bool = False +) -> torch.Tensor: + """Apply RoPE to the last ``rope_head_dim`` dims; pass the rest through.""" + nope = t[..., :-rope_head_dim] + rope = apply_dspark_rotary(t[..., -rope_head_dim:], freqs_cis, inverse=inverse) + return torch.cat([nope, rope], dim=-1) + + +def _rope_last_dims_batched( + t: torch.Tensor, rope_head_dim: int, freqs_cis: torch.Tensor, inverse: bool = False +) -> torch.Tensor: + """Per-row variant of :func:`_rope_last_dims` (``freqs_cis`` has a batch axis).""" + nope = t[..., :-rope_head_dim] + rope = apply_dspark_rotary_batched(t[..., -rope_head_dim:], freqs_cis, inverse=inverse) + return torch.cat([nope, rope], dim=-1) + + +def _rmsnorm_rope_batched( + t: torch.Tensor, + weight: torch.Tensor, + eps: float, + rope_head_dim: int, + freqs_cis: torch.Tensor, + *, + num_heads: int = 1, + apply_weight: bool = True, + apply_rmsnorm: bool = True, + inverse_rope: bool = False, +) -> torch.Tensor: + """Fuse DSpark RMSNorm and last-dimension RoPE when supported.""" + 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): + return cute_dsl_dspark_rmsnorm_rope( + t, + weight, + freqs_real, + num_heads, + rope_head_dim, + eps, + apply_weight, + apply_rmsnorm, + inverse_rope, + ) + + if apply_rmsnorm: + if apply_weight: + t = _rmsnorm(t, weight, eps) + else: + t = t * torch.rsqrt(t.square().mean(-1, keepdim=True) + eps) + elif apply_weight: + t = (t.float() * weight.float()).to(t.dtype) + if rope_head_dim > 0: + t = _rope_last_dims_batched(t, rope_head_dim, freqs_cis, inverse=inverse_rope) + return t + + +def dspark_attention_forward( + x: torch.Tensor, + main_x: torch.Tensor, + start_pos: int, + kv_cache: torch.Tensor, + *, + wq_a: torch.Tensor, + q_norm_w: torch.Tensor, + wq_b: torch.Tensor, + wkv: torch.Tensor, + kv_norm_w: torch.Tensor, + wo_a: torch.Tensor, + wo_b: torch.Tensor, + attn_sink: torch.Tensor, + n_heads: int, + head_dim: int, + rope_head_dim: int, + n_groups: int, + o_lora_rank: int, + window_size: int, + eps: float, + softmax_scale: float, + freqs_cis: torch.Tensor, + persist: bool = False, +) -> torch.Tensor: + """Captured-context DSpark draft attention (generation path, ``start_pos > 0``). + + Functional port of DeepSpec ``DSparkAttention.forward`` for the dense + (``compress_ratio == 0``) draft: low-rank Q (``wq_a`` -> ``q_norm`` -> ``wq_b``) + with a per-head RMS + RoPE, MQA K/V from ``wkv`` (shared across heads), keys + gathered from a rolling captured-context window (``kv_cache``, into which the + projected ``main_x`` context is written at ``start_pos % window_size``) plus the + block's own positions, attention-sink softmax, inverse-RoPE on the output, and a + grouped low-rank O projection (``wo_a`` einsum + ``wo_b``). + + Weights are plain tensors for ``F.linear`` (the caller supplies the loaded / + dequantized projection weights); ``wo_a`` is the raw grouped weight matrix + ``[n_groups * o_lora_rank, n_heads * head_dim // n_groups]``. ``kv_cache`` is + ``[b, window_size, head_dim]`` and is updated functionally (cloned). + + Returns: + ``[b, block_size, dim]`` attention output (residual stream contribution). + """ + assert start_pos > 0, "DSpark draft attention runs at generation (start_pos > 0)" + b, block, _ = x.shape + rd = rope_head_dim + main_freqs = freqs_cis[start_pos : start_pos + 1] + blk_freqs = freqs_cis[start_pos + 1 : start_pos + 1 + block] + + # Captured-context K/V from main_x (MQA, shared across heads). + main_kv = _rmsnorm(F.linear(main_x, wkv), kv_norm_w, eps) # [b, 1, head_dim] + main_kv = _rope_last_dims(main_kv, rd, main_freqs) + + # Query: low-rank + per-head RMS + RoPE. + q = _rmsnorm(F.linear(x, wq_a), q_norm_w, eps) + q = F.linear(q, wq_b).unflatten(-1, (n_heads, head_dim)) # [b, block, h, head_dim] + # Per-head RMS in the query dtype (matches the reference inline normalization, + # which is NOT the fp32 RMSNorm path). + q = q * torch.rsqrt(q.square().mean(-1, keepdim=True) + eps) + q = _rope_last_dims(q, rd, blk_freqs) + + # Block K/V. + kv = _rmsnorm(F.linear(x, wkv), kv_norm_w, eps) # [b, block, head_dim] + kv = _rope_last_dims(kv, rd, blk_freqs) + + # Write the context K/V into the rolling window, then attend over + # [window context | block] with the sink. ``persist=True`` writes through + # to the caller's buffer (cross-step decode, worker-owned window); the + # default clones so single-shot callers (golden / unit tests) stay pure. + cache = kv_cache if persist else kv_cache.clone() + cache[:, start_pos % window_size] = main_kv.squeeze(1) + kv_full = torch.cat([cache, kv], dim=1) # [b, window + block, head_dim] + topk = get_dspark_topk_idxs(window_size, b, block, start_pos, device=x.device) + o = dspark_sparse_attn(q, kv_full, attn_sink, topk, softmax_scale) # [b, block, h, head_dim] + o = _rope_last_dims(o, rd, blk_freqs, inverse=True) + + # Grouped low-rank O projection. + o = o.reshape(b, block, n_groups, -1) + wo_a_v = wo_a.view(n_groups, o_lora_rank, -1) + o = torch.einsum("bsgd,grd->bsgr", o, wo_a_v) + return F.linear(o.flatten(2), wo_b) + + +def dspark_attention_forward_batched( + x: torch.Tensor, + main_x: torch.Tensor, + start_pos: torch.Tensor, + kv_cache: torch.Tensor, + slots: torch.Tensor, + valid_len: torch.Tensor | None = None, + *, + wq_a: torch.Tensor, + q_norm_w: torch.Tensor, + wq_b: torch.Tensor, + wkv: torch.Tensor, + kv_norm_w: torch.Tensor, + wo_a: torch.Tensor, + wo_b: torch.Tensor, + attn_sink: torch.Tensor, + n_heads: int, + head_dim: int, + rope_head_dim: int, + n_groups: int, + o_lora_rank: int, + window_size: int, + eps: float, + softmax_scale: float, + freqs_cis: torch.Tensor, + persist: bool = False, +) -> torch.Tensor: + """Batched, CUDA-graph-safe captured-context DSpark draft attention. + + Numerically identical, per request, to :func:`dspark_attention_forward`, but + free of host syncs and data-dependent shapes so it can be captured into a CUDA + graph (the one-engine drafter runs inside the target's graph). The differences + from the scalar path are purely mechanical: + + * ``start_pos`` is a ``[G]`` int tensor (one absolute decode position per gen + request) instead of a python int; RoPE phases are *gathered* per request from + the fixed ``freqs_cis`` table rather than sliced. + * the rolling-window context K/V is written/read through the ``slots`` index + into a shared ``kv_cache`` (``persist=True`` writes through to the caller's + worker-owned buffer; otherwise a clone is used), instead of mutating a + per-request cache in place. + * the attended-position list has the fixed width ``window_size + block_size`` + with ``-1`` masking (see :func:`get_dspark_topk_idxs_batched`). + + Args: + x: ``[G, block, dim]`` block layer input (per gen request). + main_x: ``[G, 1, hidden]`` projected captured context. + start_pos: ``[G]`` int tensor of absolute decode positions (> 0). + kv_cache: ``[N, window_size, head_dim]`` rolling captured-context windows + (``N`` rows indexed by ``slots``; ``N == G`` for single-shot callers). + slots: ``[G]`` int tensor mapping each request to its ``kv_cache`` row. + valid_len: optional ``[G]`` count of actually written context entries; + masks holes left when absolute positions are bootstrapped without + receiving the corresponding DSpark rolling-window state. + freqs_cis: ``[maxlen, rope_head_dim // 2]`` precomputed plain-RoPE table; + must satisfy ``maxlen > start_pos.max() + block_size``. + + Returns: + ``[G, block, dim]`` attention output (residual stream contribution). + """ + g, block, _ = x.shape + if kv_cache.shape[1] != window_size: + raise ValueError( + f"kv_cache window extent {kv_cache.shape[1]} does not match window_size {window_size}" + ) + rd = rope_head_dim + # Per-request RoPE phases gathered from the fixed table (no host-int slicing). + main_freqs = freqs_cis[start_pos].unsqueeze(1) # [G, 1, rd//2] + blk_pos = start_pos.unsqueeze(1) + 1 + torch.arange(block, device=x.device) # [G, block] + blk_freqs = freqs_cis[blk_pos] # [G, block, rd//2] + + # Captured-context K/V from main_x (MQA, shared across heads). + main_kv = _rmsnorm_rope_batched(F.linear(main_x, wkv), kv_norm_w, eps, rd, main_freqs) + + # Query: low-rank + per-head RMS + RoPE. + q = _rmsnorm_rope_batched(F.linear(x, wq_a), q_norm_w, eps, 0, blk_freqs) + 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, + eps, + rd, + blk_freqs, + num_heads=n_heads, + apply_weight=False, + ) + + # Block K/V. + kv = _rmsnorm_rope_batched(F.linear(x, wkv), kv_norm_w, eps, rd, blk_freqs) + + # Write the context K/V into the rolling window at slot start_pos%window_size, + # then attend over [window context | block]. ``persist=True`` writes through to + # the worker-owned buffer (cross-step decode); otherwise clone so single-shot + # callers stay pure. + write_target = kv_cache if persist else kv_cache.clone() + main_kv_flat = main_kv.squeeze(1).to(write_target.dtype) + if ( + valid_len is None + and IS_CUTLASS_DSL_AVAILABLE + and is_fused_dspark_attention_supported( + q, main_kv_flat, kv, write_target, slots, start_pos, attn_sink + ) + ): + # One custom op performs the rolling-cache write/read, validity handling, + # QK, attention-sink online softmax, and PV. In particular it creates no + # topk index, gathered KV, score, or probability tensors. + o = cute_dsl_dspark_attention( + q, + main_kv_flat, + kv, + write_target, + slots, + start_pos, + attn_sink, + softmax_scale, + ) + else: + slot_pos = start_pos % window_size # [G] + write_target[slots, slot_pos] = main_kv_flat + cache_rows = write_target[slots] # [G, window, head_dim] + kv_full = torch.cat([cache_rows, kv], dim=1) # [G, window + block, head_dim] + topk = get_dspark_topk_idxs_batched(window_size, block, start_pos, valid_len) + o = dspark_sparse_attn( + q, kv_full, attn_sink, topk, softmax_scale + ) # [G, block, h, head_dim] + o = _rmsnorm_rope_batched( + o, + kv_norm_w, + eps, + rd, + blk_freqs, + num_heads=n_heads, + apply_weight=False, + apply_rmsnorm=False, + inverse_rope=True, + ) + + # Grouped low-rank O projection. + o = o.reshape(g, block, n_groups, -1) + wo_a_v = wo_a.view(n_groups, o_lora_rank, -1) + o = torch.einsum("bsgd,grd->bsgr", o, wo_a_v) + return F.linear(o.flatten(2), wo_b) + + +# ---------------------------------------------------------------------------- +# Block draft I/O. +# ---------------------------------------------------------------------------- + + +def build_draft_input_ids( + bonus_token_ids: torch.Tensor, *, block_size: int, noise_token_id: int +) -> torch.Tensor: + """``[batch] -> [batch, block_size]`` = ``[bonus, noise, noise, ...]``. + + The first position is the verified bonus token (the target's last accepted + token); the rest are the DSpark noise/mask token (id 128799 for V4-Pro). + """ + batch = bonus_token_ids.shape[0] + out = bonus_token_ids.new_full((batch, block_size), int(noise_token_id)) + out[:, 0] = bonus_token_ids + return out + + +def dspark_propose( + base_logits: torch.Tensor, + *, + bonus_token_ids: torch.Tensor, + block_hidden: torch.Tensor, + markov_head: Optional[nn.Module], + confidence_head: Optional[nn.Module], + block_size: int, + temperature: float = 0.0, + confidence_threshold: float = 0.0, + return_logits: bool = False, +) -> tuple: + """Produce DSpark draft tokens for one block (functional-first, static length). + + Args: + base_logits: ``[batch, block_size, vocab]`` from the backbone + lm_head. + bonus_token_ids: ``[batch]`` the token preceding the first draft position. + block_hidden: ``[batch, block_size, hidden]`` backbone hidden (feeds the + confidence head, and the RNN-head variant). + markov_head / confidence_head: the validated DSpark heads (may be None). + Returns: + draft_tokens: ``[batch, block_size]`` sampled tokens (full block; callers + keep the tensor fixed-width for CUDA-graph safety). + num_proposed: ``[batch]`` int32 — how many leading tokens survive the + static confidence-threshold truncation (== block_size when no head / + threshold<=0). + """ + batch = base_logits.shape[0] + # ``draft_logits`` are the per-position distributions the draft token is drawn + # from (markov-corrected when a head is present, else the raw base logits). + # Surfaced under ``return_logits`` for the §7.9 probabilistic-acceptance + # (1-TV) measurement; the normal path ignores them. + draft_logits = base_logits + if markov_head is not None: + draft_tokens, corrected = markov_head.sample_block_tokens( + base_logits, + first_prev_token_ids=bonus_token_ids, + hidden_states=block_hidden, + temperature=temperature, + ) + draft_logits = corrected + else: + from .modeling_speculative import greedy_or_sample + + draft_tokens = greedy_or_sample(base_logits, temperature) + + # Scaffolding: confidence-based dynamic drafting is NOT enabled in this PR. + # The worker always calls with confidence_threshold=0.0, so the block below is + # inert and num_proposed stays == block_size (the full block is proposed). The + # returned num_proposed is intentionally not yet consumed by the speculative + # scheduler/verifier; wiring it through is a follow-up (see PR description). + num_proposed = torch.full( + (batch,), int(block_size), dtype=torch.int32, device=base_logits.device + ) + if confidence_head is not None and confidence_threshold > 0.0: + # prev token at position k is [bonus, draft_0, ..., draft_{k-1}] + prev_ids = torch.cat([bonus_token_ids.unsqueeze(1), draft_tokens[:, :-1]], dim=1) + prev_emb = ( + markov_head.get_prev_embeddings(prev_ids) + if (markov_head is not None and getattr(confidence_head, "with_markov", False)) + else None + ) + conf_logits = ( + confidence_head(block_hidden, prev_embeddings=prev_emb) + if prev_emb is not None + else confidence_head(block_hidden) + ) + # Per-request prefix truncation (batch handled row-wise to stay simple; + # functional-first scope typically runs batch=1 for the draft). + for b in range(batch): + num_proposed[b] = confident_prefix_length( + conf_logits[b : b + 1], block_size=block_size, threshold=confidence_threshold + ) + if return_logits: + return draft_tokens, num_proposed, draft_logits + return draft_tokens, num_proposed + + +# ---------------------------------------------------------------------------- +# Draft backbone (``mtp.*`` stages of DeepSeek-V4 blocks). +# ---------------------------------------------------------------------------- + # Matches the draft namespace ``mtp..`` in the V4-Pro-DSpark # checkpoint. Each draft stage is a full DeepSeek-V4 block stored under this # prefix; the main model's keys (``layers.*``, ``embed.weight``, ``head.weight``, @@ -1275,4 +1920,16 @@ def _build_dspark_draft(model_config, draft_config, lm_head, model): "DSparkForCausalLM", "validate_dspark_eplb_layer_base", "validate_dspark_eplb_stage_layers", + # Captured-context attention primitives. + "get_dspark_topk_idxs", + "get_dspark_topk_idxs_batched", + "dspark_sparse_attn", + "precompute_dspark_freqs_cis", + "apply_dspark_rotary", + "apply_dspark_rotary_batched", + "dspark_attention_forward", + "dspark_attention_forward_batched", + # Block draft I/O. + "build_draft_input_ids", + "dspark_propose", ] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py index b73fadde3e5b..9651dd445d89 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py @@ -27,16 +27,20 @@ import torch import torch.nn.functional as F -import tensorrt_llm._torch.models.dspark.attention as dspark_attention import tensorrt_llm._torch.models.modeling_dspark as modeling_dspark -from tensorrt_llm._torch.models.dspark.attention import ( +from tensorrt_llm._torch.models.modeling_dspark import ( + DSparkDraftModel, apply_dspark_rotary, dspark_attention_forward, dspark_sparse_attn, get_dspark_topk_idxs, precompute_dspark_freqs_cis, ) -from tensorrt_llm._torch.models.modeling_dspark import DSparkDraftModel + +# The captured-context attention primitives were folded into +# modeling_dspark; keep the historical alias so monkeypatch targets +# below read unchanged. +dspark_attention = modeling_dspark def test_rmsnorm_rope_fallback_applies_weight_without_rmsnorm(monkeypatch): diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py index 9ada20bf4ef3..95fa963affed 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py @@ -27,7 +27,7 @@ import pytest import torch -from tensorrt_llm._torch.models.dspark.attention import ( +from tensorrt_llm._torch.models.modeling_dspark import ( apply_dspark_rotary, apply_dspark_rotary_batched, dspark_attention_forward, diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py index f31dab6cb12a..9d349ef72a2c 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py @@ -16,7 +16,7 @@ import torch -from tensorrt_llm._torch.models.dspark.draft import build_draft_input_ids, dspark_propose +from tensorrt_llm._torch.models.modeling_dspark import build_draft_input_ids, dspark_propose from tensorrt_llm._torch.models.modeling_speculative import DSparkConfidenceHead, build_markov_head VOCAB, HID, RANK, B, BLK = 257, 32, 16, 2, 5 diff --git a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py index 42a1c69ee2a9..06e35de621ce 100644 --- a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py +++ b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py @@ -7,7 +7,7 @@ import torch from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE -from tensorrt_llm._torch.models.dspark.attention import ( +from tensorrt_llm._torch.models.modeling_dspark import ( dspark_sparse_attn, get_dspark_topk_idxs_batched, ) @@ -172,7 +172,7 @@ def test_cute_dsl_dspark_attention_compiles_once_across_batch_sizes(): def test_dspark_attention_forward_batched_fused_matches_fallback(monkeypatch): - import tensorrt_llm._torch.models.dspark.attention as dspark_attention + import tensorrt_llm._torch.models.modeling_dspark as dspark_attention torch.manual_seed(17) device = torch.device("cuda") From 6a99209913a45e33703a2705dcbdb85e97e33a8a Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Wed, 19 Aug 2026 01:48:55 -0700 Subject: [PATCH 07/21] [None][chore] Prefix the DeepSeek-V4 DSpark draft classes with DSv4 DSparkForCausalLM, DSparkDraftModel and DSparkBlock are hard-wired to DeepSeek-V4: DSparkBlock derives from DeepseekV4DecoderLayer, the draft weights live in the target checkpoint's mtp.* namespace, and the stages carry EPLB layer-index alignment. The unqualified names overclaim, and they hold a name that the standalone DSpark drafters will need. The two inference/model.py citations keep the original spelling: they name DeepSpec's own class, not this one. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/models/modeling_dspark.py | 30 +++++++++---------- tensorrt_llm/_torch/models/modeling_utils.py | 2 +- tensorrt_llm/_torch/speculative/dspark.py | 12 ++++---- .../hw_agnostic/test_dspark_attention.py | 10 +++---- 4 files changed, 27 insertions(+), 27 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index aebafa552838..241514badbde 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -835,7 +835,7 @@ def count_dspark_stages(ckpt_dir: str) -> Optional[int]: def _rename_dspark_stage_subkey(rest: str, routed_scale: str) -> str: - """Map a per-stage checkpoint subkey to the ``DSparkBlock`` param subkey.""" + """Map a per-stage checkpoint subkey to the ``DSv4DSparkBlock`` param subkey.""" if rest == "attn_norm.weight": return "input_layernorm.weight" if rest == "ffn_norm.weight": @@ -855,7 +855,7 @@ def _rename_dspark_stage_subkey(rest: str, routed_scale: str) -> str: if rest.startswith("ffn."): return f"mlp.{_rename_deepseek_v4_ffn_subkey(rest[len('ffn.') :], routed_scale)}" # main_proj.weight, main_norm.weight, norm.weight, markov_head.*, - # confidence_head.* map 1:1 onto the DSparkBlock submodules. + # confidence_head.* map 1:1 onto the DSv4DSparkBlock submodules. return rest @@ -892,7 +892,7 @@ def remap_dspark_draft_keys(weights: Dict, num_stages: int) -> Dict: # dequantize ``wo_a`` (cos 1.0 vs ``wo_a_fp8 * scale``). Always dequantize now. -class DSparkBlock(DeepseekV4DecoderLayer): +class DSv4DSparkBlock(DeepseekV4DecoderLayer): """One DSpark draft stage = a DeepSeek-V4 decoder block + DSpark extras. ``stage_id`` in ``[0, num_stages)``; only stage 0 owns the capture projection @@ -984,7 +984,7 @@ def has_heads(self) -> bool: return self.stage_id == self.num_stages - 1 -class DSparkDraftModel(nn.Module): +class DSv4DSparkDraftModel(nn.Module): """The ``n_mtp_layers``-stage DSpark draft stacked on a DeepSeek-V4 target. Shares ``embed_tokens`` / ``lm_head`` with the target model. ``forward_embed`` @@ -1061,7 +1061,7 @@ def __init__( draft_model_config = self._derive_draft_model_config(model_config, base, self.num_stages) self.mtp_layers = nn.ModuleList( [ - DSparkBlock( + DSv4DSparkBlock( draft_model_config, base + s, aux_stream_dict, @@ -1199,7 +1199,7 @@ def cache_attn_weights_from_checkpoint(self, ckpt_dir: str, weight_map: Dict[str def cache_attn_weights_from_state_dict(self, weights: Dict) -> None: """Populate ``_dspark_attn`` from an already-loaded in-memory ``weights`` dict (no extra disk I/O); used on the one-engine load path - (``DSparkForCausalLM.load_weights``). Delegates to :meth:`_cache_attn_weights`. + (``DSv4DSparkForCausalLM.load_weights``). Delegates to :meth:`_cache_attn_weights`. """ self._cache_attn_weights(weights) @@ -1463,7 +1463,7 @@ def forward_embed( def _forward_stage( self, - stage: "DSparkBlock", + stage: "DSv4DSparkBlock", h: torch.Tensor, main_x: torch.Tensor, start_pos, @@ -1798,10 +1798,10 @@ def forward_head( ) -class DSparkForCausalLM(nn.Module): +class DSv4DSparkForCausalLM(nn.Module): """One-engine draft wrapper for DSpark (mirrors ``DFlashForCausalLM``). - Wraps :class:`DSparkDraftModel` (the ``n_mtp_layers``-stage ``mtp.*`` backbone) + Wraps :class:`DSv4DSparkDraftModel` (the ``n_mtp_layers``-stage ``mtp.*`` backbone) for the single-engine external-drafter flow: created by ``get_draft_model``, appended to the target's epilogue, and driven by ``DSparkWorker``. @@ -1815,7 +1815,7 @@ class DSparkForCausalLM(nn.Module): def __init__(self, draft_config, aux_stream_dict=None, num_stages=None, block_size=None): super().__init__() - self.dspark_model = DSparkDraftModel( + self.dspark_model = DSv4DSparkDraftModel( draft_config, aux_stream_dict, num_stages=num_stages, @@ -1902,11 +1902,11 @@ def _build_dspark_draft(model_config, draft_config, lm_head, model): model: the target model, whose aux streams the draft stages reuse. Returns: - The ``DSparkForCausalLM`` draft module. + The ``DSv4DSparkForCausalLM`` draft module. """ num_stages = count_dspark_stages(model_config.spec_config.speculative_model) validate_dspark_eplb_layer_base(model_config, draft_config) - return DSparkForCausalLM( + return DSv4DSparkForCausalLM( draft_config, getattr(model, "aux_stream_dict", None), num_stages=num_stages, @@ -1915,9 +1915,9 @@ def _build_dspark_draft(model_config, draft_config, lm_head, model): __all__ = [ - "DSparkBlock", - "DSparkDraftModel", - "DSparkForCausalLM", + "DSv4DSparkBlock", + "DSv4DSparkDraftModel", + "DSv4DSparkForCausalLM", "validate_dspark_eplb_layer_base", "validate_dspark_eplb_stage_layers", # Captured-context attention primitives. diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 6725b388eec2..16f621397a49 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -1042,7 +1042,7 @@ def _build_dspark_draft(model_config, draft_config, lm_head, model): num_stages = count_dspark_stages( model_config.spec_config.speculative_model) validate_dspark_eplb_layer_base(model_config, draft_config) - return DSparkForCausalLM( + return DSv4DSparkForCausalLM( draft_config, getattr(model, "aux_stream_dict", None), num_stages=num_stages, diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 6510f565fea9..6840200e5cbb 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -17,7 +17,7 @@ # hidden states, accept the previous block with standard verification, draft a # new block in one backbone forward), adapted to DSpark's draft model which # produces the whole block (and its confidence-truncated length) inside a single -# ``DSparkDraftModel.forward`` rather than via mask-token cross-attention. +# ``DSv4DSparkDraftModel.forward`` rather than via mask-token cross-attention. from collections import deque from dataclasses import dataclass @@ -44,7 +44,7 @@ class DSparkSpecMetadata(SpecMetadata): the target forward pass. DSpark captures the *mean over the multi-head (mHC) residual streams* at each captured layer (handled by the target-side capture hook), concatenated across layers, and feeds them to the draft - model's ``main_proj`` + ``main_norm`` (inside ``DSparkDraftModel.forward``) + model's ``main_proj`` + ``main_norm`` (inside ``DSv4DSparkDraftModel.forward``) as the captured-context attention input (``main_x``). Mirrors :class:`DFlashSpecMetadata`; the only DSpark-specific detail is that @@ -188,7 +188,7 @@ class DSparkWorker(SpecWorkerBase): """Worker for DSpark speculative decoding. DSpark drafts a whole block of ``block_size`` tokens in one backbone forward - (``DSparkDraftModel.forward``): it projects the captured target-layer hidden + (``DSv4DSparkDraftModel.forward``): it projects the captured target-layer hidden states (``main_proj`` + ``main_norm``) into the draft's captured-context attention, runs the ``num_stages`` DSpark blocks over a rolling captured window, refines the per-position logits with the Markov head, and predicts a @@ -205,7 +205,7 @@ class DSparkWorker(SpecWorkerBase): The rolling window is kept consistent across the whole decode: it is seeded from the prompt's captured context at prefill and back-filled with the intermediate accepted tokens of a multi-accept step (both via - ``DSparkDraftModel.write_context_windows``), in addition to the per-step bonus + ``DSv4DSparkDraftModel.write_context_windows``), in addition to the per-step bonus write done by the generation path. These affect draft acceptance rate only, not correctness, which the standard target verify guarantees. @@ -245,7 +245,7 @@ def __init__( self._scratch_slot = 0 # The generation draft path is the batched, host-sync-free - # ``_draft_gen_block_batched`` + ``DSparkDraftModel.forward_batched`` + + # ``_draft_gen_block_batched`` + ``DSv4DSparkDraftModel.forward_batched`` + # ``dspark_attention_forward_batched``: it is correct in eager mode AND safe # to capture into the target's CUDA graph (DSpark is a one-engine drafter — # its worker forward runs inside that graph, so the draft path MUST be @@ -414,7 +414,7 @@ def _draft_gen_block_batched( (``nacc``, the bonus, ``main_hidden``, ``start_pos``, the multi-accept back-fill) are gathered as tensors, slots come from the host-built ``_batch_to_slot`` mirror, and the backbone runs once via - ``DSparkDraftModel.forward_batched``. Returns the per-position corrected + ``DSv4DSparkDraftModel.forward_batched``. Returns the per-position corrected block logits ``[num_gens, K, vocab]`` (or ``None`` when there is nothing to draft); the worker feeds them to ``SpecWorkerBase.sample_draft_tokens``. Confidence truncation stays disabled — the full block is proposed. diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py index 9651dd445d89..40a841f3656f 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py @@ -29,7 +29,7 @@ import tensorrt_llm._torch.models.modeling_dspark as modeling_dspark from tensorrt_llm._torch.models.modeling_dspark import ( - DSparkDraftModel, + DSv4DSparkDraftModel, apply_dspark_rotary, dspark_attention_forward, dspark_sparse_attn, @@ -122,8 +122,8 @@ def test_rope_table_is_cached_once_per_device(): _freqs_table_cache={}, ) - first = DSparkDraftModel._dspark_freqs_table(model, torch.device("cpu")) - second = DSparkDraftModel._dspark_freqs_table(model, torch.device("cpu")) + first = DSv4DSparkDraftModel._dspark_freqs_table(model, torch.device("cpu")) + second = DSv4DSparkDraftModel._dspark_freqs_table(model, torch.device("cpu")) assert first.data_ptr() == second.data_ptr() assert len(model._freqs_table_cache) == 1 @@ -166,7 +166,7 @@ def fake_decoder_layer_init( spec_config=None, ) - block = modeling_dspark.DSparkBlock( + block = modeling_dspark.DSv4DSparkBlock( model_config, layer_idx=10, aux_stream_dict={}, @@ -248,7 +248,7 @@ def call(*args, **kwargs): ), ) - actual = DSparkDraftModel._forward_stage( + actual = DSv4DSparkDraftModel._forward_stage( model, stage, h, From cfd3641b089e17984d5c1f83f1ac62803edbb845 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Wed, 19 Aug 2026 02:59:25 -0700 Subject: [PATCH 08/21] [None][feat] Build standalone DSpark drafters from decoding_type DSpark DSpark shipped in two forms that had no common entry point: the embedded DeepSeek-V4-Pro draft used decoding_type DSpark, while a standalone drafter had to be configured as decoding_type DFlash, because DFlashForCausalLM was where the Markov head, the confidence head and the shift_label convention were implemented. Move that head set into DSparkDrafterForCausalLM, a DFlashForCausalLM subclass in modeling_dspark, and give _build_dspark_draft two levels of dispatch: the checkpoint's mtp.* namespace selects the embedded draft, otherwise the drafter's own model_type selects the backbone class (Qwen3DSparkForCausalLM today). DFlash now refuses a drafter that declares the DSpark heads instead of serving it without them, which would only show up as a lower acceptance rate. modeling_dflash keeps no reference to a DSpark class, so the new modeling_dspark -> modeling_dflash inheritance edge stays one-way. The sliding-window configuration stays in the DFlash base: the block decode indexes the resolved windows directly and must not reach for an attribute only a subclass defines. DSparkDecodingConfig gains attention_backend for the standalone path, and the checkpoint reader now accepts the dflash_config and plain top-level spellings alongside dspark_*; a knob the reader misses is not an error, it silently degrades the drafter. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/models/modeling_dflash.py | 236 ++++++-------- tensorrt_llm/_torch/models/modeling_dspark.py | 299 ++++++++++++++++-- tensorrt_llm/llmapi/llm_args.py | 48 ++- .../test_dspark_flavour_dispatch.py | 249 +++++++++++++++ .../test_kimi_k3_dspark_semantics.py | 47 +-- 5 files changed, 685 insertions(+), 194 deletions(-) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_flavour_dispatch.py diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index 79f65a679647..7f77ae8d659b 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -21,7 +21,6 @@ from ..pyexecutor.config_utils import _is_sliding_attention_layer, get_layer_attention_window from ..speculative.dflash_attention import get_dflash_flash_attention, get_dflash_trtllm_gen_ops from ..speculative.interface import SpeculativeDecodingMode -from .modeling_speculative import dspark_markov_chain_logits from .modeling_utils import get_model_architecture, register_draft_model @@ -30,16 +29,21 @@ def dspark_layer_window_size( ) -> tuple[int, int]: """flash-attn ``window_size`` for one draft layer of the block decode. - DSpark drafters (deepseek-ai/DeepSpec) run the draft block through HF - attention with ``sliding_window`` set on 'sliding_attention' layers and - is_causal=False. HF's flash path - (transformers/modeling_flash_attention_utils.py) translates that to - ``window_size = (sliding_window - 1, sliding_window - 1)``, i.e. each - query attends keys within ``swa_window - 1`` KV-index distance on both - sides. In the DFlash pool layout KV index == token position, so this - limits draft queries to the most recent ``swa_window`` context tokens - plus the (nearby) draft block. Full-attention layers and non-dspark - drafters keep flash-attn's default ``(-1, -1)`` (no window). + Sliding-window block decode was introduced by the DSpark drafters + (deepseek-ai/DeepSpec) -- hence the name -- but it is an attention + configuration, not part of the DSpark head set, so it stays in the DFlash + base: the block decode below indexes the resolved windows directly and must + not depend on a subclass-only attribute. + + Those drafters run the draft block through HF attention with + ``sliding_window`` set on 'sliding_attention' layers and is_causal=False. + HF's flash path (transformers/modeling_flash_attention_utils.py) translates + that to ``window_size = (sliding_window - 1, sliding_window - 1)``, i.e. + each query attends keys within ``swa_window - 1`` KV-index distance on both + sides. In the DFlash pool layout KV index == token position, so this limits + draft queries to the most recent ``swa_window`` context tokens plus the + (nearby) draft block. Full-attention layers and drafters that do not enable + the window keep flash-attn's default ``(-1, -1)`` (no window). """ if not use_swa: return (-1, -1) @@ -131,77 +135,25 @@ def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): f"attention_backend: {self.dflash_attention_backend}" ) - # DSpark drafters (DFlash + low-rank Markov head + confidence head, - # arXiv 2607.05147; reference: deepseek-ai/DeepSpec). The weights- - # independent drafter-forward semantics ARE implemented here: - # - vanilla Markov intra-block logit bias (applied by DFlashWorker - # through apply_markov_chain_logits), - # - sliding-window attention on 'sliding_attention' draft layers - # during the block decode (use_swa / swa_window_size), - # - the shift_label output convention (hidden state at block slot j - # predicts draft token j+1; slot 0 holds the anchor token). - # Confidence-scheduled verification is NOT implemented yet: the - # confidence_proj weights are loaded (for the follow-up MR) but never - # used, and drafting always proposes the full K tokens. - self._dspark_shift_label = bool(dflash_config.get("shift_label", False)) - self._dspark_use_swa = bool(dflash_config.get("use_swa", False)) - self._dspark_swa_window = int(dflash_config.get("swa_window_size", 0) or 0) - self._dspark_markov_rank = int(dflash_config.get("markov_rank", 0) or 0) - self._dspark_markov_head_type = str( - dflash_config.get("markov_head_type", "vanilla") or "vanilla" - ).lower() - self._dspark_use_confidence_head = bool(dflash_config.get("use_confidence_head", False)) - # Plain None placeholders rather than nn.Parameter/buffer: most - # DFlash checkpoints don't ship these heads, and their shapes - # ([vocab, rank]) are checkpoint-dependent, so nothing is - # pre-allocated. load_weights() fills them in only when the - # checkpoint ships them; consumers treat None as "head absent". - self.markov_w1 = None # [vocab, rank] (nn.Embedding weight layout) - self.markov_w2 = None # [vocab, rank] (nn.Linear(rank->vocab) weight) - self.confidence_proj_weight = None # loaded, unused (follow-up MR) - self.confidence_proj_bias = None - - if self._dspark_markov_rank > 0 and self._dspark_markov_head_type != "vanilla": + # Sliding-window block decode (use_swa / swa_window_size). Kept in the + # base class rather than in the DSpark subclass because the block decode + # below indexes ``self._layer_windows`` directly: a base-class forward + # must not reach for an attribute only a subclass defines. Drafters that + # leave use_swa unset get all-(-1,-1) windows, i.e. a no-op. + self._use_swa = bool(dflash_config.get("use_swa", False)) + self._swa_window = int(dflash_config.get("swa_window_size", 0) or 0) + if self._use_swa and self._swa_window < 1: raise ValueError( - f"DFlash dspark drafter declares markov_head_type=" - f"'{self._dspark_markov_head_type}'; only 'vanilla' is " - "supported (gated/rnn heads need per-step hidden features)." - ) - if self._dspark_use_swa and self._dspark_swa_window < 1: - raise ValueError( - "DFlash dspark drafter sets use_swa but swa_window_size=" + "DFlash drafter sets use_swa but swa_window_size=" f"{dflash_config.get('swa_window_size')} is invalid." ) - # causal=true is only invalid under the dspark convention. Legacy - # DFlash drafter configs (e.g. Laguna) also carry a causal field; - # their causality is handled by the legacy decode path - # (_sliding_layers_causal), so don't reject them here. - is_dspark = ( - str(dflash_config.get("projector_type", "") or "").lower() == "dspark" - or self._dspark_shift_label - or self._dspark_use_swa - or self._dspark_markov_rank > 0 - or self._dspark_use_confidence_head - ) - if is_dspark and dflash_config.get("causal"): - raise ValueError( - "DFlash dspark drafter sets causal=true; the block decode " - "only supports the non-causal dspark convention." - ) # Per-layer flash-attn window for the block decode, resolved once. num_draft_layers = getattr(pretrained_config, "num_hidden_layers", 0) layer_types = getattr(pretrained_config, "layer_types", None) - self._dspark_layer_windows = [ - dspark_layer_window_size(self._dspark_use_swa, self._dspark_swa_window, layer_types, i) + self._layer_windows = [ + dspark_layer_window_size(self._use_swa, self._swa_window, layer_types, i) for i in range(num_draft_layers) ] - if self._dspark_use_confidence_head: - logger.warning( - "DFlash dspark drafter declares use_confidence_head; " - "confidence-scheduled verification is not implemented yet " - "(confidence_proj weights are loaded but unused, drafting " - "always proposes the full K tokens)." - ) self.logits_processor = None # Set by caller after construction @@ -319,33 +271,6 @@ def project_target_hidden(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = hidden_states.to(self.fc.weight.dtype) return self.hidden_norm(self.fc(hidden_states)) - @property - def has_markov_head(self) -> bool: - return self._dspark_markov_rank > 0 and self.markov_w1 is not None - - def apply_markov_chain_logits( - self, - base_logits: torch.Tensor, - first_prev_tokens: torch.Tensor, - argmax_fn=None, - vocab_slice: slice | None = None, - ) -> torch.Tensor: - """Apply the dspark vanilla-Markov intra-block bias to block logits. - - No-op (returns ``base_logits`` unchanged) for non-dspark drafters. - See :func:`dspark_markov_chain` for the semantics; when - ``base_logits`` is a TP vocab shard, the caller must pass this - rank's ``vocab_slice`` (to shard the markov_w2 rows identically) - and an ``argmax_fn`` returning full-vocab token ids — DFlashWorker - handles both. - """ - if not self.has_markov_head: - return base_logits - markov_w2 = self.markov_w2 if vocab_slice is None else self.markov_w2[vocab_slice] - return dspark_markov_chain_logits( - base_logits, first_prev_tokens, self.markov_w1, markov_w2, argmax_fn=argmax_fn - ) - def _post_attention_gate(self, attn_output, gate_input, attn_mod, num_heads, head_dim): """Hook applied to the block-attention output before o_proj. @@ -392,43 +317,6 @@ def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): split[k] = v weights = split - # DSpark head weights: keep them out of the backbone remap (they'd - # get a 'model.' prefix and be dropped by allow_partial_loading). - # markov_w1/markov_w2 drive the intra-block logit bias; the - # confidence_proj weights are loaded for the confidence-scheduling - # follow-up MR but are not used yet. - dspark_keys = ( - "markov_w1.weight", - "markov_w2.weight", - "confidence_proj.weight", - "confidence_proj.bias", - ) - dspark_weights = {k: weights[k] for k in dspark_keys if k in weights} - if dspark_weights: - weights = {k: v for k, v in weights.items() if k not in dspark_weights} - if self._dspark_markov_rank > 0: - vocab = self.config.vocab_size - rank = self._dspark_markov_rank - for k in ("markov_w1.weight", "markov_w2.weight"): - if k not in dspark_weights: - raise ValueError( - f"DFlash dspark drafter declares markov_rank=" - f"{self._dspark_markov_rank} but the checkpoint is " - f"missing {k}." - ) - if tuple(dspark_weights[k].shape) != (vocab, rank): - raise ValueError( - f"DFlash dspark {k} has shape " - f"{tuple(dspark_weights[k].shape)}, expected " - f"[vocab, markov_rank] = ({vocab}, {rank})." - ) - self.markov_w1 = dspark_weights["markov_w1.weight"].to("cuda") - self.markov_w2 = dspark_weights["markov_w2.weight"].to("cuda") - if "confidence_proj.weight" in dspark_weights: - 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") - # Remap: add 'model.' prefix where needed, and extract DFlash-specific weights remapped = {} for key, value in weights.items(): @@ -1058,13 +946,11 @@ def dflash_forward( # Per-layer view into the pooled ctx cache. causal, window_size = self._get_attention_mask_args(layer_idx) - dspark_window = ( - self._dspark_layer_windows[layer_idx] - if layer_idx < len(self._dspark_layer_windows) - else (-1, -1) + swa_window = ( + self._layer_windows[layer_idx] if layer_idx < len(self._layer_windows) else (-1, -1) ) - if dspark_window != (-1, -1): - window_size = dspark_window + if swa_window != (-1, -1): + window_size = swa_window if self.dflash_attention_backend == "TRTLLM": layer_cache = ctx_kv_cache[layer_idx] trtllm_gen_ops.append_paged_kv_cache( @@ -1320,13 +1206,73 @@ def _post_attention_gate(self, attn_output, gate_input, attn_mod, num_heads, hea return (attn_output.unflatten(-1, (num_heads, head_dim)) * gate.unsqueeze(-1)).flatten(-2) +# Published DSpark drafters spell the head switches four different ways, and a +# reader that knows only one of them degrades silently: the heads are skipped, +# their weights are dropped, and nothing raises. Resolution order matches +# ``TorchLlmArgs.validate_speculative_config`` so the model and the user-visible +# spec_config can never disagree about whether a head is on. +_DSPARK_HEAD_KEY_ALIASES = { + # RadixArk/Kimi-K3-DSpark ships ``enable_confidence_head`` top-level. + "use_confidence_head": ("use_confidence_head", "enable_confidence_head"), +} + + +def resolve_dspark_head_config(pretrained_config, key): + """Resolve one DSpark head switch across every spelling in the wild. + + Looks in ``dspark_config``, then ``dflash_config``, then the top level as + ``dspark_`` and as ````, for each accepted alias of ``key``. + Returns ``None`` when the drafter declares the switch nowhere. + """ + dspark_cfg = getattr(pretrained_config, "dspark_config", None) or {} + dflash_cfg = getattr(pretrained_config, "dflash_config", None) or {} + for name in _DSPARK_HEAD_KEY_ALIASES.get(key, (key,)): + for value in ( + dspark_cfg.get(name), + dflash_cfg.get(name), + getattr(pretrained_config, f"dspark_{name}", None), + getattr(pretrained_config, name, None), + ): + if value is not None: + return value + return None + + +def declares_dspark_heads(pretrained_config) -> bool: + """True when a drafter config asks for the DSpark head set. + + Resolved here rather than in the DSpark module: this module must not import + the DSpark drafters, or the ``modeling_dspark -> modeling_dflash`` + inheritance edge would become a cycle. + """ + return bool( + str(resolve_dspark_head_config(pretrained_config, "projector_type") or "").lower() + == "dspark" + or resolve_dspark_head_config(pretrained_config, "shift_label") + or resolve_dspark_head_config(pretrained_config, "use_confidence_head") + or int(resolve_dspark_head_config(pretrained_config, "markov_rank") or 0) > 0 + ) + + @register_draft_model(SpeculativeDecodingMode.DFLASH) def _build_dflash_draft(model_config, draft_config, lm_head, model): """Build the DFlash drafter. Selects the Laguna variant by detecting its architecture in the draft - checkpoint's own config. + checkpoint's own config. A drafter that declares the DSpark heads is + rejected rather than silently served without them: DFlash no longer + implements the Markov / confidence / shift_label semantics, so building one + here would degrade the drafter with no error and show up only as a lower + acceptance rate. """ + if declares_dspark_heads(draft_config.pretrained_config): + raise ValueError( + "This drafter checkpoint declares the DSpark head set " + "(dflash_config with any of projector_type='dspark', shift_label, " + "use_confidence_head, markov_rank > 0), which decoding_type " + "'DFlash' does not implement. Set speculative_config.decoding_type " + "to 'DSpark' to run it." + ) draft_arches = getattr(draft_config.pretrained_config, "architectures", None) or [] dflash_attention_backend = model_config.spec_config.attention_backend if any("Laguna" in arch for arch in draft_arches): diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 241514badbde..47bd600b4935 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -25,18 +25,28 @@ # # The draft I/O stages are ported from the same reference # (`inference/model.py`: DSparkBlock.forward_embed / forward_head). -"""DeepSeek-V4-Pro DSpark speculative-decoding draft. +"""DSpark speculative-decoding drafters. ``DSpark`` names the speculative-decoding *algorithm*: a parallel block draft over captured target hidden states, refined by a low-rank Markov head and -scheduled by a confidence head. This module holds its **in-checkpoint** flavour, -where the draft weights ship inside the DeepSeek-V4-Pro *target* checkpoint under -the ``mtp.*`` namespace and reuse the V4 decoder block, so the draft inherits the -target's EPLB layer namespace and fp8/NVFP4 quantization. Standalone DSpark -drafters — shipped as their own checkpoint, backbone resolved through the model -registry — are built by :mod:`modeling_dflash` instead. +scheduled by a confidence head. Every ``decoding_type: DSpark`` drafter is built +here, in one of two flavours that differ only in how the draft is delivered: -Three parts live here: +* **embedded** — the draft weights ship inside the DeepSeek-V4-Pro *target* + checkpoint under the ``mtp.*`` namespace and reuse the V4 decoder block, so the + draft inherits the target's EPLB layer namespace and fp8/NVFP4 quantization. + Most of this module is this flavour. +* **standalone** — the drafter ships as its own checkpoint and shares nothing + with the target but the vocabulary and the captured hidden states. Its block + decode is DFlash's, so :class:`DSparkDrafterForCausalLM` subclasses + ``DFlashForCausalLM`` and adds the Markov head, the confidence head and the + shift_label convention. See "Standalone DSpark drafters" near the bottom. + +The dependency runs one way, ``modeling_dspark -> modeling_dflash``: DFlash is +DSpark minus those three heads, and :mod:`modeling_dflash` must not name a DSpark +class or the edge would become a cycle. + +Four parts live here: 1. **Draft backbone** — ``n_mtp_layers`` (3 for V4-Pro) full DeepSeek-V4 blocks (MLA attention + MoE + manifold Hyper-Connections), plus: @@ -57,6 +67,10 @@ ``[bonus_token, noise, ...]`` block input) and ``dspark_propose`` (Markov refinement + static confidence truncation). +4. **Standalone DSpark drafters** — :class:`DSparkDrafterForCausalLM` and its + per-backbone subclasses, plus the two-level ``_build_dspark_draft`` dispatch + that picks between the two flavours. + The per-stage *backbone* forward (block attention whose K/V derive from ``main_x``, + MoE + mHC) is brought up and numerically validated against the real fp8 weights separately; ``forward_embed`` (capture) and ``forward_head`` (block @@ -94,7 +108,13 @@ _rename_deepseek_v4_attn_subkey, _rename_deepseek_v4_ffn_subkey, ) -from .modeling_speculative import DSparkConfidenceHead, build_markov_head, confident_prefix_length +from .modeling_dflash import DFlashForCausalLM, resolve_dspark_head_config +from .modeling_speculative import ( + DSparkConfidenceHead, + build_markov_head, + confident_prefix_length, + dspark_markov_chain_logits, +) from .modeling_utils import register_draft_model if IS_CUTLASS_DSL_AVAILABLE: @@ -1888,12 +1908,230 @@ def load_weights_from_target_model(self, target_model): self.dspark_model.lm_head = target_model.lm_head +# ---------------------------------------------------------------------------- +# Standalone DSpark drafters. +# +# The other DSpark flavour: the drafter ships as its own checkpoint instead of +# living in the target's ``mtp.*`` namespace, so it shares nothing with the +# target but the vocabulary and the captured hidden states. Its block decode is +# DFlash's -- DSpark is DFlash plus a Markov logit bias, a confidence head and +# the shift_label slot convention -- so these subclass ``DFlashForCausalLM`` and +# add exactly those three. +# ---------------------------------------------------------------------------- + +# Both published drafters (RadixArk/Kimi-K3-DSpark, Inferact/Kimi-K3-DSpark) +# name the head tensors after the submodules that own them: markov_head is a +# VanillaMarkov, confidence_head an AcceptRatePredictor whose linear is ``proj``. +# The bare spellings are kept for drafters exported without that nesting. +_DSPARK_HEAD_WEIGHT_ALIASES = { + "markov_w1.weight": ("markov_head.markov_w1.weight", "markov_w1.weight"), + "markov_w2.weight": ("markov_head.markov_w2.weight", "markov_w2.weight"), + "confidence_proj.weight": ("confidence_head.proj.weight", "confidence_proj.weight"), + "confidence_proj.bias": ("confidence_head.proj.bias", "confidence_proj.bias"), +} + + +class DSparkDrafterForCausalLM(DFlashForCausalLM): + """DSpark drafter built from a standalone draft checkpoint. + + Adds the DSpark head set on top of the generic DFlash block decode: + + - the vanilla Markov intra-block logit bias, applied by ``DFlashWorker`` + 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. + + Confidence-scheduled verification is not implemented yet: ``confidence_proj`` + is loaded but unused, and drafting always proposes the full K tokens. + + The draft backbone itself is whatever the drafter config resolves to through + the model registry, so this class is backbone-agnostic; per-backbone + subclasses exist to carry backbone-specific block-decode overrides. + + 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) + + cfg = draft_config.pretrained_config + # 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) + # runs one slot past the block and reads the next request's anchor. + # An explicit false still selects the legacy layout. + shift_label = resolve_dspark_head_config(cfg, "shift_label") + self._dspark_shift_label = True if shift_label is None else bool(shift_label) + self._dspark_markov_rank = int(resolve_dspark_head_config(cfg, "markov_rank") or 0) + self._dspark_markov_head_type = str( + resolve_dspark_head_config(cfg, "markov_head_type") or "vanilla" + ).lower() + self._dspark_use_confidence_head = bool( + resolve_dspark_head_config(cfg, "use_confidence_head") or False + ) + # Plain None placeholders rather than nn.Parameter/buffer: the shapes + # ([vocab, rank]) are checkpoint-dependent, so nothing is pre-allocated + # and nothing is constructed in the module's default dtype. Using the + # checkpoint tensor as-is is also what keeps the head in the checkpoint's + # dtype instead of an nn.Module default. load_weights() fills them in; + # consumers treat None as "head absent". + self.markov_w1 = None # [vocab, rank] (nn.Embedding weight layout) + self.markov_w2 = None # [vocab, rank] (nn.Linear(rank->vocab) weight) + self.confidence_proj_weight = None # loaded, unused (follow-up MR) + self.confidence_proj_bias = None + + if self._dspark_markov_rank > 0 and self._dspark_markov_head_type != "vanilla": + raise ValueError( + f"DSpark drafter declares markov_head_type=" + f"'{self._dspark_markov_head_type}'; only 'vanilla' is " + "supported (gated/rnn heads need per-step hidden features)." + ) + # The block decode only supports the non-causal DSpark convention. + # Legacy DFlash drafter configs (e.g. Laguna) also carry a causal field + # and handle it in the legacy decode path, which is why this check lives + # here rather than in the DFlash base. + if resolve_dspark_head_config(cfg, "causal"): + raise ValueError( + "DSpark drafter sets causal=true; the block decode only " + "supports the non-causal DSpark convention." + ) + if self._dspark_use_confidence_head: + logger.warning( + "DSpark drafter declares use_confidence_head; " + "confidence-scheduled verification is not implemented yet " + "(confidence_proj weights are loaded but unused, drafting " + "always proposes the full K tokens)." + ) + + @property + def has_markov_head(self) -> bool: + return self._dspark_markov_rank > 0 and self.markov_w1 is not None + + def apply_markov_chain_logits( + self, + base_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + argmax_fn=None, + vocab_slice: Optional[slice] = None, + ) -> torch.Tensor: + """Apply the vanilla-Markov intra-block bias to block logits. + + No-op (returns ``base_logits`` unchanged) when the checkpoint ships no + Markov head. See :func:`dspark_markov_chain` for the semantics; when + ``base_logits`` is a TP vocab shard the caller must pass this rank's + ``vocab_slice`` (to shard the markov_w2 rows identically) and an + ``argmax_fn`` returning full-vocab token ids -- ``DFlashWorker`` handles + both. + """ + if not self.has_markov_head: + return base_logits + markov_w2 = self.markov_w2 if vocab_slice is None else self.markov_w2[vocab_slice] + return dspark_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. + + 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. + """ + dspark_weights = {} + consumed = set() + for canonical, aliases in _DSPARK_HEAD_WEIGHT_ALIASES.items(): + for name in aliases: + if name in weights: + dspark_weights[canonical] = weights[name] + consumed.add(name) + break + if consumed: + weights = {k: v for k, v in weights.items() if k not in consumed} + # The inverse of the missing-weights check below. Without it, a config + # whose head switches this build cannot resolve loads the drafter with + # the heads silently dropped -- correct output, lower acceptance. + if self._dspark_markov_rank <= 0 and "markov_w1.weight" in dspark_weights: + raise ValueError( + "DSpark drafter ships markov_w1/markov_w2 but markov_rank resolved to 0. " + "The checkpoint's head switches were not found in dspark_config, " + "dflash_config, or at the top level; loading it would drop the Markov " + "head silently." + ) + if self._dspark_markov_rank > 0: + vocab = self.config.vocab_size + rank = self._dspark_markov_rank + for k in ("markov_w1.weight", "markov_w2.weight"): + if k not in dspark_weights: + raise ValueError( + f"DSpark drafter declares markov_rank=" + f"{self._dspark_markov_rank} but the checkpoint is " + f"missing {k}." + ) + if tuple(dspark_weights[k].shape) != (vocab, rank): + raise ValueError( + f"DSpark {k} has shape " + f"{tuple(dspark_weights[k].shape)}, expected " + f"[vocab, markov_rank] = ({vocab}, {rank})." + ) + self.markov_w1 = dspark_weights["markov_w1.weight"].to("cuda") + self.markov_w2 = dspark_weights["markov_w2.weight"].to("cuda") + if "confidence_proj.weight" in dspark_weights: + 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 super().load_weights(weights, weight_mapper=weight_mapper, **kwargs) + + +class Qwen3DSparkForCausalLM(DSparkDrafterForCausalLM): + """DSpark drafter on a Qwen3-style GQA draft backbone. + + Overrides nothing today: the backbone is built from the drafter config + through the model registry, and the DSpark head set is backbone-independent, + so ``DSparkDrafterForCausalLM`` already covers this combination end to end. + + It exists as the explicit dispatch target for ``model_type: qwen3``, which + keeps the supported matrix visible in class names rather than buried in a + builder, and as the seat for backbone-specific overrides when they arrive. + They will: an MLA-backboned drafter cannot reuse this block decode, which + assumes a fused ``qkv_proj`` and one uniform head dim across Q/K/V. + """ + + +# Standalone DSpark drafters by the draft checkpoint's ``model_type``. The key +# is the backbone family, not the target model: the same drafter class serves +# any target, and a target-specific one would have nothing to hold -- the +# target-side half of DSpark is the hidden-state capture, which lives in each +# target's own modeling file. +_DSPARK_DRAFTERS_BY_MODEL_TYPE = { + "qwen3": Qwen3DSparkForCausalLM, +} + + +def draft_is_embedded_in_target(model_config, draft_config) -> bool: + """True when the DSpark draft weights live inside the target checkpoint. + + That is the DeepSeek-V4-Pro layout: the draft is ``mtp.*`` inside the target + checkpoint and inherits its block definition, EPLB layer namespace and + quantization. The probe is the weight index rather than a config field + because the index is authoritative and cannot be left unset; the model_type + check is the fallback for a checkpoint whose index file is absent. + """ + if count_dspark_stages(model_config.spec_config.speculative_model) is not None: + return True + return getattr(draft_config.pretrained_config, "model_type", None) == "deepseek_v4" + + @register_draft_model(SpeculativeDecodingMode.DSPARK) def _build_dspark_draft(model_config, draft_config, lm_head, model): - """Build the DSpark drafter, reusing the target's aux streams. + """Build the DSpark drafter for either flavour. - The draft stage count (``n_mtp_layers``) is not in the HF config, so it is - derived from the checkpoint's ``mtp.*`` namespace. + Two levels of dispatch: + + 1. Are the draft weights embedded in the target checkpoint? If so this is + the DeepSeek-V4-Pro draft, whose stage count (``n_mtp_layers``) is not in + the HF config and is derived from the ``mtp.*`` namespace. + 2. Otherwise the drafter is standalone, and its own ``model_type`` selects + the backbone-specific class. Args: model_config: the target engine's ``ModelConfig``. @@ -1902,22 +2140,45 @@ def _build_dspark_draft(model_config, draft_config, lm_head, model): model: the target model, whose aux streams the draft stages reuse. Returns: - The ``DSv4DSparkForCausalLM`` draft module. + The draft ``nn.Module`` for this drafter. """ - num_stages = count_dspark_stages(model_config.spec_config.speculative_model) - validate_dspark_eplb_layer_base(model_config, draft_config) - return DSv4DSparkForCausalLM( + if draft_is_embedded_in_target(model_config, draft_config): + num_stages = count_dspark_stages(model_config.spec_config.speculative_model) + validate_dspark_eplb_layer_base(model_config, draft_config) + return DSv4DSparkForCausalLM( + draft_config, + getattr(model, "aux_stream_dict", None), + num_stages=num_stages, + block_size=model_config.spec_config.block_size, + ) + + model_type = getattr(draft_config.pretrained_config, "model_type", None) + drafter_cls = _DSPARK_DRAFTERS_BY_MODEL_TYPE.get(model_type) + if drafter_cls is None: + supported = ", ".join(sorted(_DSPARK_DRAFTERS_BY_MODEL_TYPE)) + raise NotImplementedError( + f"No standalone DSpark drafter for draft model_type {model_type!r}. " + f"Supported draft model_type values: {supported}. The dispatch keys " + "on the draft backbone, so a drafter is supported once its backbone " + "has a block-decode implementation here; MLA-backboned drafters " + "(e.g. Inferact/Kimi-K3-DSpark, model_type 'k3_dspark') need one and " + "are a follow-up." + ) + return drafter_cls( draft_config, - getattr(model, "aux_stream_dict", None), - num_stages=num_stages, - block_size=model_config.spec_config.block_size, + dflash_attention_backend=model_config.spec_config.attention_backend, ) __all__ = [ + # Embedded (DeepSeek-V4-Pro) flavour. "DSv4DSparkBlock", "DSv4DSparkDraftModel", "DSv4DSparkForCausalLM", + # Standalone flavour. + "DSparkDrafterForCausalLM", + "Qwen3DSparkForCausalLM", + "draft_is_embedded_in_target", "validate_dspark_eplb_layer_base", "validate_dspark_eplb_stage_layers", # Captured-context attention primitives. diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 3b95780bd822..a57597c286a3 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2965,6 +2965,19 @@ 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.") + @model_validator(mode="after") def set_max_total_draft_tokens(self): self.max_total_draft_tokens = self.max_draft_len @@ -5993,26 +6006,39 @@ def validate_speculative_config(self): if spec_cfg.speculative_model is None: raise ValueError( "DSpark requires speculative_config.speculative_model " - "to point at the checkpoint directory containing the " - "mtp.* draft weights (for DeepSeek-V4-Pro-DSpark this " - "is the target checkpoint directory itself).") + "to point at the drafter's checkpoint directory: a " + "standalone DSpark drafter repository, or -- for the " + "embedded DeepSeek-V4-Pro flavour, whose draft weights " + "live in the mtp.* namespace -- the target checkpoint " + "directory itself.") # Resolve target_layer_ids / mask_token_id / block_size / # markov_rank from the draft (or main) model config if not set. - # DSpark ships these as top-level ``dspark_*`` keys in the - # DeepSeek-V4-Pro config.json; also accept a nested - # ``dspark_config`` dict for forward compatibility. + # Three checkpoint spellings are in the wild for the same knobs + # and all are accepted here, because a key the reader misses is + # not an error -- it silently falls back to a default and the + # drafter degrades (a markov_rank read as 0 skips the Markov + # head entirely, costing acceptance with no warning): + # - top-level ``dspark_*`` (DeepSeek-V4-Pro) + # - nested ``dflash_config`` (SpecForge / RadixArk drafters) + # - nested ``dspark_config`` (forward compatibility) + # - plain top-level keys (TorchSpec drafters) draft_config_path = os.path.join(spec_cfg.speculative_model, "config.json") if os.path.exists(draft_config_path): with open(draft_config_path) as f: draft_cfg = json.load(f) - dspark_cfg = draft_cfg.get("dspark_config", {}) + dspark_cfg = draft_cfg.get("dspark_config") or {} + dflash_cfg = draft_cfg.get("dflash_config") or {} def _dspark_get(key, top_level_key): - value = dspark_cfg.get(key) - if value is None: - value = draft_cfg.get(top_level_key) - return value + for source, name in ((dspark_cfg, key), (dflash_cfg, + key), + (draft_cfg, + top_level_key), (draft_cfg, key)): + value = source.get(name) + if value is not None: + return value + return None # The checkpoint's ``dspark_target_layer_ids`` is # authoritative: it fixes both which target hidden states diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_flavour_dispatch.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_flavour_dispatch.py new file mode 100644 index 000000000000..72f171dedcfb --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_flavour_dispatch.py @@ -0,0 +1,249 @@ +# 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. +"""Which DSpark drafter ``decoding_type: DSpark`` builds, and what DFlash refuses. + +``DSpark`` ships in two flavours -- embedded in the target checkpoint +(DeepSeek-V4-Pro's ``mtp.*``) or standalone with its own checkpoint -- and one +builder picks between them, then picks the standalone backbone by the draft +checkpoint's ``model_type``. + +The DFlash half matters just as much: DFlash no longer implements the DSpark +head set, so a drafter that declares it must be refused rather than served +without it. Silently dropping the Markov head does not fail anything; it lowers +the acceptance rate, which no test would attribute back to this function. + +Everything here asserts *which class is selected*, never the object it builds: +constructing a real drafter needs GPUs and checkpoints. +""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.models import modeling_dflash, modeling_dspark +from tensorrt_llm._torch.models.modeling_dflash import declares_dspark_heads + +_DSV4_SENTINEL = object() +_QWEN3_SENTINEL = object() +_DFLASH_SENTINEL = object() +_LAGUNA_SENTINEL = object() + +_DSPARK_HEADS = { + "markov_rank": 256, + "markov_head_type": "vanilla", + "use_confidence_head": True, + "shift_label": True, + "projector_type": "dspark", +} + + +def _configs( + *, model_type="qwen3", dflash_config=None, architectures=None, attention_backend="TRTLLM" +): + """Duck-typed (target ModelConfig, draft ModelConfig) for dispatch-only asserts.""" + model_config = SimpleNamespace( + spec_config=SimpleNamespace( + speculative_model="/nonexistent/drafter", + block_size=7, + attention_backend=attention_backend, + ) + ) + draft_config = SimpleNamespace( + pretrained_config=SimpleNamespace( + model_type=model_type, + architectures=architectures, + dflash_config=dflash_config, + ) + ) + return model_config, draft_config + + +@pytest.fixture +def stub_dspark(monkeypatch): + """Replace the DSpark drafter classes with sentinel-returning stubs.""" + monkeypatch.setattr(modeling_dspark, "DSv4DSparkForCausalLM", lambda *a, **k: _DSV4_SENTINEL) + monkeypatch.setattr( + modeling_dspark, + "_DSPARK_DRAFTERS_BY_MODEL_TYPE", + {"qwen3": lambda *a, **k: _QWEN3_SENTINEL}, + ) + monkeypatch.setattr(modeling_dspark, "validate_dspark_eplb_layer_base", lambda *a, **k: None) + + +@pytest.fixture +def stub_dflash(monkeypatch): + monkeypatch.setattr(modeling_dflash, "DFlashForCausalLM", lambda *a, **k: _DFLASH_SENTINEL) + monkeypatch.setattr( + modeling_dflash, "DFlashLagunaForCausalLM", lambda *a, **k: _LAGUNA_SENTINEL + ) + + +def _embedded(monkeypatch, stages): + """Force the level-1 probe: ``stages`` is the mtp.* count, None if standalone.""" + monkeypatch.setattr(modeling_dspark, "count_dspark_stages", lambda _p: stages) + + +# -------------------------------------------------------------------------- +# decoding_type: DSpark +# -------------------------------------------------------------------------- + + +def test_standalone_qwen3_drafter_selects_qwen3_dspark(monkeypatch, stub_dspark): + _embedded(monkeypatch, None) + model_config, draft_config = _configs(model_type="qwen3") + + built = modeling_dspark._build_dspark_draft(model_config, draft_config, None, None) + + assert built is _QWEN3_SENTINEL + + +def test_embedded_draft_selects_dsv4_dspark(monkeypatch, stub_dspark): + # The mtp.* namespace in the checkpoint index is the level-1 probe. + _embedded(monkeypatch, 3) + model_config, draft_config = _configs(model_type="deepseek_v4") + + built = modeling_dspark._build_dspark_draft(model_config, draft_config, None, None) + + assert built is _DSV4_SENTINEL + + +def test_deepseek_v4_without_weight_index_still_selects_dsv4(monkeypatch, stub_dspark): + # Fallback arm of the probe: a V4 checkpoint whose index file is absent must + # not fall through to the standalone lineage, which has no V4 drafter. + _embedded(monkeypatch, None) + model_config, draft_config = _configs(model_type="deepseek_v4") + + built = modeling_dspark._build_dspark_draft(model_config, draft_config, None, None) + + assert built is _DSV4_SENTINEL + + +def test_unknown_standalone_model_type_lists_supported(monkeypatch, stub_dspark): + _embedded(monkeypatch, None) + model_config, draft_config = _configs(model_type="llama") + + with pytest.raises(NotImplementedError) as excinfo: + modeling_dspark._build_dspark_draft(model_config, draft_config, None, None) + + message = str(excinfo.value) + assert "llama" in message + assert "qwen3" in message, "the error must list the supported draft model_type values" + + +def test_standalone_drafter_receives_the_attention_backend(monkeypatch, stub_dspark): + _embedded(monkeypatch, None) + seen = {} + + def _capture(draft_config, *, dflash_attention_backend): + seen["backend"] = dflash_attention_backend + return _QWEN3_SENTINEL + + monkeypatch.setattr(modeling_dspark, "_DSPARK_DRAFTERS_BY_MODEL_TYPE", {"qwen3": _capture}) + model_config, draft_config = _configs(attention_backend="TRTLLM") + + modeling_dspark._build_dspark_draft(model_config, draft_config, None, None) + + assert seen["backend"] == "TRTLLM" + + +# -------------------------------------------------------------------------- +# decoding_type: DFlash +# -------------------------------------------------------------------------- + + +def test_dflash_refuses_a_dspark_drafter(stub_dflash): + model_config, draft_config = _configs(dflash_config=dict(_DSPARK_HEADS)) + + with pytest.raises(ValueError) as excinfo: + modeling_dflash._build_dflash_draft(model_config, draft_config, None, None) + + message = str(excinfo.value) + assert "DSpark" in message + assert "decoding_type" in message, "the error must say how to fix the config" + + +@pytest.mark.parametrize( + "field,value", + [ + ("markov_rank", 256), + ("use_confidence_head", True), + ("shift_label", True), + ("projector_type", "dspark"), + ], +) +def test_any_single_dspark_field_is_enough_to_refuse(stub_dflash, field, value): + # Each field alone means the drafter was trained under the DSpark + # convention; serving it as plain DFlash degrades it silently. + model_config, draft_config = _configs(dflash_config={"mask_token_id": 7, field: value}) + + with pytest.raises(ValueError, match="DSpark"): + modeling_dflash._build_dflash_draft(model_config, draft_config, None, None) + + +def test_plain_dflash_drafter_is_unchanged(stub_dflash): + model_config, draft_config = _configs( + dflash_config={"mask_token_id": 7, "target_layer_ids": [0, 1]} + ) + + built = modeling_dflash._build_dflash_draft(model_config, draft_config, None, None) + + assert built is _DFLASH_SENTINEL + + +def test_laguna_drafter_is_unchanged(stub_dflash): + model_config, draft_config = _configs( + architectures=["DFlashLagunaForCausalLM"], + dflash_config={"mask_token_id": 7}, + ) + + built = modeling_dflash._build_dflash_draft(model_config, draft_config, None, None) + + assert built is _LAGUNA_SENTINEL + + +def test_legacy_causal_dflash_config_is_not_mistaken_for_dspark(stub_dflash): + # Laguna configs carry ``causal`` without any DSpark field; the legacy + # decode path handles it, so this must not trip the refusal. + model_config, draft_config = _configs(dflash_config={"mask_token_id": 7, "causal": True}) + + built = modeling_dflash._build_dflash_draft(model_config, draft_config, None, None) + + assert built is _DFLASH_SENTINEL + + +# -------------------------------------------------------------------------- +# The predicate itself +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "dflash_config,expected", + [ + (None, False), + ({}, False), + ({"mask_token_id": 7}, False), + ({"causal": True}, False), + ({"markov_rank": 0}, False), + ({"shift_label": False}, False), + ({"markov_rank": 256}, True), + ({"shift_label": True}, True), + ({"use_confidence_head": True}, True), + ({"projector_type": "dspark"}, True), + ({"projector_type": "DSpark"}, True), + ], +) +def test_declares_dspark_heads(dflash_config, expected): + config = SimpleNamespace(dflash_config=dflash_config) + assert declares_dspark_heads(config) is expected 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 3768d4198efe..e22aea1e7442 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 @@ -24,6 +24,7 @@ import torch.nn.functional as F from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM, dspark_layer_window_size +from tensorrt_llm._torch.models.modeling_dspark import Qwen3DSparkForCausalLM from tensorrt_llm._torch.models.modeling_speculative import ( dspark_markov_chain_logits, dspark_markov_step_bias, @@ -263,7 +264,9 @@ def _build_drafter(dspark: bool, weights): from tensorrt_llm._torch.model_config import ModelConfig model_config = ModelConfig(pretrained_config=_tiny_config(dspark), attn_backend="TRTLLM") - drafter = DFlashForCausalLM(model_config).to("cuda") + # The DSpark head set lives in the DSpark drafter, not in the DFlash base. + drafter_cls = Qwen3DSparkForCausalLM if dspark else DFlashForCausalLM + drafter = drafter_cls(model_config).to("cuda") # Drop dspark head tensors for the plain drafter (schema without them). if not dspark: weights = {k: v for k, v in weights.items() if not k.startswith(("markov_", "confidence_"))} @@ -355,9 +358,9 @@ def _has_flash_attn(): def test_dspark_drafter_loads_head_weights_and_parses_config(): weights = _tiny_weights() drafter = _build_drafter(True, weights) - assert drafter._dspark_shift_label and drafter._dspark_use_swa - assert drafter._dspark_swa_window == SWA_WINDOW - assert drafter._dspark_layer_windows == [(SWA_WINDOW - 1, SWA_WINDOW - 1)] * 2 + assert drafter._dspark_shift_label and drafter._use_swa + assert drafter._swa_window == SWA_WINDOW + assert drafter._layer_windows == [(SWA_WINDOW - 1, SWA_WINDOW - 1)] * 2 assert drafter.has_markov_head torch.testing.assert_close(drafter.markov_w1.cpu(), weights["markov_w1.weight"]) torch.testing.assert_close(drafter.markov_w2.cpu(), weights["markov_w2.weight"]) @@ -371,16 +374,22 @@ def test_dspark_drafter_loads_head_weights_and_parses_config(): @needs_gpu def test_plain_dflash_drafter_keeps_old_gates(): """No-regression: a config WITHOUT dspark fields resolves to the exact - old code path (no window, no markov, mask slots 1..K).""" + old code path (no window, no markov, mask slots 1..K). + + The DFlash base no longer carries the DSpark head set at all, so the + assertions are that those attributes are absent rather than inert. + """ drafter = _build_drafter(False, _tiny_weights()) - assert not drafter._dspark_shift_label - assert not drafter._dspark_use_swa - assert drafter._dspark_layer_windows == [(-1, -1)] * 2 - assert not drafter.has_markov_head - assert drafter.markov_w1 is None and drafter.confidence_proj_weight is None - x = torch.randn(2, 3, VOCAB) - t = torch.zeros(2, dtype=torch.long) - assert drafter.apply_markov_chain_logits(x, t) is x + assert not drafter._use_swa + assert drafter._layer_windows == [(-1, -1)] * 2 + for absent in ( + "_dspark_shift_label", + "has_markov_head", + "markov_w1", + "confidence_proj_weight", + "apply_markov_chain_logits", + ): + assert not hasattr(drafter, absent), f"DFlash base still carries {absent}" @needs_gpu @@ -393,8 +402,8 @@ def test_legacy_causal_dflash_config_constructs(): cfg = _tiny_config(False) cfg.dflash_config = dict(cfg.dflash_config, causal=True) drafter = DFlashForCausalLM(ModelConfig(pretrained_config=cfg, attn_backend="TRTLLM")) - assert drafter._dspark_layer_windows == [(-1, -1)] * 2 - assert not drafter.has_markov_head + assert drafter._layer_windows == [(-1, -1)] * 2 + assert not hasattr(drafter, "has_markov_head") @needs_gpu @@ -404,8 +413,8 @@ def test_dspark_causal_config_rejected(): cfg = _tiny_config(True) cfg.dflash_config = dict(cfg.dflash_config, causal=True) - with pytest.raises(ValueError, match="non-causal dspark convention"): - DFlashForCausalLM(ModelConfig(pretrained_config=cfg, attn_backend="TRTLLM")) + with pytest.raises(ValueError, match="non-causal DSpark convention"): + Qwen3DSparkForCausalLM(ModelConfig(pretrained_config=cfg, attn_backend="TRTLLM")) @needs_gpu @@ -416,8 +425,8 @@ def test_dspark_projector_type_alone_rejects_causal(): cfg = _tiny_config(False) cfg.dflash_config = dict(cfg.dflash_config, projector_type="dspark", causal=True) - with pytest.raises(ValueError, match="non-causal dspark convention"): - DFlashForCausalLM(ModelConfig(pretrained_config=cfg, attn_backend="TRTLLM")) + with pytest.raises(ValueError, match="non-causal DSpark convention"): + Qwen3DSparkForCausalLM(ModelConfig(pretrained_config=cfg, attn_backend="TRTLLM")) def _run_block_decode(drafter, weights, captured, noise_embed): From 0e8c0ff05338fb0967d70a4282a49e777804f35a Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Wed, 19 Aug 2026 04:27:08 -0700 Subject: [PATCH 09/21] [None][fix] Admit DSpark in the Kimi K3 spec-dec mode gate The gate whitelisted SA and DFlash only, so a K3 engine configured with decoding_type DSpark aborted at model construction. The target side is identical for both modes -- the hidden-state capture in KimiLinearModel.forward is unconditional. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/models/modeling_kimi_linear.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index c22ece009bff..077f714001a7 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -1958,16 +1958,20 @@ def __init__(self, model_config: ModelConfig): # separate dense checkpoint (K2.7-Code-DFlash schema) consumed by # the generic DFlashForCausalLM wrapper, and the target only has # to expose per-layer hidden states via maybe_capture_hidden_states - # (see KimiLinearModel.forward). No trained K3 drafter exists yet; - # this path is exercised with synthetic weights - # (examples/kimi_k3/make_synthetic_dflash_drafter.py). + # (see KimiLinearModel.forward). + # - DSpark: the same external-drafter flow with the Markov and + # confidence heads enabled (RadixArk/Kimi-K3-DSpark and friends). + # The target side is identical -- the capture in + # KimiLinearModel.forward is unconditional -- so this gate is the + # only place the mode has to be admitted. # Modes needing draft heads (MTP/Eagle) are blocked until a # draft-head checkpoint exists. assert ( spec_config is None or spec_config.spec_dec_mode.is_sa() or spec_config.spec_dec_mode.is_dflash() - ), "Kimi K3 supports speculative decoding only with SA or DFlash" + or spec_config.spec_dec_mode.is_dspark() + ), "Kimi K3 supports speculative decoding only with SA, DFlash or DSpark" super().__init__( KimiLinearModel(model_config), model_config, From 1f6233b54656171387845b153fc93a36a52757e9 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Wed, 19 Aug 2026 08:07:35 -0700 Subject: [PATCH 10/21] [None][fix] Route the DSpark worker and metadata by draft form decoding_type DSpark now builds one of two draft models -- the embedded DeepSeek-V4-Pro draft or a standalone drafter -- but worker and spec metadata selection stayed one-to-one on the mode. A standalone drafter therefore reached DSparkWorker, which reads V4-draft-only attributes, and died on first contact with AttributeError: 'Qwen3DSparkForCausalLM' object has no attribute 'num_stages' Every dispatch that has to tell the two apart now reads one resolved-once flag on DSparkDecodingConfig, so a builder and a worker cannot disagree. The flag lives on the config because _torch/speculative imports nothing from _torch/models. Standalone drafters also stop inheriting the target's EPLB namespace, which only the embedded draft's stages belong to. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/models/modeling_dspark.py | 18 +- .../_torch/models/modeling_speculative.py | 7 +- tensorrt_llm/_torch/speculative/interface.py | 10 +- tensorrt_llm/_torch/speculative/utils.py | 17 +- tensorrt_llm/llmapi/llm_args.py | 61 ++++++ .../hw_agnostic/test_dspark_eplb_config.py | 19 +- .../test_dspark_flavour_dispatch.py | 202 +++++++++++++++--- 7 files changed, 291 insertions(+), 43 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 47bd600b4935..2d1a2bac6d00 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -2107,18 +2107,20 @@ class Qwen3DSparkForCausalLM(DSparkDrafterForCausalLM): } -def draft_is_embedded_in_target(model_config, draft_config) -> bool: +def draft_is_embedded_in_target(model_config) -> bool: """True when the DSpark draft weights live inside the target checkpoint. That is the DeepSeek-V4-Pro layout: the draft is ``mtp.*`` inside the target checkpoint and inherits its block definition, EPLB layer namespace and - quantization. The probe is the weight index rather than a config field - because the index is authoritative and cannot be left unset; the model_type - check is the fallback for a checkpoint whose index file is absent. + quantization. + + The answer comes from ``DSparkDecodingConfig.draft_is_embedded_in_target`` + rather than being re-derived here, because the worker and the spec metadata + have to make the same call from ``_torch/speculative/`` -- which cannot + import this package -- and a builder that disagreed with them would hand + the worker a draft model whose attributes it does not have. """ - if count_dspark_stages(model_config.spec_config.speculative_model) is not None: - return True - return getattr(draft_config.pretrained_config, "model_type", None) == "deepseek_v4" + return bool(model_config.spec_config.draft_is_embedded_in_target) @register_draft_model(SpeculativeDecodingMode.DSPARK) @@ -2142,7 +2144,7 @@ def _build_dspark_draft(model_config, draft_config, lm_head, model): Returns: The draft ``nn.Module`` for this drafter. """ - if draft_is_embedded_in_target(model_config, draft_config): + if draft_is_embedded_in_target(model_config): num_stages = count_dspark_stages(model_config.spec_config.speculative_model) validate_dspark_eplb_layer_base(model_config, draft_config) return DSv4DSparkForCausalLM( diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 33f6ded534b9..d95cf64349c8 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1467,7 +1467,12 @@ def external_drafter_config_kwargs(model_config, spec_config) -> dict: max_num_tokens=model_config.max_num_tokens, moe_max_num_tokens=model_config.moe_max_num_tokens, ) - if spec_config.spec_dec_mode.is_dspark(): + # Only the embedded DSpark draft shares the target's EPLB namespace (its + # stages are target decoder blocks registered into the target's balancer). + # A standalone DSpark drafter is an independent checkpoint, so it falls + # under the "other external drafters" rule above. + if (spec_config.spec_dec_mode.is_dspark() + and spec_config.draft_is_embedded_in_target): kwargs["moe_load_balancer"] = model_config.moe_load_balancer return kwargs diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 313a081b87de..ad5de9d05980 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -117,9 +117,13 @@ def should_use_separate_draft_kv_cache(spec_config) -> bool: return False if spec_config._use_shared_kv_cache: return False - # DSpark owns a dedicated rolling-window cache in DSparkWorker. Its draft - # model does not read the paged draft KV cache managed by attention metadata. - if spec_config.spec_dec_mode.is_dspark(): + # The embedded DSpark draft owns a dedicated rolling-window cache in + # DSparkWorker and never reads the paged draft KV cache that attention + # metadata manages. A standalone DSpark drafter runs on DFlashWorker, which + # does read it, so it keeps the default -- hence a flavour check, not a + # mode check (see DSparkDecodingConfig.draft_is_embedded_in_target). + if (spec_config.spec_dec_mode.is_dspark() + and spec_config.draft_is_embedded_in_target): return False return spec_config._allow_separate_draft_kv_cache diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 488d0412548b..7627f93df580 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -465,7 +465,13 @@ def _build_spec_metadata(spec_config, vocab_size=vocab_size, draft_vocab_size=draft_vocab_size, ) - if spec_config.spec_dec_mode.is_dflash(): + # A standalone DSpark drafter is drafted by DFlashWorker, so it needs the + # DFlash metadata (paged draft KV, DFlash capture buffer). Only the + # embedded DeepSeek-V4-Pro draft uses DSparkSpecMetadata and its rolling + # window. See DSparkDecodingConfig.draft_is_embedded_in_target. + if spec_config.spec_dec_mode.is_dflash() or ( + spec_config.spec_dec_mode.is_dspark() + and not spec_config.draft_is_embedded_in_target): target_layer_ids = getattr(spec_config, 'target_layer_ids', None) return DFlashSpecMetadata( max_draft_len=spec_config.max_draft_len, @@ -788,7 +794,14 @@ def get_spec_worker(spec_config, use_separate_draft_kv_cache=use_separate_draft_kv_cache) if spec_dec_mode.is_pard(): return PARDWorker(spec_config, mapping, use_separate_draft_kv_cache) - if spec_dec_mode.is_dflash(): + # Only the embedded DeepSeek-V4-Pro draft is served by DSparkWorker, whose + # rolling-window plumbing reads V4-draft-only attributes (num_stages, + # write_context_windows, forward_batched). A standalone DSpark drafter is a + # DFlash-lineage model and is served by DFlashWorker, which already probes + # the DSpark heads defensively (getattr has_markov_head / _dspark_shift_label). + if spec_dec_mode.is_dflash() or ( + spec_dec_mode.is_dspark() + and not spec_config.draft_is_embedded_in_target): return DFlashWorker(spec_config, mapping, use_separate_draft_kv_cache) if spec_dec_mode.is_dspark(): return DSparkWorker(spec_config, mapping, use_separate_draft_kv_cache) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index a57597c286a3..2112ce8c3919 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2995,6 +2995,62 @@ def tokens_per_gen_step(self) -> int: def supports_backend(self, backend: str) -> bool: return backend == "pytorch" + @functools.cached_property + def draft_is_embedded_in_target(self) -> bool: + """True for the embedded (DeepSeek-V4-Pro) flavour of the DSpark draft. + + DSpark ships in two shapes, and they need different runtime plumbing: + + - embedded: the draft is the ``mtp.*`` namespace of the *target* + checkpoint, built from full target decoder blocks, and served by + ``DSparkWorker`` with its own rolling captured-context window. + - standalone: the draft is its own checkpoint with a registry-resolved + backbone, served by ``DFlashWorker`` and its paged draft KV cache. + + Both are ``decoding_type: DSpark``, so every dispatch that must tell + them apart -- draft-model builder, worker, spec metadata, and the + separate-draft-KV-cache decision -- reads this one flag instead of + re-deriving it. That is what keeps those decisions from drifting apart: + a builder and a worker that disagree produce a draft model whose + attributes the worker does not have. + + The probe is the weight index rather than a config field because the + index is authoritative and cannot be left unset; ``model_type`` is the + fallback for a checkpoint whose index file is absent. Resolution is + memoized here and warmed during ``TorchLlmArgs`` validation, so the + filesystem probe happens once in the main process -- not per rank, and + never at forward or CUDA-graph-capture time. + """ + ckpt_dir = self.speculative_model + if ckpt_dir is None: + return False + ckpt_dir = str(ckpt_dir) + + for name in ("model.safetensors.index.json", + "pytorch_model.bin.index.json"): + index = os.path.join(ckpt_dir, name) + if not os.path.isfile(index): + continue + try: + with open(index, encoding="utf-8") as f: + weight_map = json.load(f).get("weight_map", {}) + except (OSError, ValueError): + break + # An index that parsed is authoritative both ways. Falling through + # to model_type here would classify a standalone V4-shaped drafter + # as embedded, and that only surfaces much later, inside + # count_dspark_stages. + return any(re.match(r"^mtp\.\d+\.", key) for key in weight_map) + + config_path = os.path.join(ckpt_dir, "config.json") + if os.path.isfile(config_path): + try: + with open(config_path, encoding="utf-8") as f: + return json.load(f).get("model_type") == "deepseek_v4" + except (OSError, ValueError): + return False + return False + @functools.cached_property def spec_dec_mode(self): from tensorrt_llm._torch.speculative.interface import \ @@ -6011,6 +6067,11 @@ def validate_speculative_config(self): "embedded DeepSeek-V4-Pro flavour, whose draft weights " "live in the mtp.* namespace -- the target checkpoint " "directory itself.") + # Warm the embedded-vs-standalone probe here, while we are in + # the main process and the checkpoint path is known local. The + # flag then travels with the config, so no rank repeats the + # filesystem read and nothing probes at forward time. + _ = spec_cfg.draft_is_embedded_in_target # Resolve target_layer_ids / mask_token_id / block_size / # markov_rank from the draft (or main) model config if not set. # Three checkpoint spellings are in the wild for the same knobs 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 c59c2c5909ae..b52ba3a552e1 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 @@ -68,8 +68,12 @@ def _model_config(lb_config=None, num_hidden_layers=NUM_HIDDEN_LAYERS): ) -def _spec_config(mode): - return SimpleNamespace(spec_dec_mode=mode) +def _spec_config(mode, *, embedded=True): + # Only the embedded DSpark draft shares the target's EPLB layer namespace: + # its stages are target decoder blocks registered into the target's + # balancer. A standalone DSpark drafter is an independent checkpoint and is + # treated like the other external drafters. + return SimpleNamespace(spec_dec_mode=mode, draft_is_embedded_in_target=embedded) @pytest.fixture @@ -124,6 +128,17 @@ def test_non_dspark_external_drafters_do_not_inherit_load_balancer(mode): assert "moe_load_balancer" not in kwargs +def test_standalone_dspark_drafter_does_not_inherit_load_balancer(): + # The flavour, not the mode, decides: a standalone DSpark drafter is its own + # checkpoint, so forwarding the target's EPLB config would key its experts + # against a layer namespace that is not the drafter's. + kwargs = external_drafter_config_kwargs( + _model_config(_lb_config(DSPARK_LAYERS)), + _spec_config(SpeculativeDecodingMode.DSPARK, embedded=False), + ) + assert "moe_load_balancer" not in kwargs + + def test_external_drafter_kwargs_are_stable_across_modes(): common = external_drafter_config_kwargs( _model_config(_lb_config(DSPARK_LAYERS)), _spec_config(SpeculativeDecodingMode.PARD) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_flavour_dispatch.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_flavour_dispatch.py index 72f171dedcfb..76b798534eb6 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_flavour_dispatch.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_flavour_dispatch.py @@ -24,16 +24,28 @@ without it. Silently dropping the Markov head does not fail anything; it lowers the acceptance rate, which no test would attribute back to this function. -Everything here asserts *which class is selected*, never the object it builds: -constructing a real drafter needs GPUs and checkpoints. +Selection is all this file checks -- which class each factory returns, with the +classes stubbed. That is worth pinning, but it cannot catch a worker and a +drafter that agree on paper and diverge on first contact; for that see +``test_dspark_drafter_worker_contract.py``, which builds the real drafter and +drives the real worker's lazy init. The flavour probe at the bottom is the one +exception here -- it reads checkpoints written to ``tmp_path``, because it is +the single source every dispatch above consults. """ +import json from types import SimpleNamespace import pytest from tensorrt_llm._torch.models import modeling_dflash, modeling_dspark from tensorrt_llm._torch.models.modeling_dflash import declares_dspark_heads +from tensorrt_llm._torch.speculative import utils as spec_utils +from tensorrt_llm._torch.speculative.interface import ( + SpeculativeDecodingMode, + should_use_separate_draft_kv_cache, +) +from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig _DSV4_SENTINEL = object() _QWEN3_SENTINEL = object() @@ -50,7 +62,12 @@ def _configs( - *, model_type="qwen3", dflash_config=None, architectures=None, attention_backend="TRTLLM" + *, + model_type="qwen3", + dflash_config=None, + architectures=None, + attention_backend="TRTLLM", + embedded=False, ): """Duck-typed (target ModelConfig, draft ModelConfig) for dispatch-only asserts.""" model_config = SimpleNamespace( @@ -58,6 +75,7 @@ def _configs( speculative_model="/nonexistent/drafter", block_size=7, attention_backend=attention_backend, + draft_is_embedded_in_target=embedded, ) ) draft_config = SimpleNamespace( @@ -90,19 +108,13 @@ def stub_dflash(monkeypatch): ) -def _embedded(monkeypatch, stages): - """Force the level-1 probe: ``stages`` is the mtp.* count, None if standalone.""" - monkeypatch.setattr(modeling_dspark, "count_dspark_stages", lambda _p: stages) - - # -------------------------------------------------------------------------- # decoding_type: DSpark # -------------------------------------------------------------------------- def test_standalone_qwen3_drafter_selects_qwen3_dspark(monkeypatch, stub_dspark): - _embedded(monkeypatch, None) - model_config, draft_config = _configs(model_type="qwen3") + model_config, draft_config = _configs(model_type="qwen3", embedded=False) built = modeling_dspark._build_dspark_draft(model_config, draft_config, None, None) @@ -110,20 +122,8 @@ def test_standalone_qwen3_drafter_selects_qwen3_dspark(monkeypatch, stub_dspark) def test_embedded_draft_selects_dsv4_dspark(monkeypatch, stub_dspark): - # The mtp.* namespace in the checkpoint index is the level-1 probe. - _embedded(monkeypatch, 3) - model_config, draft_config = _configs(model_type="deepseek_v4") - - built = modeling_dspark._build_dspark_draft(model_config, draft_config, None, None) - - assert built is _DSV4_SENTINEL - - -def test_deepseek_v4_without_weight_index_still_selects_dsv4(monkeypatch, stub_dspark): - # Fallback arm of the probe: a V4 checkpoint whose index file is absent must - # not fall through to the standalone lineage, which has no V4 drafter. - _embedded(monkeypatch, None) - model_config, draft_config = _configs(model_type="deepseek_v4") + monkeypatch.setattr(modeling_dspark, "count_dspark_stages", lambda _p: 3) + model_config, draft_config = _configs(model_type="deepseek_v4", embedded=True) built = modeling_dspark._build_dspark_draft(model_config, draft_config, None, None) @@ -131,8 +131,7 @@ def test_deepseek_v4_without_weight_index_still_selects_dsv4(monkeypatch, stub_d def test_unknown_standalone_model_type_lists_supported(monkeypatch, stub_dspark): - _embedded(monkeypatch, None) - model_config, draft_config = _configs(model_type="llama") + model_config, draft_config = _configs(model_type="llama", embedded=False) with pytest.raises(NotImplementedError) as excinfo: modeling_dspark._build_dspark_draft(model_config, draft_config, None, None) @@ -143,7 +142,6 @@ def test_unknown_standalone_model_type_lists_supported(monkeypatch, stub_dspark) def test_standalone_drafter_receives_the_attention_backend(monkeypatch, stub_dspark): - _embedded(monkeypatch, None) seen = {} def _capture(draft_config, *, dflash_attention_backend): @@ -247,3 +245,153 @@ def test_legacy_causal_dflash_config_is_not_mistaken_for_dspark(stub_dflash): def test_declares_dspark_heads(dflash_config, expected): config = SimpleNamespace(dflash_config=dflash_config) assert declares_dspark_heads(config) is expected + + +# -------------------------------------------------------------------------- +# Runtime routing +# +# The worker, the spec metadata and the separate-draft-KV-cache decision must +# follow the same flavour flag the builder follows. When they disagree, +# DSparkWorker gets a standalone drafter and dies reaching for V4-draft-only +# attributes (num_stages, write_context_windows) -- only at the first forward, +# long after the engine reported a successful build. +# -------------------------------------------------------------------------- + +_WORKER_SENTINELS = {"DFlashWorker": object(), "DSparkWorker": object()} +# SimpleNamespace rather than object(): get_spec_metadata assigns +# ``metadata.enable_penalty`` on whatever it built, which a bare object rejects. +_METADATA_SENTINELS = { + "DFlashSpecMetadata": SimpleNamespace(), + "DSparkSpecMetadata": SimpleNamespace(), +} + + +@pytest.fixture +def stub_runtime(monkeypatch): + """Stub the worker/metadata classes: constructing the real ones needs CUDA.""" + for name, sentinel in _WORKER_SENTINELS.items(): + monkeypatch.setattr(spec_utils, name, lambda *a, _s=sentinel, **k: _s) + for name, sentinel in _METADATA_SENTINELS.items(): + monkeypatch.setattr(spec_utils, name, lambda *a, _s=sentinel, **k: _s) + + +def _spec_config(mode, *, embedded, allow_separate_kv=True): + """Duck-typed spec config carrying only what the routing functions read.""" + return SimpleNamespace( + spec_dec_mode=mode, + draft_is_embedded_in_target=embedded, + _use_shared_kv_cache=False, + _allow_separate_draft_kv_cache=allow_separate_kv, + max_draft_len=7, + max_total_draft_tokens=7, + tokens_per_gen_step=8, + target_layer_ids=[7, 23, 51, 67, 83], + advanced_sampling_mode=None, + # Read by get_spec_metadata for the occurrence-penalty workspace; the + # routing under test does not depend on it. + enable_penalty=False, + ) + + +@pytest.mark.parametrize( + "mode,embedded,expected", + [ + (SpeculativeDecodingMode.DSPARK, True, "DSparkWorker"), + (SpeculativeDecodingMode.DSPARK, False, "DFlashWorker"), + (SpeculativeDecodingMode.DFLASH, False, "DFlashWorker"), + ], +) +def test_worker_follows_the_flavour_not_the_mode(stub_runtime, mode, embedded, expected): + worker = spec_utils.get_spec_worker( + _spec_config(mode, embedded=embedded), + model_config=None, + mapping=None, + use_separate_draft_kv_cache=False, + ) + assert worker is _WORKER_SENTINELS[expected] + + +@pytest.mark.parametrize( + "mode,embedded,expected", + [ + (SpeculativeDecodingMode.DSPARK, True, "DSparkSpecMetadata"), + (SpeculativeDecodingMode.DSPARK, False, "DFlashSpecMetadata"), + (SpeculativeDecodingMode.DFLASH, False, "DFlashSpecMetadata"), + ], +) +def test_spec_metadata_follows_the_flavour_not_the_mode(stub_runtime, mode, embedded, expected): + metadata = spec_utils.get_spec_metadata( + _spec_config(mode, embedded=embedded), + SimpleNamespace(hidden_size=7168, torch_dtype=None, vocab_size=163840), + max_num_requests=8, + max_num_tokens=4096, + ) + assert metadata is _METADATA_SENTINELS[expected] + + +@pytest.mark.parametrize( + "mode,embedded,expected", + [ + # The embedded draft opts out: it owns a rolling captured-context window. + (SpeculativeDecodingMode.DSPARK, True, False), + # A standalone DSpark drafter runs on DFlashWorker's paged draft KV -- + # the path K3 used before its decoding_type moved to DSpark. + (SpeculativeDecodingMode.DSPARK, False, True), + (SpeculativeDecodingMode.DFLASH, False, True), + ], +) +def test_separate_draft_kv_cache_follows_the_flavour(mode, embedded, expected): + config = _spec_config(mode, embedded=embedded) + assert should_use_separate_draft_kv_cache(config) is expected + + +# -------------------------------------------------------------------------- +# The flavour probe itself. This is the single source every dispatch above +# reads, so it is the one place the embedded/standalone question is decided. +# -------------------------------------------------------------------------- + + +def _write_ckpt(tmp_path, *, weight_map=None, model_type=None): + if weight_map is not None: + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": weight_map}) + ) + if model_type is not None: + (tmp_path / "config.json").write_text(json.dumps({"model_type": model_type})) + return DSparkDecodingConfig(max_draft_len=7, speculative_model=str(tmp_path)) + + +def test_probe_reads_the_mtp_namespace_from_the_weight_index(tmp_path): + config = _write_ckpt( + tmp_path, + weight_map={ + "mtp.0.attn.wq_a.weight": "x.safetensors", + "layers.0.q.weight": "x.safetensors", + }, + model_type="deepseek_v4", + ) + assert config.draft_is_embedded_in_target is True + + +def test_probe_treats_a_standalone_drafter_index_as_standalone(tmp_path): + config = _write_ckpt( + tmp_path, + weight_map={"layers.0.self_attn.q_proj.weight": "x.safetensors"}, + model_type="qwen3", + ) + assert config.draft_is_embedded_in_target is False + + +def test_probe_falls_back_to_model_type_without_an_index(tmp_path): + # A V4 checkpoint whose index file is absent must not be read as + # standalone: the standalone lineage has no V4 drafter. + config = _write_ckpt(tmp_path, model_type="deepseek_v4") + assert config.draft_is_embedded_in_target is True + + +def test_probe_is_standalone_when_nothing_can_be_read(tmp_path): + # Fail soft: an unreadable or not-yet-downloaded checkpoint must not crash + # config validation, and standalone is the safe default (it is the flavour + # whose worker probes the DSpark heads defensively). + config = DSparkDecodingConfig(max_draft_len=7, speculative_model=str(tmp_path / "missing")) + assert config.draft_is_embedded_in_target is False From 0c6331ba5efe0d0333cd504d907236e782f9dc5f Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Wed, 19 Aug 2026 08:08:17 -0700 Subject: [PATCH 11/21] [None][test] Exercise the DSpark drafter and the K3 mode gate for real The dispatch tests assert which class a factory returns, with the classes stubbed. Two real failures walked past them: the K3 target's spec-dec mode whitelist, and a worker handed a draft model whose attributes it does not have. Both needed a 16-GPU run to surface, one of them five minutes in. Add the two tests that catch them on one GPU in under two seconds. The contract test builds the real Qwen3DSparkForCausalLM and drives the real DFlashWorker's lazy init, which is where the draft-model interface is actually consumed; it also pins the mis-route, so reverting the routing to a mode check fails here instead of in production. The gate test asks only whether construction stopped at the whitelist, since everything past it builds the full K3 model. The drafter is built on the TRTLLM block-decode backend, matching the K3 serving config; the VANILLA default would pull in flash-attn. Signed-off-by: Zhenhuan Chen --- .../test_dspark_drafter_worker_contract.py | 223 ++++++++++++++++++ .../test_kimi_k3_spec_mode_gate.py | 114 +++++++++ 2 files changed, 337 insertions(+) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_drafter_worker_contract.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_spec_mode_gate.py diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_drafter_worker_contract.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_drafter_worker_contract.py new file mode 100644 index 000000000000..232a713c6761 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_drafter_worker_contract.py @@ -0,0 +1,223 @@ +# 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. +"""Contract tests: a real standalone DSpark drafter driven by the real worker +that the routing selects for it. + +``test_dspark_flavour_dispatch.py`` asserts *which class* each factory returns, +using stubs. That cannot catch a worker handed a draft model whose attributes it +does not have: the two agree on paper and diverge on first contact. The routing +bug this file pins surfaced only as + + AttributeError: 'Qwen3DSparkForCausalLM' object has no attribute 'num_stages' + +five minutes into a 16-GPU run, after the weights had loaded -- because nothing +below the factory was ever exercised. + +So these tests build the drafter for real and drive ``DFlashWorker``'s lazy +init, which is where the drafter contract actually lives: it reaches for +``fc.weight``, ``block_size``, ``_build_fused_kv_buffers``, ``_num_attn_layers``, +``_num_heads``, ``_num_kv_heads``, ``_head_dim`` and ``_get_attention_mask_args``. +A worker routed by mode instead of by flavour fails here, in seconds, on one GPU. + +The drafter is built with the TRTLLM block-decode backend, matching the K3 +serving config; the VANILLA default would pull in flash-attn. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.speculative import utils as spec_utils +from tensorrt_llm._torch.speculative.dflash import DFlashWorker +from tensorrt_llm._torch.speculative.dspark import DSparkWorker +from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode +from tensorrt_llm.mapping import Mapping + +needs_gpu = pytest.mark.skipif( + not torch.cuda.is_available(), reason="the drafter and the worker buffers are CUDA-resident" +) + +VOCAB = 256 +RANK = 8 +BLOCK_SIZE = 4 +MAX_REQUESTS = 4 +MAX_SEQ_LEN = 128 + +TINY = dict( + architectures=["DSparkDraftModel"], + model_type="qwen3", + block_size=BLOCK_SIZE, + hidden_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + # The fused QK-norm-RoPE kernel the bf16 block decode uses rejects small + # head dims; 128 is also the real K3 drafter's head_dim. + head_dim=128, + intermediate_size=128, + hidden_act="silu", + rms_norm_eps=1e-6, + vocab_size=VOCAB, + max_position_embeddings=2048, + rope_theta=10000.0, + rope_scaling=None, + attention_bias=False, + torch_dtype="bfloat16", + num_target_layers=2, + tie_word_embeddings=False, +) + +NUM_CAPTURE = 2 + + +def _drafter_config(): + """A standalone DSpark drafter config: qwen3 backbone plus the head set.""" + from transformers import Qwen3Config + + cfg = dict(TINY) + cfg["dflash_config"] = { + "mask_token_id": VOCAB - 2, + "target_layer_ids": [0, 1], + "projector_type": "dspark", + "causal": False, + "shift_label": True, + "markov_rank": RANK, + "markov_head_type": "vanilla", + "use_confidence_head": True, + } + return Qwen3Config.from_dict(cfg) + + +def _drafter_weights(seed=11): + g = torch.Generator().manual_seed(seed) + + def rnd(*shape): + return (torch.randn(*shape, generator=g) * 0.05).to(torch.bfloat16) + + h, inter = TINY["hidden_size"], TINY["intermediate_size"] + nh, nkv, hd = (TINY["num_attention_heads"], TINY["num_key_value_heads"], TINY["head_dim"]) + weights = { + "fc.weight": rnd(h, h * NUM_CAPTURE), + "hidden_norm.weight": rnd(h) + 1.0, + "norm.weight": rnd(h) + 1.0, + "markov_w1.weight": rnd(VOCAB, RANK), + "markov_w2.weight": rnd(VOCAB, RANK), + "confidence_proj.weight": rnd(1, h + RANK), + "confidence_proj.bias": rnd(1), + } + for i in range(TINY["num_hidden_layers"]): + p = f"layers.{i}." + weights[p + "self_attn.q_proj.weight"] = rnd(nh * hd, h) + weights[p + "self_attn.k_proj.weight"] = rnd(nkv * hd, h) + weights[p + "self_attn.v_proj.weight"] = rnd(nkv * hd, h) + weights[p + "self_attn.o_proj.weight"] = rnd(h, nh * hd) + weights[p + "self_attn.q_norm.weight"] = rnd(hd) + 1.0 + weights[p + "self_attn.k_norm.weight"] = rnd(hd) + 1.0 + weights[p + "input_layernorm.weight"] = rnd(h) + 1.0 + weights[p + "post_attention_layernorm.weight"] = rnd(h) + 1.0 + weights[p + "mlp.gate_proj.weight"] = rnd(inter, h) + weights[p + "mlp.up_proj.weight"] = rnd(inter, h) + weights[p + "mlp.down_proj.weight"] = rnd(h, inter) + return weights + + +@pytest.fixture(scope="module") +def standalone_drafter(): + """The real ``Qwen3DSparkForCausalLM``, weights loaded, on the device.""" + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models.modeling_dspark import Qwen3DSparkForCausalLM + + model_config = ModelConfig(pretrained_config=_drafter_config(), attn_backend="TRTLLM") + drafter = Qwen3DSparkForCausalLM(model_config, dflash_attention_backend="TRTLLM").to("cuda") + drafter.load_weights(_drafter_weights()) + return drafter + + +def _spec_config(*, embedded): + """Only the fields the routing and the worker read.""" + return SimpleNamespace( + spec_dec_mode=SpeculativeDecodingMode.DSPARK, + draft_is_embedded_in_target=embedded, + _use_shared_kv_cache=False, + _allow_separate_draft_kv_cache=True, + # K == block_size under the dspark shift_label convention, which is + # also what DSparkWorker validates before it touches anything else. + max_draft_len=BLOCK_SIZE, + attention_backend="TRTLLM", + ) + + +def _lazy_init_args(): + spec_metadata = SimpleNamespace(max_num_requests=MAX_REQUESTS) + attn_metadata = SimpleNamespace(max_seq_len=MAX_SEQ_LEN) + return spec_metadata, attn_metadata + + +@needs_gpu +def test_routed_worker_initializes_against_a_real_standalone_drafter(standalone_drafter): + """The end-to-end contract, and the regression test for the routing bug. + + ``get_spec_worker`` picks the worker; the drafter is the real one the + builder would produce for the same config. Driving lazy init proves the two + agree on the draft-model interface. Route by mode instead of by flavour and + this raises ``AttributeError: ... has no attribute 'num_stages'``. + """ + spec_config = _spec_config(embedded=False) + worker = spec_utils.get_spec_worker( + spec_config, model_config=None, mapping=Mapping(), use_separate_draft_kv_cache=True + ) + assert isinstance(worker, DFlashWorker) + + worker.set_draft_model(standalone_drafter) + worker._lazy_init_ctx_buffers(standalone_drafter, *_lazy_init_args()) + + assert worker._ctx_buf_inited + # One scratch slot on top of the request slots, so dummy/padded writes + # cannot land on a real request's context. + assert worker._ctx_len.shape == (MAX_REQUESTS + 1,) + assert worker._dummy_slot == MAX_REQUESTS + assert worker._batch_to_slot.shape == (MAX_REQUESTS,) + assert worker._resolved_block_size == BLOCK_SIZE + assert sorted(worker._free_slots) == list(range(MAX_REQUESTS)) + + +@needs_gpu +def test_routed_worker_sees_the_dspark_heads(standalone_drafter): + """The heads moved to the DSpark subclass must stay visible to the worker. + + ``DFlashWorker`` probes them defensively (``getattr(..., False)``), so a + drafter that lost them degrades to plain DFlash silently -- lower acceptance, + no error. These are the two probes the block-draft step makes. + """ + assert getattr(standalone_drafter, "has_markov_head", False) is True + assert getattr(standalone_drafter, "_dspark_shift_label", False) is True + assert standalone_drafter.markov_w1.shape == (VOCAB, RANK) + assert standalone_drafter.markov_w2.shape == (VOCAB, RANK) + + +@needs_gpu +def test_embedded_worker_cannot_drive_a_standalone_drafter(standalone_drafter): + """Witness for why the routing has to follow the flavour. + + ``DSparkWorker`` serves the embedded DeepSeek-V4-Pro draft and reads + V4-draft-only attributes. Handing it a standalone drafter is the exact + mis-route that reached production, so pin the failure rather than trusting + the routing test alone to stay correct. + """ + worker = DSparkWorker(_spec_config(embedded=True), Mapping()) + spec_metadata, _ = _lazy_init_args() + with pytest.raises(AttributeError, match="num_stages"): + worker._lazy_init(standalone_drafter, spec_metadata) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_spec_mode_gate.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_spec_mode_gate.py new file mode 100644 index 000000000000..2fe70dbadac3 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_spec_mode_gate.py @@ -0,0 +1,114 @@ +# 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. +"""The Kimi K3 target's spec-dec mode gate. + +``KimiLinearForCausalLM.__init__`` whitelists the speculative-decoding modes the +target will serve. It is a plain assert on the *target* side, so nothing that +exercises the drafter, the builder or the worker can reach it -- which is how a +K3 engine configured with ``decoding_type: DSpark`` got through every unit test +and then failed at model construction with + + AssertionError: Kimi K3 supports speculative decoding only with SA or DFlash + +The gate is the second statement in ``__init__``, and everything past it builds +the real K3 model, which needs weights and 16 GPUs. So rather than stub the +framework out from under it -- the base initializer's arguments construct +``KimiLinearModel`` before the base is even called, so stubbing the base does +not help -- these tests ask a narrower question: did construction fail *at the +gate*, or did it get past it? Anything that fails later has passed the gate, +which is the whole of what is under test here. +""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM +from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode + +# Admitted: SA drafts in-forward with no draft weights; DFlash and DSpark are +# the external-drafter flow, and the target side is identical for both (the +# hidden-state capture in KimiLinearModel.forward is unconditional). +ADMITTED = [ + SpeculativeDecodingMode.SA, + SpeculativeDecodingMode.DFLASH, + SpeculativeDecodingMode.DSPARK, +] +# Refused: these need draft heads that no K3 checkpoint ships. +REFUSED = [ + SpeculativeDecodingMode.MTP, + SpeculativeDecodingMode.EAGLE3_ONE_MODEL, +] + +_SPEC_GATE = "speculative decoding" +_PP_GATE = "pipeline parallelism" + + +def _model_config(mode, *, pp_size=1): + """The minimum ``__init__`` reads before it reaches the gate.""" + return SimpleNamespace( + pretrained_config=SimpleNamespace(model_type="kimi_linear", linear_attn_config={}), + mapping=SimpleNamespace(pp_size=pp_size), + spec_config=None if mode is None else SimpleNamespace(spec_dec_mode=mode), + ) + + +def _rejected_by(model_config) -> str | None: + """Which gate rejected this config, or None if construction got past them. + + Only the two guard asserts count as a rejection; construction is expected + to fail afterwards on the real model, and that failure means the config was + admitted. An AssertionError from anywhere else is a genuine problem and is + re-raised rather than silently read as a rejection. + """ + try: + KimiLinearForCausalLM(model_config) + except AssertionError as exc: + message = str(exc) + for gate in (_SPEC_GATE, _PP_GATE): + if gate in message: + return gate + raise + except Exception: + return None + return None + + +@pytest.mark.parametrize("mode", ADMITTED, ids=lambda m: m.name) +def test_admitted_modes_pass_the_gate(mode): + assert _rejected_by(_model_config(mode)) is None + + +def test_no_spec_config_passes_the_gate(): + assert _rejected_by(_model_config(None)) is None + + +@pytest.mark.parametrize("mode", REFUSED, ids=lambda m: m.name) +def test_refused_modes_are_rejected_at_the_spec_gate(mode): + assert _rejected_by(_model_config(mode)) == _SPEC_GATE + + +def test_the_refusal_message_names_the_admitted_modes(): + # The message is the only guidance a user gets, so it has to name the modes + # that would work -- that is what turns "not supported" into an action. + with pytest.raises(AssertionError, match="SA, DFlash or DSpark"): + KimiLinearForCausalLM(_model_config(SpeculativeDecodingMode.MTP)) + + +def test_pipeline_parallelism_is_still_rejected(): + # The pp guard sits ahead of the spec gate; pin the order so a future edit + # to the mode list cannot let a pp>1 config through. + config = _model_config(SpeculativeDecodingMode.DSPARK, pp_size=2) + assert _rejected_by(config) == _PP_GATE From 67aaecfbdefa0ab56cf032831380e99c8e36060f Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Wed, 19 Aug 2026 19:13:07 -0700 Subject: [PATCH 12/21] [None][refactor] Split the DSpark workers by deployment form Worker classification now matches the draft-model side. DSparkWorker becomes DSv4DSparkWorker, and standalone DSpark drafters get StandaloneDSparkWorker, a DFlashWorker subclass carrying the only two policies that differ: the shift_label block-output slot convention and the Markov intra-block logit bias. DFlashWorker had been probing both defensively through getattr, so a plain DFlash drafter paid for DSpark bookkeeping it never used and a DSpark drafter that lost a head degraded silently. The probes now live in the subclass and reach the drafter through explicit hooks (_draft_slot_ids, _refine_block_logits). Workers are named by deployment form, never by draft backbone: the worker only allocates against shapes the draft model reports and sequences calls it owns, so an MLA drafter reuses StandaloneDSparkWorker unchanged. The rationale is recorded on the classes themselves, and the names share vocabulary with DSparkDecodingConfig.draft_is_embedded_in_target, which is what get_spec_worker branches on. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/models/modeling_dspark.py | 2 +- tensorrt_llm/_torch/speculative/dflash.py | 87 +++----- tensorrt_llm/_torch/speculative/dspark.py | 156 +++++++++++++- tensorrt_llm/_torch/speculative/interface.py | 9 +- tensorrt_llm/_torch/speculative/utils.py | 20 +- tensorrt_llm/llmapi/llm_args.py | 2 +- .../test_dspark_drafter_worker_contract.py | 204 ++++++++++++++++-- .../test_dspark_flavour_dispatch.py | 33 ++- .../hw_agnostic/test_dspark_worker.py | 6 +- 9 files changed, 423 insertions(+), 96 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 2d1a2bac6d00..1db272fc0e59 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -1823,7 +1823,7 @@ class DSv4DSparkForCausalLM(nn.Module): Wraps :class:`DSv4DSparkDraftModel` (the ``n_mtp_layers``-stage ``mtp.*`` backbone) for the single-engine external-drafter flow: created by ``get_draft_model``, - appended to the target's epilogue, and driven by ``DSparkWorker``. + appended to the target's epilogue, and driven by ``DSv4DSparkWorker``. ``embed_tokens`` / ``lm_head`` are shared with the target model (:meth:`load_weights_from_target_model`). The draft weights live in the SAME diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 14644ff3e774..b16ff043a6a1 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -698,13 +698,10 @@ def _forward_impl( # Gather K logits per gen request from the block outputs. # hidden_states_out is flat: [num_gens * block_size, hidden_dim]. - # Plain DFlash reads mask slots 1..K; dspark shift_label reads - # slots 0..K-1 (see dflash_draft_slot_ids). + # Which block slots carry them is a drafter-family convention, + # resolved through _draft_slot_ids. block_size = self._compute_block_size - shift_label = getattr(draft_model, "_dspark_shift_label", False) - gen_gather_ids = dflash_draft_slot_ids( - num_gens, block_size, K, shift_label, device="cuda" - ) + gen_gather_ids = self._draft_slot_ids(draft_model, num_gens, block_size, K) gen_gather_ids = gen_gather_ids.clamp(max=hidden_states_out.shape[0] - 1) gen_logits = draft_model.logits_processor( @@ -714,12 +711,9 @@ def _forward_impl( vocab_size = gen_logits.shape[-1] gen_logits = gen_logits.reshape(num_gens, K, vocab_size) - # DSpark Markov head: add the greedy-chained intra-block - # logit bias before sampling (no-op for plain DFlash). - if getattr(draft_model, "has_markov_head", False): - gen_logits = self._apply_dspark_markov_bias( - draft_model, gen_logits, inputs["first_prev_tokens"], spec_metadata - ) + gen_logits = self._refine_block_logits( + draft_model, gen_logits, inputs, spec_metadata + ) gen_draft_tokens = self.sample_draft_tokens( gen_logits, @@ -791,60 +785,33 @@ def _forward_impl( "next_new_tokens": next_new_tokens, } - def _apply_dspark_markov_bias( + def _draft_slot_ids( + self, draft_model, num_gens: int, block_size: int, num_draft_tokens: int + ) -> torch.Tensor: + """Block-output slots whose hidden states produce the K draft logits. + + Plain DFlash uses the K2.7 convention: mask slots 1..K. Drafter + families with another convention override this — see + :meth:`DSparkWorker._draft_slot_ids` for the shift_label + variant, which reads slots 0..K-1 instead. + """ + return dflash_draft_slot_ids(num_gens, block_size, num_draft_tokens, False, device="cuda") + + def _refine_block_logits( self, draft_model, gen_logits: torch.Tensor, - first_prev_tokens: torch.Tensor, + inputs: dict, spec_metadata, ) -> torch.Tensor: - """Apply the dspark vanilla-Markov intra-block bias to block logits. - - Reference (DeepSpec VanillaMarkov.sample_block_tokens, temperature 0): - step i adds bias = markov_w2 @ markov_w1[prev_i] to the shared-lm_head - logits, where prev_0 is the anchor (last accepted) token and prev_{i>0} - is the greedy token from step i-1's biased logits. Greedy per-position - argmax of the returned logits therefore reproduces the reference - sampled chain; the rejection-sampling path samples from the same - biased distributions (proposal conditioned on the greedy chain). - - Handles a TP vocab-sharded draft lm_head by slicing markov_w2's rows - to this rank's contiguous shard and chaining through the TP-aware - global argmax. + """Refine the block logits between the draft forward and sampling. + + Plain DFlash proposes the backbone's logits unchanged. Drafter + families carrying extra heads override this — see + :meth:`DSparkWorker._refine_block_logits` for the Markov + intra-block bias. """ - if self._d2t is not None: - raise NotImplementedError( - "DSpark Markov head requires a shared draft/target vocab " - "(d2t vocab mapping is not supported)." - ) - full_vocab = draft_model.markov_w2.shape[0] - shard = gen_logits.shape[-1] - vocab_slice = None - if shard != full_vocab: - mapping = self.mapping - if ( - mapping is None - or getattr(mapping, "enable_attention_dp", False) - or shard * mapping.tp_size != full_vocab - ): - raise NotImplementedError( - f"DSpark Markov head: draft logits width {shard} does not " - f"match the drafter vocab {full_vocab} and is not a plain " - "TP column shard of it." - ) - vocab_slice = slice(mapping.tp_rank * shard, (mapping.tp_rank + 1) * shard) - - def argmax_fn(step_logits): - # Full-vocab token ids (TP-aware when sharded); tokens stay in - # draft-vocab space, which is what markov_w1 indexes. - return self.greedy_sample_draft_with_tp_gather(step_logits, spec_metadata).long() - - return draft_model.apply_markov_chain_logits( - gen_logits, - first_prev_tokens, - argmax_fn=argmax_fn, - vocab_slice=vocab_slice, - ) + return gen_logits def prepare_1st_drafter_inputs( self, diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 6840200e5cbb..bbdaddd1043d 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -30,6 +30,7 @@ from tensorrt_llm.mapping import Mapping from ..pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID +from .dflash import DFlashWorker, dflash_draft_slot_ids from .interface import SpecMetadata, SpecWorkerBase if TYPE_CHECKING: @@ -184,7 +185,7 @@ def get_hidden_states(self, num_tokens: int) -> Optional[torch.Tensor]: ] -class DSparkWorker(SpecWorkerBase): +class DSv4DSparkWorker(SpecWorkerBase): """Worker for DSpark speculative decoding. DSpark drafts a whole block of ``block_size`` tokens in one backbone forward @@ -209,6 +210,14 @@ class DSparkWorker(SpecWorkerBase): write done by the generation path. These affect draft acceptance rate only, not correctness, which the standard target verify guarantees. + Naming: workers are classified by *deployment form*, not by draft + backbone (see :class:`DSparkWorker`). This one is form-specific + because it owns a rolling captured-context window and drives the draft + through attributes only an embedded DeepSeek-V4-Pro draft has -- + ``num_stages``, ``_attn_params``, ``write_context_windows``, + ``write_context_windows_batched`` and ``forward_batched``. A standalone + drafter has none of them and is served by :class:`DSparkWorker`. + Reference: DeepSeek DeepSpec (https://github.com/deepseek-ai/DeepSpec). """ @@ -252,7 +261,7 @@ def __init__( # capture-safe whenever ``cuda_graph_config`` is set). logger.info( - f"DSparkWorker initialized with " + f"DSv4DSparkWorker initialized with " f"use_separate_draft_kv_cache={use_separate_draft_kv_cache}" ) @@ -696,3 +705,146 @@ def _forward_impl( "next_draft_tokens": next_draft_tokens, "next_new_tokens": next_new_tokens, } + + +class DSparkWorker(DFlashWorker): + """Worker for a *standalone* DSpark drafter (DFlash lineage). + + DSpark is DFlash plus two extra heads, so the drafting plumbing is + inherited wholesale from :class:`DFlashWorker` -- paged context K/V, + slot management, the mask-token block forward -- and only the two + head-driven policies are overridden here: the block-output slot + convention (``shift_label``) and the Markov intra-block logit bias. + + Mirrors the model side, where ``DSparkDrafterForCausalLM`` extends + ``DFlashForCausalLM`` with the same two heads. + + Naming: this is the unqualified DSpark worker because a separately + shipped drafter is the ordinary case; :class:`DSv4DSparkWorker` carries + the qualifier because a draft embedded in the target checkpoint is the + special one. Workers are classified by *deployment form*, never by draft + backbone -- so there is no ``Qwen3DSparkWorker``. Note the name meant the + embedded worker before this split; both the rebind and the rename to + ``DSv4DSparkWorker`` land in one commit so the swap reads as a unit. + + A worker is agnostic to the draft backbone: everything backbone-shaped is + supplied by the draft model, which reports its own shapes + (``_num_attn_layers``, ``_num_heads``, ``_num_kv_heads``, ``_head_dim``) + and owns the operators (``_build_fused_kv_buffers``, + ``precompute_context_kv``, ``dflash_forward``, + ``apply_markov_chain_logits``, ``project_target_hidden``). The worker only + allocates against the reported shapes and sequences the calls. An MLA + drafter therefore reuses this class unchanged; its differences (fused-QKV + assumptions, a 576-latent K/V layout) land in its own draft-model + subclass. Naming workers by backbone would produce N classes with + identical bodies. + + Deployment form is the axis the runtime state actually splits on: paged + draft K/V here, a worker-owned rolling window in + :class:`DSv4DSparkWorker`. + """ + + def set_draft_model(self, draft_model) -> None: + """Reject an unsupported vocab mapping here rather than mid-decode. + + ``d2t`` is model-static, so a config mistake should surface at load and + not as a ``NotImplementedError`` raised per decode step, possibly during + CUDA-graph capture. + """ + super().set_draft_model(draft_model) + if self._d2t is not None and getattr(draft_model, "has_markov_head", False): + raise NotImplementedError( + "DSpark Markov head requires a shared draft/target vocab " + "(d2t vocab mapping is not supported); drafter " + f"{type(draft_model).__name__} declares one." + ) + + def _draft_slot_ids( + self, draft_model, num_gens: int, block_size: int, num_draft_tokens: int + ) -> torch.Tensor: + """Block-output slots under the dspark ``shift_label`` convention. + + The drafter checkpoint declares the convention, so it is read off the + draft model rather than assumed: a DSpark drafter trained with the + legacy DFlash slot layout keeps the base class' slots 1..K. + """ + shift_label = getattr(draft_model, "_dspark_shift_label", False) + return dflash_draft_slot_ids( + num_gens, block_size, num_draft_tokens, shift_label, device="cuda" + ) + + def _refine_block_logits( + self, + draft_model, + gen_logits: torch.Tensor, + inputs: dict, + spec_metadata, + ) -> torch.Tensor: + """Add the greedy-chained Markov intra-block bias to the block logits. + + A DSpark drafter checkpoint may omit the Markov head (``markov_rank`` + 0), which loads as a drafter without one; that case falls through to + the unmodified backbone logits. + """ + if not getattr(draft_model, "has_markov_head", False): + return gen_logits + return self._apply_dspark_markov_bias( + draft_model, gen_logits, inputs["first_prev_tokens"], spec_metadata + ) + + def _apply_dspark_markov_bias( + self, + draft_model, + gen_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + spec_metadata, + ) -> torch.Tensor: + """Apply the dspark vanilla-Markov intra-block bias to block logits. + + Reference (DeepSpec VanillaMarkov.sample_block_tokens, temperature 0): + step i adds bias = markov_w2 @ markov_w1[prev_i] to the shared-lm_head + logits, where prev_0 is the anchor (last accepted) token and prev_{i>0} + is the greedy token from step i-1's biased logits. Greedy per-position + argmax of the returned logits therefore reproduces the reference + sampled chain; the rejection-sampling path samples from the same + biased distributions (proposal conditioned on the greedy chain). + + Handles a TP vocab-sharded draft lm_head by slicing markov_w2's rows + to this rank's contiguous shard and chaining through the TP-aware + global argmax. + """ + # The d2t guard lives in set_draft_model: it is model-static, so raising + # it here would surface a load-time config error per decode step. + # Unlike the d2t guard this one cannot move to set_draft_model: it + # keys on the runtime logits width, and reproducing that at init would + # duplicate the draft head's sharding rules. A standalone drafter + # borrows the target lm_head, whose gather_output defaults to True, so + # the logits normally arrive full-vocab and this branch is skipped. + full_vocab = draft_model.markov_w2.shape[0] + shard = gen_logits.shape[-1] + vocab_slice = None + if shard != full_vocab: + mapping = self.mapping + if ( + mapping is None + or getattr(mapping, "enable_attention_dp", False) + or shard * mapping.tp_size != full_vocab + ): + raise NotImplementedError( + f"DSpark Markov head: draft logits width {shard} does not " + f"match the drafter vocab {full_vocab} and is not a plain " + "TP column shard of it." + ) + vocab_slice = slice(mapping.tp_rank * shard, (mapping.tp_rank + 1) * shard) + + def argmax_fn(step_logits): + # Full-vocab token ids (TP-aware when sharded); tokens stay in + # draft-vocab space, which is what markov_w1 indexes. + return self.greedy_sample_draft_with_tp_gather(step_logits, spec_metadata).long() + + return draft_model.apply_markov_chain_logits( + gen_logits, + first_prev_tokens, + argmax_fn=argmax_fn, + vocab_slice=vocab_slice, + ) diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index ad5de9d05980..4084fbc3ffc3 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -118,10 +118,11 @@ def should_use_separate_draft_kv_cache(spec_config) -> bool: if spec_config._use_shared_kv_cache: return False # The embedded DSpark draft owns a dedicated rolling-window cache in - # DSparkWorker and never reads the paged draft KV cache that attention - # metadata manages. A standalone DSpark drafter runs on DFlashWorker, which - # does read it, so it keeps the default -- hence a flavour check, not a - # mode check (see DSparkDecodingConfig.draft_is_embedded_in_target). + # DSv4DSparkWorker and never reads the paged draft KV cache that attention + # metadata manages. A standalone DSpark drafter runs on DSparkWorker + # (DFlash lineage), which does read it, so it keeps the default -- hence a + # form check, not a mode check + # (see DSparkDecodingConfig.draft_is_embedded_in_target). if (spec_config.spec_dec_mode.is_dspark() and spec_config.draft_is_embedded_in_target): return False diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 7627f93df580..8efbac572925 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -21,7 +21,7 @@ from .dflash import DFlashSpecMetadata, DFlashWorker from .draft_target import (DraftTargetOneModelSpecMetadata, DraftTargetOneModelWorker) -from .dspark import DSparkSpecMetadata, DSparkWorker +from .dspark import DSparkSpecMetadata, DSparkWorker, DSv4DSparkWorker from .eagle3 import (Eagle3OneModelDynamicTreeResourceManager, Eagle3OneModelSpecMetadata, Eagle3OneModelWorker, Eagle3ResourceManager, Eagle3SpecMetadata, MTPEagleWorker) @@ -794,16 +794,18 @@ def get_spec_worker(spec_config, use_separate_draft_kv_cache=use_separate_draft_kv_cache) if spec_dec_mode.is_pard(): return PARDWorker(spec_config, mapping, use_separate_draft_kv_cache) - # Only the embedded DeepSeek-V4-Pro draft is served by DSparkWorker, whose - # rolling-window plumbing reads V4-draft-only attributes (num_stages, - # write_context_windows, forward_batched). A standalone DSpark drafter is a - # DFlash-lineage model and is served by DFlashWorker, which already probes - # the DSpark heads defensively (getattr has_markov_head / _dspark_shift_label). - if spec_dec_mode.is_dflash() or ( - spec_dec_mode.is_dspark() - and not spec_config.draft_is_embedded_in_target): + if spec_dec_mode.is_dflash(): return DFlashWorker(spec_config, mapping, use_separate_draft_kv_cache) + # DSpark splits by deployment form, mirroring the draft-model side. The + # embedded DeepSeek-V4-Pro draft needs DSv4DSparkWorker, whose rolling-window + # plumbing reads V4-draft-only attributes (num_stages, write_context_windows, + # forward_batched). A standalone drafter is DFlash lineage and is served by + # DSparkWorker, which adds only the Markov bias and the shift_label + # slot convention on top of DFlashWorker. if spec_dec_mode.is_dspark(): + if spec_config.draft_is_embedded_in_target: + return DSv4DSparkWorker(spec_config, mapping, + use_separate_draft_kv_cache) return DSparkWorker(spec_config, mapping, use_separate_draft_kv_cache) if spec_dec_mode.is_sa(): return SAWorker(spec_config, model_config) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 2112ce8c3919..4350be676078 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3003,7 +3003,7 @@ def draft_is_embedded_in_target(self) -> bool: - embedded: the draft is the ``mtp.*`` namespace of the *target* checkpoint, built from full target decoder blocks, and served by - ``DSparkWorker`` with its own rolling captured-context window. + ``DSv4DSparkWorker`` with its own rolling captured-context window. - standalone: the draft is its own checkpoint with a registry-resolved backbone, served by ``DFlashWorker`` and its paged draft KV cache. diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_drafter_worker_contract.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_drafter_worker_contract.py index 232a713c6761..07015c803878 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_drafter_worker_contract.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_drafter_worker_contract.py @@ -42,7 +42,7 @@ from tensorrt_llm._torch.speculative import utils as spec_utils from tensorrt_llm._torch.speculative.dflash import DFlashWorker -from tensorrt_llm._torch.speculative.dspark import DSparkWorker +from tensorrt_llm._torch.speculative.dspark import DSparkWorker, DSv4DSparkWorker from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode from tensorrt_llm.mapping import Mapping @@ -101,7 +101,14 @@ def _drafter_config(): return Qwen3Config.from_dict(cfg) -def _drafter_weights(seed=11): +def _drafter_weights(seed=11, legacy_head_keys=False): + """Synthetic drafter weights in the published key namespace. + + Head tensors are named the way both public checkpoints ship them + (``markov_head.*`` / ``confidence_head.proj.*``, after the submodules that + own them); ``legacy_head_keys`` switches to the bare spellings so the + aliasing is covered from both sides. + """ g = torch.Generator().manual_seed(seed) def rnd(*shape): @@ -109,15 +116,25 @@ def rnd(*shape): h, inter = TINY["hidden_size"], TINY["intermediate_size"] nh, nkv, hd = (TINY["num_attention_heads"], TINY["num_key_value_heads"], TINY["head_dim"]) - weights = { - "fc.weight": rnd(h, h * NUM_CAPTURE), - "hidden_norm.weight": rnd(h) + 1.0, - "norm.weight": rnd(h) + 1.0, + head = { "markov_w1.weight": rnd(VOCAB, RANK), "markov_w2.weight": rnd(VOCAB, RANK), "confidence_proj.weight": rnd(1, h + RANK), "confidence_proj.bias": rnd(1), } + if not legacy_head_keys: + head = { + "markov_head.markov_w1.weight": head["markov_w1.weight"], + "markov_head.markov_w2.weight": head["markov_w2.weight"], + "confidence_head.proj.weight": head["confidence_proj.weight"], + "confidence_head.proj.bias": head["confidence_proj.bias"], + } + weights = { + "fc.weight": rnd(h, h * NUM_CAPTURE), + "hidden_norm.weight": rnd(h) + 1.0, + "norm.weight": rnd(h) + 1.0, + **head, + } for i in range(TINY["num_hidden_layers"]): p = f"layers.{i}." weights[p + "self_attn.q_proj.weight"] = rnd(nh * hd, h) @@ -134,6 +151,25 @@ def rnd(*shape): return weights +def _drafter_config_top_level_spelling(): + """RadixArk/Kimi-K3-DSpark as published: head switches at the top level. + + ``dflash_config`` carries only mask_token_id / target_layer_ids and the + confidence flag is spelled ``enable_confidence_head``. + """ + from transformers import Qwen3Config + + cfg = dict(TINY) + # Verbatim shape of the published config: no shift_label, no + # projector_type, and dflash_config holding nothing but these two keys. + cfg["dflash_config"] = {"mask_token_id": VOCAB - 2, "target_layer_ids": [0, 1]} + cfg["markov_rank"] = RANK + cfg["markov_head_type"] = "vanilla" + cfg["enable_confidence_head"] = True + cfg["confidence_head_with_markov"] = True + return Qwen3Config.from_dict(cfg) + + @pytest.fixture(scope="module") def standalone_drafter(): """The real ``Qwen3DSparkForCausalLM``, weights loaded, on the device.""" @@ -146,6 +182,67 @@ def standalone_drafter(): return drafter +@needs_gpu +def test_top_level_head_spelling_activates_the_heads(): + # The published RadixArk revision spells the switches top-level. Reading + # only dflash_config resolves markov_rank to 0, which drops markov_w1/w2 + # on the floor and falls back to the DFlash slot convention -- correct + # output, lower acceptance, nothing raised. + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models.modeling_dspark import Qwen3DSparkForCausalLM + + model_config = ModelConfig( + pretrained_config=_drafter_config_top_level_spelling(), attn_backend="TRTLLM" + ) + drafter = Qwen3DSparkForCausalLM(model_config, dflash_attention_backend="TRTLLM").to("cuda") + drafter.load_weights(_drafter_weights()) + + assert drafter._dspark_markov_rank == RANK + assert drafter._dspark_use_confidence_head is True + assert drafter.has_markov_head, "markov weights were dropped despite being in the checkpoint" + # Declared nowhere in the published config, so it rides on the DSpark + # default. False here would run slots 1..7 on a block_size-7 drafter and + # read the next request's anchor slot. + assert drafter._dspark_shift_label is True + + +@needs_gpu +def test_published_head_weight_names_load(): + # Both public drafters nest the head tensors under the module that owns + # them. Matching only the bare names leaves markov_w1/w2 in the dict handed + # to DFlash, which drops them -- or, once the rank resolves, raises + # "missing markov_w1.weight" on a checkpoint that plainly ships it. + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models.modeling_dspark import Qwen3DSparkForCausalLM + + for legacy in (False, True): + model_config = ModelConfig( + pretrained_config=_drafter_config_top_level_spelling(), attn_backend="TRTLLM" + ) + drafter = Qwen3DSparkForCausalLM(model_config, dflash_attention_backend="TRTLLM").to("cuda") + drafter.load_weights(_drafter_weights(legacy_head_keys=legacy)) + assert drafter.has_markov_head, f"legacy_head_keys={legacy}" + assert drafter.confidence_proj_weight is not None, f"legacy_head_keys={legacy}" + + +@needs_gpu +def test_markov_weights_without_a_resolvable_rank_raise(): + # The inverse of the missing-weights check: weights present, rank resolved + # to 0. Loading silently would cost acceptance with no signal. + from transformers import Qwen3Config + + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models.modeling_dspark import Qwen3DSparkForCausalLM + + cfg = dict(TINY) + cfg["dflash_config"] = {"mask_token_id": VOCAB - 2, "target_layer_ids": [0, 1]} + model_config = ModelConfig(pretrained_config=Qwen3Config.from_dict(cfg), attn_backend="TRTLLM") + drafter = Qwen3DSparkForCausalLM(model_config, dflash_attention_backend="TRTLLM").to("cuda") + + with pytest.raises(ValueError, match="markov_rank resolved to 0"): + drafter.load_weights(_drafter_weights()) + + def _spec_config(*, embedded): """Only the fields the routing and the worker read.""" return SimpleNamespace( @@ -154,7 +251,7 @@ def _spec_config(*, embedded): _use_shared_kv_cache=False, _allow_separate_draft_kv_cache=True, # K == block_size under the dspark shift_label convention, which is - # also what DSparkWorker validates before it touches anything else. + # also what DSv4DSparkWorker validates before it touches anything else. max_draft_len=BLOCK_SIZE, attention_backend="TRTLLM", ) @@ -179,6 +276,8 @@ def test_routed_worker_initializes_against_a_real_standalone_drafter(standalone_ worker = spec_utils.get_spec_worker( spec_config, model_config=None, mapping=Mapping(), use_separate_draft_kv_cache=True ) + # DFlash lineage (all the drafting plumbing), DSpark leaf (the two heads). + assert isinstance(worker, DSparkWorker) assert isinstance(worker, DFlashWorker) worker.set_draft_model(standalone_drafter) @@ -198,9 +297,9 @@ def test_routed_worker_initializes_against_a_real_standalone_drafter(standalone_ def test_routed_worker_sees_the_dspark_heads(standalone_drafter): """The heads moved to the DSpark subclass must stay visible to the worker. - ``DFlashWorker`` probes them defensively (``getattr(..., False)``), so a - drafter that lost them degrades to plain DFlash silently -- lower acceptance, - no error. These are the two probes the block-draft step makes. + ``DSparkWorker`` probes them defensively (``getattr(..., False)``), + so a drafter that lost them degrades to plain DFlash silently -- lower + acceptance, no error. These are the two probes the block-draft step makes. """ assert getattr(standalone_drafter, "has_markov_head", False) is True assert getattr(standalone_drafter, "_dspark_shift_label", False) is True @@ -212,12 +311,93 @@ def test_routed_worker_sees_the_dspark_heads(standalone_drafter): def test_embedded_worker_cannot_drive_a_standalone_drafter(standalone_drafter): """Witness for why the routing has to follow the flavour. - ``DSparkWorker`` serves the embedded DeepSeek-V4-Pro draft and reads + ``DSv4DSparkWorker`` serves the embedded DeepSeek-V4-Pro draft and reads V4-draft-only attributes. Handing it a standalone drafter is the exact mis-route that reached production, so pin the failure rather than trusting the routing test alone to stay correct. """ - worker = DSparkWorker(_spec_config(embedded=True), Mapping()) + worker = DSv4DSparkWorker(_spec_config(embedded=True), Mapping()) spec_metadata, _ = _lazy_init_args() with pytest.raises(AttributeError, match="num_stages"): worker._lazy_init(standalone_drafter, spec_metadata) + + +@needs_gpu +def test_dspark_overrides_read_the_drafter_not_a_hardcoded_policy(standalone_drafter): + """The two DSpark policies must come off the drafter, not the class. + + ``DSparkWorker`` exists only to override the block-output slot + convention and the Markov bias. Both are declared by the checkpoint, so a + DSpark drafter trained with the legacy DFlash slot layout (or without a + Markov head) must still get the base-class behaviour. Hardcoding the + dspark answer in the subclass would pass every mode/flavour dispatch test + while silently mis-slotting such a drafter. + """ + worker = DSparkWorker(_spec_config(embedded=False), Mapping()) + worker.set_draft_model(standalone_drafter) + + # shift_label on -> slots 0..K-1 (anchor slot predicts the first draft + # token); the DFlash base would return slots 1..K for the same inputs. + ids = worker._draft_slot_ids( + standalone_drafter, num_gens=2, block_size=BLOCK_SIZE, num_draft_tokens=3 + ) + assert ids.tolist() == [0, 1, 2, BLOCK_SIZE, BLOCK_SIZE + 1, BLOCK_SIZE + 2] + base_ids = DFlashWorker._draft_slot_ids( + worker, standalone_drafter, num_gens=2, block_size=BLOCK_SIZE, num_draft_tokens=3 + ) + assert base_ids.tolist() == [1, 2, 3, BLOCK_SIZE + 1, BLOCK_SIZE + 2, BLOCK_SIZE + 3] + + # A drafter without a Markov head falls through untouched, even on the + # DSpark subclass. + logits = torch.randn(2, 3, 8, device="cuda") + object.__setattr__(standalone_drafter, "_dspark_markov_rank", 0) + standalone_drafter.markov_w1 = None + assert standalone_drafter.has_markov_head is False + out = worker._refine_block_logits(standalone_drafter, logits, {}, None) + assert out is logits + + +@needs_gpu +def test_tp_sharded_markov_chain_matches_the_unsharded_result(): + # The sharded branch is unreachable on the validated TEP16/attention-DP + # config -- there the draft head is full-vocab, so vocab_slice is None and + # no gather runs. Exercise it here with a fake two-rank split: each rank + # sees its own logit columns and a markov_w2 sliced the same way, and the + # chain must agree column-for-column with the unsharded computation. + from tensorrt_llm._torch.models.modeling_speculative import dspark_markov_chain_logits + + torch.manual_seed(0) + vocab, rank, batch, block, tp = 32, 4, 3, 5, 2 + shard = vocab // tp + w1 = torch.randn(vocab, rank, device="cuda") + w2 = torch.randn(vocab, rank, device="cuda") + base = torch.randn(batch, block, vocab, device="cuda") + anchor = torch.randint(0, vocab, (batch,), device="cuda") + + full = dspark_markov_chain_logits(base, anchor, w1, w2) + + for tp_rank in range(tp): + vocab_slice = slice(tp_rank * shard, (tp_rank + 1) * shard) + # Stands in for greedy_sample_draft_with_tp_gather: the real one + # all-gathers local (index, value) pairs, so the chain always advances + # on the global argmax rather than this rank's local one. + sharded = dspark_markov_chain_logits( + base[..., vocab_slice], + anchor, + w1, + w2[vocab_slice], + argmax_fn=_make_global_argmax(full, tp_rank, shard), + ) + torch.testing.assert_close(sharded, full[..., vocab_slice]) + + +def _make_global_argmax(full_logits, tp_rank, shard): + """Replay the global argmax the TP gather would have produced.""" + step = {"i": 0} + + def argmax_fn(_step_logits): + i = step["i"] + step["i"] += 1 + return full_logits[:, i].argmax(dim=-1) + + return argmax_fn diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_flavour_dispatch.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_flavour_dispatch.py index 76b798534eb6..8b57cc6c6b53 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_flavour_dispatch.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_flavour_dispatch.py @@ -68,6 +68,7 @@ def _configs( architectures=None, attention_backend="TRTLLM", embedded=False, + top_level=None, ): """Duck-typed (target ModelConfig, draft ModelConfig) for dispatch-only asserts.""" model_config = SimpleNamespace( @@ -83,6 +84,7 @@ def _configs( model_type=model_type, architectures=architectures, dflash_config=dflash_config, + **(top_level or {}), ) ) return model_config, draft_config @@ -172,6 +174,25 @@ def test_dflash_refuses_a_dspark_drafter(stub_dflash): assert "decoding_type" in message, "the error must say how to fix the config" +def test_dflash_refuses_a_top_level_spelling_dspark_drafter(stub_dflash): + # RadixArk/Kimi-K3-DSpark as published: the head switches sit at the top + # level and dflash_config carries only mask_token_id / target_layer_ids. + # A reader that looks in dflash_config alone misses exactly the drafter + # this guard exists to catch. + model_config, draft_config = _configs( + dflash_config={"mask_token_id": 163824, "target_layer_ids": [7, 23, 51, 67, 83]}, + top_level={ + "markov_rank": 256, + "markov_head_type": "vanilla", + "enable_confidence_head": True, + "block_size": 7, + }, + ) + + with pytest.raises(ValueError, match="DSpark"): + modeling_dflash._build_dflash_draft(model_config, draft_config, None, None) + + @pytest.mark.parametrize( "field,value", [ @@ -252,12 +273,16 @@ def test_declares_dspark_heads(dflash_config, expected): # # The worker, the spec metadata and the separate-draft-KV-cache decision must # follow the same flavour flag the builder follows. When they disagree, -# DSparkWorker gets a standalone drafter and dies reaching for V4-draft-only +# DSv4DSparkWorker gets a standalone drafter and dies reaching for V4-draft-only # attributes (num_stages, write_context_windows) -- only at the first forward, # long after the engine reported a successful build. # -------------------------------------------------------------------------- -_WORKER_SENTINELS = {"DFlashWorker": object(), "DSparkWorker": object()} +_WORKER_SENTINELS = { + "DFlashWorker": object(), + "DSparkWorker": object(), + "DSv4DSparkWorker": object(), +} # SimpleNamespace rather than object(): get_spec_metadata assigns # ``metadata.enable_penalty`` on whatever it built, which a bare object rejects. _METADATA_SENTINELS = { @@ -296,8 +321,8 @@ def _spec_config(mode, *, embedded, allow_separate_kv=True): @pytest.mark.parametrize( "mode,embedded,expected", [ - (SpeculativeDecodingMode.DSPARK, True, "DSparkWorker"), - (SpeculativeDecodingMode.DSPARK, False, "DFlashWorker"), + (SpeculativeDecodingMode.DSPARK, True, "DSv4DSparkWorker"), + (SpeculativeDecodingMode.DSPARK, False, "DSparkWorker"), (SpeculativeDecodingMode.DFLASH, False, "DFlashWorker"), ], ) 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 220aea0656d5..bd917658f9b3 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -16,7 +16,7 @@ Covers the framework-side logic that does NOT need the full draft model: ``DSparkSpecMetadata`` hidden-state capture (incl. the mHC hc-mean reduction) -and ``DSparkWorker`` slot / rolling-KV-window management. The end-to-end block +and ``DSv4DSparkWorker`` slot / rolling-KV-window management. The end-to-end block draft and acceptance path is covered by the DSpark test in ``integration/defs/accuracy/test_llm_api_pytorch.py``. """ @@ -26,7 +26,7 @@ import pytest import torch -from tensorrt_llm._torch.speculative.dspark import DSparkSpecMetadata, DSparkWorker +from tensorrt_llm._torch.speculative.dspark import DSparkSpecMetadata, DSv4DSparkWorker from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode pytestmark = pytest.mark.skipif( @@ -101,7 +101,7 @@ def _make_worker(): ) from tensorrt_llm.mapping import Mapping - return DSparkWorker(cfg, Mapping()) + return DSv4DSparkWorker(cfg, Mapping()) def _fake_draft_model(num_stages=3, window_size=128, head_dim=64): From 22cb5d00ffed96ca24618f2f9f35d8fa6dfce8c4 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Wed, 19 Aug 2026 20:20:47 -0700 Subject: [PATCH 13/21] [None][chore] Report the actual worker class in the DFlash init log StandaloneDSparkWorker does not override __init__, so the hardcoded name made a standalone-DSpark run indistinguishable from a plain DFlash one in the logs. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/speculative/dflash.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index b16ff043a6a1..10b2dc77af52 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -239,8 +239,12 @@ def __init__( f"DFlash: acceptance-statistics recording enabled -> {self._accept_stats.path}" ) + # type(self), not a literal: DSparkWorker does not override + # __init__, so a hardcoded name reports the base class and the log + # cannot evidence which worker a run actually used. logger.info( - f"DFlashWorker initialized with use_separate_draft_kv_cache={use_separate_draft_kv_cache}" + f"{type(self).__name__} initialized with " + f"use_separate_draft_kv_cache={use_separate_draft_kv_cache}" ) @property From 927a6d5f4c4862a12ad4bfd1dd56b8e3d402c344 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Wed, 19 Aug 2026 20:33:05 -0700 Subject: [PATCH 14/21] [None][chore] Document what the draft-slot clamp masks At block_size == K with shift_label off the slot ids run 1..K, so each request reads the next one's slot 0 and the last overruns the block. The clamp turns that misconfiguration into lost acceptance, never an error. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/speculative/dflash.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 10b2dc77af52..1c1320e776e7 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -706,6 +706,9 @@ def _forward_impl( # resolved through _draft_slot_ids. block_size = self._compute_block_size gen_gather_ids = self._draft_slot_ids(draft_model, num_gens, block_size, K) + # Shields only the last request: at block_size == K with + # shift_label off, slots run 1..K, so every request reads the + # next one's slot 0 and the last overruns. Degrades, never raises. gen_gather_ids = gen_gather_ids.clamp(max=hidden_states_out.shape[0] - 1) gen_logits = draft_model.logits_processor( From 02b90a8382359f11bac2611d89fadb55a6d84445 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Thu, 20 Aug 2026 23:27:45 -0700 Subject: [PATCH 15/21] [None][fix] Capture the aggregated stream for the K3 DSpark drafter The drafter is distilled on the pre-norm softmax mixture its next consumer sees, not on the raw prefix sum a layer returns. Layer i+1 already computes that tensor as its own attention-side mixture, so the tap rides along. Ground truth: SGLang kimi_k3.py:2589 _dspark_capture_stream, whose attn_res-is-None fallback is what we were capturing; attn_residual.py:285 aggregate_stream matches _apply_attn_res row-for-row. Cross-check on a separate harness (0-shot chat, RadixArk drafter, 1319 questions): AR 66.9% -> 71.4%, acceptance length 5.683 -> 6.005, against SGLang's 6.089 on the same checkpoints. Accuracy unchanged, as expected for a draft-side fix. Caveats: every arm carried a scratch acceptance-histogram patch, and there is no clean same-config repeat, so this is an argument from magnitude rather than a measured interval. Signed-off-by: Zhenhuan Chen --- .../_torch/models/modeling_kimi_linear.py | 65 ++++++++-- tensorrt_llm/_torch/speculative/dflash.py | 8 +- .../test_kimi_k3_dflash_scaffold.py | 118 ++++++++++++++++++ 3 files changed, 179 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 077f714001a7..1fd5e8c47c87 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -1682,12 +1682,20 @@ def forward( block_residual: torch.Tensor, num_snapshots: int, attn_metadata: AttentionMetadata, + capture: Optional[Tuple[Any, int]] = None, ) -> Tuple[torch.Tensor, int]: """Port of HF ``KimiDecoderLayer._forward_attn_residual`` (per token). ``block_residual`` is a preallocated snapshot bank in kernel-native ``[K_max, M, H]`` layout. Returns the running prefix sum and the number of valid bank rows. + + ``capture`` is ``(spec_metadata, layer_id)`` and taps the DSpark aux + stream for the layer BEFORE this one: the aggregated stream for layer j + is by definition what its next consumer sees, so the mixture computed + below already is it. Reading it here beats recomputing it, and is only + possible because K3 asserts pp_size == 1 -- layer j+1 is always local. + PP support would need a recompute at the rank boundary. """ prefix_sum = hidden_states valid_block_residual = block_residual[:num_snapshots] @@ -1699,6 +1707,8 @@ def forward( self.self_attention_res_proj, self.self_attention_res_norm, ) + if capture is not None: + capture[0].maybe_capture_hidden_states(capture[1], hidden_states, None) if self.layer_idx % self.attn_res_block_size == 0: block_residual[num_snapshots].copy_(prefix_sum) @@ -1813,19 +1823,52 @@ def forward( hidden_states.shape[1], ) num_snapshots = 0 - for layer in self.layers: + capture_set = ( + getattr(spec_metadata, "_capture_layer_set", None) + if spec_metadata is not None + else None + ) + for i, layer in enumerate(self.layers): + # DFlash/DSpark hidden-state capture. The drafter is distilled on + # the aggregated stream value -- the pre-norm softmax mixture its + # next consumer sees -- not on the raw prefix sum a layer returns, + # which is SGLang's fallback for models without the + # attention-residual scheme. Capturing the prefix sum costs 4.5pt + # of draft acceptance on K3 + RadixArk DSpark (AR 66.9% -> 71.4%). + # The tap fires inside layer i+1, which computes that tensor + # anyway; see its forward docstring. Ground truth: SGLang + # kimi_k3.py:2697 _dspark_capture_stream, attn_residual.py:313 + # aggregate_stream_torch. + capture = None + if ( + spec_metadata is not None + and i > 0 + and (capture_set is None or self.layers[i - 1].layer_idx in capture_set) + ): + capture = (spec_metadata, self.layers[i - 1].layer_idx) hidden_states, num_snapshots = layer( - hidden_states, block_residual, num_snapshots, attn_metadata + hidden_states, block_residual, num_snapshots, attn_metadata, capture=capture ) - if spec_metadata is not None: - # DFlash hidden-state capture. K3's attn-residual scheme - # already folds the residual into the running prefix sum - # returned by each layer, so unlike Qwen3/Llama we pass the - # full hidden state with residual=None. Whether the drafter - # is trained against this prefix sum or some other tap point - # must be confirmed against the K3 drafter training recipe - # before real weights are used. - spec_metadata.maybe_capture_hidden_states(layer.layer_idx, hidden_states, None) + + # The last layer has no successor, so this one recompute is + # unavoidable -- output-side score weights, matching SGLang's + # layer_idx + 1 >= end_layer branch. Unreachable for K3's capture set + # against 93 layers; kept so a set that does include the final layer + # gets the right tensor rather than the raw prefix sum. + if spec_metadata is not None and len(self.layers) > 0: + last = self.layers[-1] + if capture_set is None or last.layer_idx in capture_set: + tail = ( + _apply_attn_res( + hidden_states, + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + ) + if num_snapshots > 0 + else hidden_states + ) + spec_metadata.maybe_capture_hidden_states(last.layer_idx, tail, None) hidden_states = _apply_attn_res( hidden_states, diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 1c1320e776e7..40dc16711f51 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -146,7 +146,13 @@ def is_layer_capture(self, layer_id: int) -> bool: def maybe_capture_hidden_states( self, layer_id: int, hidden_states: torch.Tensor, residual: Optional[torch.Tensor] = None ) -> None: - """Capture hidden states from a target model layer into the buffer.""" + """Capture hidden states from a target model layer into the buffer. + + The ``residual`` convention is model-specific: Qwen3/Llama-style callers + pass the pre-add pair and this folds them, while K3 hands in an + already-mixed aggregated stream value and passes ``residual=None`` + (see the DSpark tap in ``modeling_kimi_linear.py``). + """ if self.captured_hidden_states is None: return i = self._layer_to_idx.get(layer_id) 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 ef778488ad8a..51af02a03888 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 @@ -342,3 +342,121 @@ def test_dflash_spec_metadata_capture_prefix_sum_convention(): assert captured.shape == (max_tokens, 2 * hidden_size) torch.testing.assert_close(captured[:, :hidden_size], h1) torch.testing.assert_close(captured[:, hidden_size:], h3) + + +def test_capture_taps_the_next_layers_aggregated_stream(): + """The tap reads layer j+1's mixture, not layer j's prefix sum. + + Worth 4.5pt of draft acceptance on K3 + RadixArk DSpark (AR 66.9% -> + 71.4%), and it degrades silently: capturing the prefix sum instead is + still a valid tensor of the right shape, so nothing raises and only + acceptance moves. The two neighbouring tests cover buffer routing and + the forward signature; both still pass if the tap regresses. + + Layers here are stubs that reproduce the one contract the real layer has + with this loop -- compute the attention-side mixture, and when handed a + ``capture`` tuple, record it under the layer id in that tuple. Giving + each stub its own proj/norm weights is what makes the assertions able to + tell "layer j+1's mixture" from "layer j's", and the model's output-side + weights from either. Also covers the no-successor tail branch, which uses + output_attn_res_proj/norm and no layer can produce. + """ + pytest.importorskip("fla") + from types import SimpleNamespace + + from tensorrt_llm._torch.models.modeling_kimi_linear import ( + KimiK3RMSNorm, + KimiLinearModel, + _apply_attn_res, + ) + + torch.manual_seed(0) + num_layers, num_tokens, hidden = 4, 3, 8 + + def _weights(scale): + proj = torch.nn.Linear(hidden, 1, bias=False, dtype=torch.float32) + norm = KimiK3RMSNorm(hidden, eps=1e-6, dtype=torch.float32) + with torch.no_grad(): + proj.weight.copy_(torch.randn(1, hidden) * scale) + norm.weight.copy_(torch.randn(hidden) * scale + 1.0) + return proj, norm + + layer_w = [_weights(0.5 + i) for i in range(num_layers)] + out_proj, out_norm = _weights(9.0) + + seen = {} + + class _StubLayer: + def __init__(self, idx): + self.layer_idx = idx + self.proj, self.norm = layer_w[idx] + + def __call__(self, hidden_states, block_residual, num_snapshots, attn_metadata, capture): + # What the real layer computes on the way into its attention, and + # hands to the tap: the aggregated stream the previous layer's + # consumer sees. + mixture = ( + _apply_attn_res(hidden_states, block_residual[:num_snapshots], self.proj, self.norm) + if num_snapshots > 0 + else hidden_states + ) + if capture is not None: + spec_md, layer_id = capture + spec_md.maybe_capture_hidden_states(layer_id, mixture, None) + block_residual[num_snapshots] = hidden_states * (self.layer_idx + 1) + # A prefix sum that is never equal to the mixture above. + return hidden_states + (self.layer_idx + 1), num_snapshots + 1 + + layers = [_StubLayer(i) for i in range(num_layers)] + last_idx = num_layers - 1 + spec_md = SimpleNamespace( + _capture_layer_set=frozenset({0, 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, + ) + + KimiLinearModel.forward( + fake, + attn_metadata=SimpleNamespace(num_tokens=num_tokens), + input_ids=torch.zeros(num_tokens, dtype=torch.int32), + spec_metadata=spec_md, + ) + + assert set(seen) == {0, last_idx}, "capture_set not honoured" + + # Replay layer 0 to rebuild what each candidate tensor would have been. + br = torch.empty(num_layers, num_tokens, hidden) + br[0] = embeds * 1 + after_l0 = embeds + 1 + + proj1, norm1 = layer_w[1] + torch.testing.assert_close(seen[0], _apply_attn_res(after_l0, br[:1], proj1, norm1)) + # The two ways this regresses, both silent: + proj0, norm0 = layer_w[0] + assert not torch.allclose(seen[0], _apply_attn_res(after_l0, br[:1], proj0, norm0)), ( + "tap used layer j's own weights instead of its successor's" + ) + assert not torch.allclose(seen[0], after_l0), "tap captured the raw prefix sum" + + # Tail: no successor exists, so it must use the model's output-side + # weights -- not the last layer's, and not the bare prefix sum. + h = embeds + for i in range(num_layers): + br[i] = h * (i + 1) + h = h + (i + 1) + torch.testing.assert_close( + seen[last_idx], _apply_attn_res(h, br[:num_layers], out_proj, out_norm) + ) + proj_last, norm_last = layer_w[last_idx] + assert not torch.allclose( + 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" From 980f31421e9090ac0afc1bf8e9b44731fa80d08c Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Fri, 21 Aug 2026 01:04:30 -0700 Subject: [PATCH 16/21] [None][refactor] Collapse the standalone DSpark drafter to one shape-named class DSparkDrafterForCausalLM + the empty Qwen3DSparkForCausalLM become GQADSparkForCausalLM. The constraint is the attention shape, not the model: DFlash already runs one block decode over qwen3, llama and gpt_oss drafters, so a per-model subclass is empty by construction. _DSPARK_DRAFTERS_BY_MODEL_TYPE is gone. It keyed on model_type a second time after DFlashForCausalLM.__init__ had already resolved the backbone through the model registry; what it actually guarded was the block decode's GQA precondition. That check now lives in the DFlash base, where DFlash needs it too, and fails at construction with the offending layer rather than deep inside _build_fused_kv_buffers. Also drops two unit tests fully covered by others: the per-field DFlash refusal (subsumed by the declares_dspark_heads truth table plus the builder-level refusal) and the no-spec-config gate case. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/models/modeling_dflash.py | 37 ++++++++ tensorrt_llm/_torch/models/modeling_dspark.py | 77 ++++++----------- tensorrt_llm/_torch/speculative/dspark.py | 2 +- .../test_dspark_drafter_worker_contract.py | 20 ++--- .../test_dspark_flavour_dispatch.py | 84 ++++++++++--------- .../test_kimi_k3_dspark_semantics.py | 8 +- .../test_kimi_k3_spec_mode_gate.py | 4 - 7 files changed, 124 insertions(+), 108 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index 7f77ae8d659b..4ea0222e0c52 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -181,8 +181,45 @@ def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"): # non-causal block attention). Subclasses opt in. self._context_input_layernorm = False self._sliding_layers_causal = False + self._validate_gqa_shape() self._warn_inferred_attention_windows() + def _validate_gqa_shape(self): + """Reject a backbone the block decode cannot express. + + The hand-written block decode is GQA-shaped throughout: it splits a + fused ``qkv_proj`` by (num_heads, num_kv_heads, head_dim), fuses K/V + across layers on one uniform head dim, and shares a single RoPE cache. + The registry will happily build a backbone that violates any of that -- + an MLA drafter has no per-head K/V at all -- and without this the + failure surfaces much later inside ``_build_fused_kv_buffers`` or, worse, + as silently mis-sliced weights. Raise at construction instead. + """ + layers = getattr(self.model, "layers", None) + if not layers: + return + for idx, layer in enumerate(layers): + attn = getattr(layer, "self_attn", None) + if attn is None or not hasattr(attn, "qkv_proj"): + raise ValueError( + f"DFlash block decode requires a fused self_attn.qkv_proj on " + f"every draft layer, but layer {idx} of draft backbone " + f"{type(self.config).__name__} has none. Backbones that do not " + "project per-head Q/K/V (e.g. MLA) need their own block decode." + ) + num_kv_heads = layers[0].self_attn.num_key_value_heads + mismatched = [ + idx + for idx, layer in enumerate(layers[1:], start=1) + if layer.self_attn.num_key_value_heads != num_kv_heads + ] + if mismatched: + raise ValueError( + "DFlash fuses draft K/V across layers and needs one uniform " + f"num_key_value_heads, but layers {mismatched} differ from layer 0 " + f"({num_kv_heads})." + ) + @staticmethod def _rope_signature(attn): """Return the effective RoPE configuration used by an attention layer.""" diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 1db272fc0e59..c85dd570db93 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -38,7 +38,7 @@ Most of this module is this flavour. * **standalone** — the drafter ships as its own checkpoint and shares nothing with the target but the vocabulary and the captured hidden states. Its block - decode is DFlash's, so :class:`DSparkDrafterForCausalLM` subclasses + decode is DFlash's, so :class:`GQADSparkForCausalLM` subclasses ``DFlashForCausalLM`` and adds the Markov head, the confidence head and the shift_label convention. See "Standalone DSpark drafters" near the bottom. @@ -67,9 +67,9 @@ class or the edge would become a cycle. ``[bonus_token, noise, ...]`` block input) and ``dspark_propose`` (Markov refinement + static confidence truncation). -4. **Standalone DSpark drafters** — :class:`DSparkDrafterForCausalLM` and its - per-backbone subclasses, plus the two-level ``_build_dspark_draft`` dispatch - that picks between the two flavours. +4. **Standalone DSpark drafters** — :class:`GQADSparkForCausalLM`, plus the + ``_build_dspark_draft`` dispatch that picks between the two flavours on + deployment form alone. The per-stage *backbone* forward (block attention whose K/V derive from ``main_x``, + MoE + mHC) is brought up and numerically validated against the real @@ -1931,12 +1931,12 @@ def load_weights_from_target_model(self, target_model): } -class DSparkDrafterForCausalLM(DFlashForCausalLM): - """DSpark drafter built from a standalone draft checkpoint. +class GQADSparkForCausalLM(DFlashForCausalLM): + """DSpark drafter on a GQA-shaped backbone, from a standalone checkpoint. - Adds the DSpark head set on top of the generic DFlash block decode: + Adds the DSpark head set on top of the DFlash block decode: - - the vanilla Markov intra-block logit bias, applied by ``DFlashWorker`` + - 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); @@ -1945,9 +1945,14 @@ class DSparkDrafterForCausalLM(DFlashForCausalLM): Confidence-scheduled verification is not implemented yet: ``confidence_proj`` is loaded but unused, and drafting always proposes the full K tokens. - The draft backbone itself is whatever the drafter config resolves to through - the model registry, so this class is backbone-agnostic; per-backbone - subclasses exist to carry backbone-specific block-decode overrides. + 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. Reference: arXiv 2607.05147; deepseek-ai/DeepSpec. """ @@ -2082,31 +2087,6 @@ def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): return super().load_weights(weights, weight_mapper=weight_mapper, **kwargs) -class Qwen3DSparkForCausalLM(DSparkDrafterForCausalLM): - """DSpark drafter on a Qwen3-style GQA draft backbone. - - Overrides nothing today: the backbone is built from the drafter config - through the model registry, and the DSpark head set is backbone-independent, - so ``DSparkDrafterForCausalLM`` already covers this combination end to end. - - It exists as the explicit dispatch target for ``model_type: qwen3``, which - keeps the supported matrix visible in class names rather than buried in a - builder, and as the seat for backbone-specific overrides when they arrive. - They will: an MLA-backboned drafter cannot reuse this block decode, which - assumes a fused ``qkv_proj`` and one uniform head dim across Q/K/V. - """ - - -# Standalone DSpark drafters by the draft checkpoint's ``model_type``. The key -# is the backbone family, not the target model: the same drafter class serves -# any target, and a target-specific one would have nothing to hold -- the -# target-side half of DSpark is the hidden-state capture, which lives in each -# target's own modeling file. -_DSPARK_DRAFTERS_BY_MODEL_TYPE = { - "qwen3": Qwen3DSparkForCausalLM, -} - - def draft_is_embedded_in_target(model_config) -> bool: """True when the DSpark draft weights live inside the target checkpoint. @@ -2154,19 +2134,15 @@ def _build_dspark_draft(model_config, draft_config, lm_head, model): block_size=model_config.spec_config.block_size, ) - model_type = getattr(draft_config.pretrained_config, "model_type", None) - drafter_cls = _DSPARK_DRAFTERS_BY_MODEL_TYPE.get(model_type) - if drafter_cls is None: - supported = ", ".join(sorted(_DSPARK_DRAFTERS_BY_MODEL_TYPE)) - raise NotImplementedError( - f"No standalone DSpark drafter for draft model_type {model_type!r}. " - f"Supported draft model_type values: {supported}. The dispatch keys " - "on the draft backbone, so a drafter is supported once its backbone " - "has a block-decode implementation here; MLA-backboned drafters " - "(e.g. Inferact/Kimi-K3-DSpark, model_type 'k3_dspark') need one and " - "are a follow-up." - ) - return drafter_cls( + # 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( draft_config, dflash_attention_backend=model_config.spec_config.attention_backend, ) @@ -2178,8 +2154,7 @@ def _build_dspark_draft(model_config, draft_config, lm_head, model): "DSv4DSparkDraftModel", "DSv4DSparkForCausalLM", # Standalone flavour. - "DSparkDrafterForCausalLM", - "Qwen3DSparkForCausalLM", + "GQADSparkForCausalLM", "draft_is_embedded_in_target", "validate_dspark_eplb_layer_base", "validate_dspark_eplb_stage_layers", diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index bbdaddd1043d..e6ddbc5904e1 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -716,7 +716,7 @@ class DSparkWorker(DFlashWorker): head-driven policies are overridden here: the block-output slot convention (``shift_label``) and the Markov intra-block logit bias. - Mirrors the model side, where ``DSparkDrafterForCausalLM`` extends + Mirrors the model side, where ``GQADSparkForCausalLM`` extends ``DFlashForCausalLM`` with the same two heads. Naming: this is the unqualified DSpark worker because a separately diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_drafter_worker_contract.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_drafter_worker_contract.py index 07015c803878..d2250602d7c1 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_drafter_worker_contract.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_drafter_worker_contract.py @@ -20,7 +20,7 @@ does not have: the two agree on paper and diverge on first contact. The routing bug this file pins surfaced only as - AttributeError: 'Qwen3DSparkForCausalLM' object has no attribute 'num_stages' + AttributeError: 'GQADSparkForCausalLM' object has no attribute 'num_stages' five minutes into a 16-GPU run, after the weights had loaded -- because nothing below the factory was ever exercised. @@ -172,12 +172,12 @@ def _drafter_config_top_level_spelling(): @pytest.fixture(scope="module") def standalone_drafter(): - """The real ``Qwen3DSparkForCausalLM``, weights loaded, on the device.""" + """The real ``GQADSparkForCausalLM``, weights loaded, on the device.""" from tensorrt_llm._torch.model_config import ModelConfig - from tensorrt_llm._torch.models.modeling_dspark import Qwen3DSparkForCausalLM + from tensorrt_llm._torch.models.modeling_dspark import GQADSparkForCausalLM model_config = ModelConfig(pretrained_config=_drafter_config(), attn_backend="TRTLLM") - drafter = Qwen3DSparkForCausalLM(model_config, dflash_attention_backend="TRTLLM").to("cuda") + drafter = GQADSparkForCausalLM(model_config, dflash_attention_backend="TRTLLM").to("cuda") drafter.load_weights(_drafter_weights()) return drafter @@ -189,12 +189,12 @@ def test_top_level_head_spelling_activates_the_heads(): # on the floor and falls back to the DFlash slot convention -- correct # output, lower acceptance, nothing raised. from tensorrt_llm._torch.model_config import ModelConfig - from tensorrt_llm._torch.models.modeling_dspark import Qwen3DSparkForCausalLM + from tensorrt_llm._torch.models.modeling_dspark import GQADSparkForCausalLM model_config = ModelConfig( pretrained_config=_drafter_config_top_level_spelling(), attn_backend="TRTLLM" ) - drafter = Qwen3DSparkForCausalLM(model_config, dflash_attention_backend="TRTLLM").to("cuda") + drafter = GQADSparkForCausalLM(model_config, dflash_attention_backend="TRTLLM").to("cuda") drafter.load_weights(_drafter_weights()) assert drafter._dspark_markov_rank == RANK @@ -213,13 +213,13 @@ def test_published_head_weight_names_load(): # to DFlash, which drops them -- or, once the rank resolves, raises # "missing markov_w1.weight" on a checkpoint that plainly ships it. from tensorrt_llm._torch.model_config import ModelConfig - from tensorrt_llm._torch.models.modeling_dspark import Qwen3DSparkForCausalLM + from tensorrt_llm._torch.models.modeling_dspark import GQADSparkForCausalLM for legacy in (False, True): model_config = ModelConfig( pretrained_config=_drafter_config_top_level_spelling(), attn_backend="TRTLLM" ) - drafter = Qwen3DSparkForCausalLM(model_config, dflash_attention_backend="TRTLLM").to("cuda") + drafter = GQADSparkForCausalLM(model_config, dflash_attention_backend="TRTLLM").to("cuda") drafter.load_weights(_drafter_weights(legacy_head_keys=legacy)) assert drafter.has_markov_head, f"legacy_head_keys={legacy}" assert drafter.confidence_proj_weight is not None, f"legacy_head_keys={legacy}" @@ -232,12 +232,12 @@ def test_markov_weights_without_a_resolvable_rank_raise(): from transformers import Qwen3Config from tensorrt_llm._torch.model_config import ModelConfig - from tensorrt_llm._torch.models.modeling_dspark import Qwen3DSparkForCausalLM + from tensorrt_llm._torch.models.modeling_dspark import GQADSparkForCausalLM cfg = dict(TINY) cfg["dflash_config"] = {"mask_token_id": VOCAB - 2, "target_layer_ids": [0, 1]} model_config = ModelConfig(pretrained_config=Qwen3Config.from_dict(cfg), attn_backend="TRTLLM") - drafter = Qwen3DSparkForCausalLM(model_config, dflash_attention_backend="TRTLLM").to("cuda") + drafter = GQADSparkForCausalLM(model_config, dflash_attention_backend="TRTLLM").to("cuda") with pytest.raises(ValueError, match="markov_rank resolved to 0"): drafter.load_weights(_drafter_weights()) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_flavour_dispatch.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_flavour_dispatch.py index 8b57cc6c6b53..b3ee48d4eb10 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_flavour_dispatch.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_flavour_dispatch.py @@ -16,8 +16,10 @@ ``DSpark`` ships in two flavours -- embedded in the target checkpoint (DeepSeek-V4-Pro's ``mtp.*``) or standalone with its own checkpoint -- and one -builder picks between them, then picks the standalone backbone by the draft -checkpoint's ``model_type``. +builder picks between them on deployment form alone. There is no second +dispatch on ``model_type``: the standalone drafter's backbone is resolved one +layer down by the model registry, and the block decode's GQA precondition is +what rejects a backbone it cannot express. The DFlash half matters just as much: DFlash no longer implements the DSpark head set, so a drafter that declares it must be refused rather than served @@ -48,7 +50,7 @@ from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig _DSV4_SENTINEL = object() -_QWEN3_SENTINEL = object() +_GQA_SENTINEL = object() _DFLASH_SENTINEL = object() _LAGUNA_SENTINEL = object() @@ -94,11 +96,7 @@ def _configs( def stub_dspark(monkeypatch): """Replace the DSpark drafter classes with sentinel-returning stubs.""" monkeypatch.setattr(modeling_dspark, "DSv4DSparkForCausalLM", lambda *a, **k: _DSV4_SENTINEL) - monkeypatch.setattr( - modeling_dspark, - "_DSPARK_DRAFTERS_BY_MODEL_TYPE", - {"qwen3": lambda *a, **k: _QWEN3_SENTINEL}, - ) + monkeypatch.setattr(modeling_dspark, "GQADSparkForCausalLM", lambda *a, **k: _GQA_SENTINEL) monkeypatch.setattr(modeling_dspark, "validate_dspark_eplb_layer_base", lambda *a, **k: None) @@ -115,12 +113,15 @@ def stub_dflash(monkeypatch): # -------------------------------------------------------------------------- -def test_standalone_qwen3_drafter_selects_qwen3_dspark(monkeypatch, stub_dspark): - model_config, draft_config = _configs(model_type="qwen3", embedded=False) +@pytest.mark.parametrize("model_type", ["qwen3", "llama", "gpt_oss"]) +def test_standalone_drafter_selects_gqa_dspark(monkeypatch, stub_dspark, model_type): + # No per-model_type table: every GQA backbone the registry can build gets + # the same drafter class. A new GQA family must not need an entry here. + model_config, draft_config = _configs(model_type=model_type, embedded=False) built = modeling_dspark._build_dspark_draft(model_config, draft_config, None, None) - assert built is _QWEN3_SENTINEL + assert built is _GQA_SENTINEL def test_embedded_draft_selects_dsv4_dspark(monkeypatch, stub_dspark): @@ -132,15 +133,40 @@ def test_embedded_draft_selects_dsv4_dspark(monkeypatch, stub_dspark): assert built is _DSV4_SENTINEL -def test_unknown_standalone_model_type_lists_supported(monkeypatch, stub_dspark): - model_config, draft_config = _configs(model_type="llama", embedded=False) +def test_non_gqa_backbone_is_rejected_at_construction(): + # The builder no longer whitelists model_type; the block decode's GQA + # precondition is what rejects an MLA-shaped drafter, and it must do so at + # construction rather than deep inside _build_fused_kv_buffers. + from types import SimpleNamespace + + drafter = modeling_dflash.DFlashForCausalLM.__new__(modeling_dflash.DFlashForCausalLM) + drafter.model = SimpleNamespace( + layers=[SimpleNamespace(self_attn=SimpleNamespace())] # no qkv_proj -> MLA-like + ) + drafter.config = SimpleNamespace() - with pytest.raises(NotImplementedError) as excinfo: - modeling_dspark._build_dspark_draft(model_config, draft_config, None, None) + with pytest.raises(ValueError) as excinfo: + drafter._validate_gqa_shape() - message = str(excinfo.value) - assert "llama" in message - assert "qwen3" in message, "the error must list the supported draft model_type values" + assert "qkv_proj" in str(excinfo.value) + + +def test_layers_with_mismatched_kv_heads_are_rejected(): + from types import SimpleNamespace + + def _layer(nkv): + return SimpleNamespace( + self_attn=SimpleNamespace(qkv_proj=object(), num_key_value_heads=nkv) + ) + + drafter = modeling_dflash.DFlashForCausalLM.__new__(modeling_dflash.DFlashForCausalLM) + drafter.model = SimpleNamespace(layers=[_layer(8), _layer(8), _layer(4)]) + drafter.config = SimpleNamespace() + + with pytest.raises(ValueError) as excinfo: + drafter._validate_gqa_shape() + + assert "[2]" in str(excinfo.value) def test_standalone_drafter_receives_the_attention_backend(monkeypatch, stub_dspark): @@ -148,9 +174,9 @@ def test_standalone_drafter_receives_the_attention_backend(monkeypatch, stub_dsp def _capture(draft_config, *, dflash_attention_backend): seen["backend"] = dflash_attention_backend - return _QWEN3_SENTINEL + return _GQA_SENTINEL - monkeypatch.setattr(modeling_dspark, "_DSPARK_DRAFTERS_BY_MODEL_TYPE", {"qwen3": _capture}) + monkeypatch.setattr(modeling_dspark, "GQADSparkForCausalLM", _capture) model_config, draft_config = _configs(attention_backend="TRTLLM") modeling_dspark._build_dspark_draft(model_config, draft_config, None, None) @@ -193,24 +219,6 @@ def test_dflash_refuses_a_top_level_spelling_dspark_drafter(stub_dflash): modeling_dflash._build_dflash_draft(model_config, draft_config, None, None) -@pytest.mark.parametrize( - "field,value", - [ - ("markov_rank", 256), - ("use_confidence_head", True), - ("shift_label", True), - ("projector_type", "dspark"), - ], -) -def test_any_single_dspark_field_is_enough_to_refuse(stub_dflash, field, value): - # Each field alone means the drafter was trained under the DSpark - # convention; serving it as plain DFlash degrades it silently. - model_config, draft_config = _configs(dflash_config={"mask_token_id": 7, field: value}) - - with pytest.raises(ValueError, match="DSpark"): - modeling_dflash._build_dflash_draft(model_config, draft_config, None, None) - - def test_plain_dflash_drafter_is_unchanged(stub_dflash): model_config, draft_config = _configs( dflash_config={"mask_token_id": 7, "target_layer_ids": [0, 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 e22aea1e7442..1a061c8daa00 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 @@ -24,7 +24,7 @@ import torch.nn.functional as F from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM, dspark_layer_window_size -from tensorrt_llm._torch.models.modeling_dspark import Qwen3DSparkForCausalLM +from tensorrt_llm._torch.models.modeling_dspark import GQADSparkForCausalLM from tensorrt_llm._torch.models.modeling_speculative import ( dspark_markov_chain_logits, dspark_markov_step_bias, @@ -265,7 +265,7 @@ def _build_drafter(dspark: bool, weights): model_config = ModelConfig(pretrained_config=_tiny_config(dspark), attn_backend="TRTLLM") # The DSpark head set lives in the DSpark drafter, not in the DFlash base. - drafter_cls = Qwen3DSparkForCausalLM if dspark else DFlashForCausalLM + drafter_cls = GQADSparkForCausalLM if dspark else DFlashForCausalLM drafter = drafter_cls(model_config).to("cuda") # Drop dspark head tensors for the plain drafter (schema without them). if not dspark: @@ -414,7 +414,7 @@ def test_dspark_causal_config_rejected(): cfg = _tiny_config(True) cfg.dflash_config = dict(cfg.dflash_config, causal=True) with pytest.raises(ValueError, match="non-causal DSpark convention"): - Qwen3DSparkForCausalLM(ModelConfig(pretrained_config=cfg, attn_backend="TRTLLM")) + GQADSparkForCausalLM(ModelConfig(pretrained_config=cfg, attn_backend="TRTLLM")) @needs_gpu @@ -426,7 +426,7 @@ def test_dspark_projector_type_alone_rejects_causal(): cfg = _tiny_config(False) cfg.dflash_config = dict(cfg.dflash_config, projector_type="dspark", causal=True) with pytest.raises(ValueError, match="non-causal DSpark convention"): - Qwen3DSparkForCausalLM(ModelConfig(pretrained_config=cfg, attn_backend="TRTLLM")) + GQADSparkForCausalLM(ModelConfig(pretrained_config=cfg, attn_backend="TRTLLM")) def _run_block_decode(drafter, weights, captured, noise_embed): diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_spec_mode_gate.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_spec_mode_gate.py index 2fe70dbadac3..d98fdbae4572 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_spec_mode_gate.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_spec_mode_gate.py @@ -91,10 +91,6 @@ def test_admitted_modes_pass_the_gate(mode): assert _rejected_by(_model_config(mode)) is None -def test_no_spec_config_passes_the_gate(): - assert _rejected_by(_model_config(None)) is None - - @pytest.mark.parametrize("mode", REFUSED, ids=lambda m: m.name) def test_refused_modes_are_rejected_at_the_spec_gate(mode): assert _rejected_by(_model_config(mode)) == _SPEC_GATE From 57a0eb113849d6214a9b7dbf420549df043d681a Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Fri, 21 Aug 2026 01:30:56 -0700 Subject: [PATCH 17/21] [None][test] Fold the DSpark coverage into the existing test files Replaces four new files (~57 tests, much of it mechanism trivia) with nine tests in the files that already own each subject, reusing their fixtures instead of rebuilding them. Eight of the nine guard silent degradation: the published config and weight spellings activating the heads, head weights with an unresolvable rank, the DFlash refusal, the GQA precondition, form-based worker routing, the worker policies coming off the drafter, and the deployment-form probe. A dropped Markov head lowers acceptance without failing anything, which is exactly what a gsm8k accuracy run cannot see. The ninth, spec-mode index drift, is there for a different reason: this PR adds SPEC_MODE_TO_MODULE as the fourth hand-maintained table in _arch_index, and every other one already has a drift test in test_lazy_model_zoo. Its failure is loud but misattributed -- a missing index entry surfaces as "unsupported speculative decoding mode" to whoever next runs that mode, not as a missing line to whoever omitted it. Dropped as redundant or trivia: the TP-sharded Markov chain (already covered by test_markov_chain_sharded_matches_full_vocab), registry internals, the K3 spec-mode gate (its failure is a loud AssertionError at startup), and the per-field declares_dspark_heads truth table. Signed-off-by: Zhenhuan Chen --- .../hw_agnostic/test_draft_model_registry.py | 226 --------- .../test_dspark_drafter_worker_contract.py | 403 ---------------- .../hw_agnostic/test_dspark_eplb_config.py | 44 ++ .../test_dspark_flavour_dispatch.py | 430 ------------------ .../hw_agnostic/test_dspark_worker.py | 73 ++- .../test_kimi_k3_dspark_semantics.py | 156 ++++++- .../test_kimi_k3_spec_mode_gate.py | 110 ----- tests/unittest/others/test_lazy_model_zoo.py | 39 +- 8 files changed, 276 insertions(+), 1205 deletions(-) delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_draft_model_registry.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_drafter_worker_contract.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_flavour_dispatch.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_spec_mode_gate.py diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_model_registry.py b/tests/unittest/_torch/speculative/hw_agnostic/test_draft_model_registry.py deleted file mode 100644 index e05e82c079f9..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_model_registry.py +++ /dev/null @@ -1,226 +0,0 @@ -# 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. -"""Dispatch contract of the draft-model builder registry. - -``get_draft_model`` used to be one if/elif chain whose *branch order* carried -unwritten rules. Registry dispatch has no inherent order, so what the chain -implied is pinned here explicitly — breaking it silently swaps the draft model -for ``AutoModelForCausalLM``, which no accuracy test would attribute back to -this function. - -The external-draft pre-check runs ahead of the registry without a mode guard, -which is only safe because ``uses_external_draft_model`` implies -``is_mtp_one_model()``. That mutual exclusion is an invariant of ``llm_args``, not -of this module, so it is asserted here rather than assumed. - -Everything here asserts *which builder is selected*, never the object it -builds: constructing a real drafter needs GPUs and checkpoints, and the -selection is the whole contract of this layer. -""" - -from types import SimpleNamespace - -import pytest - -from tensorrt_llm._torch.models import modeling_speculative, modeling_utils -from tensorrt_llm._torch.models._arch_index import SPEC_MODE_TO_MODULE -from tensorrt_llm._torch.models.modeling_utils import ( - _REGISTERED_SPEC_MODES_ATTR, - DRAFT_MODEL_BUILDER_MAPPING, - get_registered_draft_model_builder, - register_draft_model, -) -from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode - -_EXTERNAL_DRAFT_SENTINEL = object() -_BUILDER_SENTINEL = object() - - -class _StubAutoModel: - """Stands in for ``AutoModelForCausalLM`` on the external-draft path.""" - - @staticmethod - def from_config(draft_config): - return _EXTERNAL_DRAFT_SENTINEL - - -def _model_config(mode, *, uses_external_draft_model=False, eagle3_model_arch="llama3"): - """Minimal duck-typed ``ModelConfig`` for dispatch-only assertions.""" - spec_config = SimpleNamespace( - spec_dec_mode=mode, - uses_external_draft_model=uses_external_draft_model, - eagle3_model_arch=eagle3_model_arch, - ) - return SimpleNamespace( - spec_config=spec_config, - pretrained_config=SimpleNamespace(num_hidden_layers=4), - ) - - -def _stub_builder(monkeypatch, mode): - """Replace ``mode``'s registered builder with a sentinel-returning stub.""" - monkeypatch.setitem( - DRAFT_MODEL_BUILDER_MAPPING, mode, lambda *args, **kwargs: _BUILDER_SENTINEL - ) - - -def test_external_draft_model_bypasses_the_registry(monkeypatch): - # An external draft model is loaded from its own checkpoint, so the - # pre-check must short-circuit before the registry is consulted at all. - monkeypatch.setattr(modeling_speculative, "AutoModelForCausalLM", _StubAutoModel) - monkeypatch.setattr( - modeling_speculative, - "get_registered_draft_model_builder", - lambda mode: pytest.fail(f"registry consulted for {mode.name} under external draft"), - ) - - result = modeling_speculative.get_draft_model( - _model_config(SpeculativeDecodingMode.MTP, uses_external_draft_model=True), - draft_config=object(), - lm_head=None, - model=None, - ) - - assert result is _EXTERNAL_DRAFT_SENTINEL - - -def test_eagle3_is_unaffected_by_the_external_draft_flag(monkeypatch): - # `uses_external_draft_model` implies `is_mtp_one_model()`, so it can never - # be true for EAGLE3. This pins the mutual exclusion that lets the pre-check run - # without a mode guard: were the property ever widened, EAGLE3 would start - # building an AutoModel drafter and this test would catch it. - monkeypatch.setattr(modeling_speculative, "AutoModelForCausalLM", _StubAutoModel) - _stub_builder(monkeypatch, SpeculativeDecodingMode.EAGLE3_ONE_MODEL) - - result = modeling_speculative.get_draft_model( - _model_config(SpeculativeDecodingMode.EAGLE3_ONE_MODEL, uses_external_draft_model=True), - draft_config=object(), - lm_head=None, - model=None, - ) - - assert result is _BUILDER_SENTINEL, "external-draft flag hijacked the EAGLE3 builder" - - -def test_external_draft_model_without_draft_config_raises(monkeypatch): - monkeypatch.setattr(modeling_speculative, "AutoModelForCausalLM", _StubAutoModel) - - with pytest.raises(ValueError, match="requires its model config"): - modeling_speculative.get_draft_model( - _model_config(SpeculativeDecodingMode.MTP, uses_external_draft_model=True), - draft_config=None, - lm_head=None, - model=None, - ) - - -def test_unregistered_mode_raises_not_implemented(): - # NGRAM is a drafter-loop mode with no one-engine draft model, so it is - # absent from both the index and the registry. - assert SpeculativeDecodingMode.NGRAM.name not in SPEC_MODE_TO_MODULE - - with pytest.raises(NotImplementedError, match="does not support speculative decoding mode"): - modeling_speculative.get_draft_model( - _model_config(SpeculativeDecodingMode.NGRAM), - draft_config=object(), - lm_head=None, - model=None, - ) - - -def test_every_indexed_mode_resolves_to_a_declaring_builder(): - # Index -> decorator direction: each indexed mode must resolve through the - # single entry point, and the builder must itself declare that mode. The - # declaration is read off the function, never by scanning the mapping by - # identity (a built-in overridden externally keeps the attribute but loses - # its slot). - for mode_name in SPEC_MODE_TO_MODULE: - mode = getattr(SpeculativeDecodingMode, mode_name, None) - assert mode is not None, f"{mode_name} is not a SpeculativeDecodingMode member" - builder = get_registered_draft_model_builder(mode) - assert builder is not None, f"no builder resolved for {mode_name}" - assert mode in getattr(builder, _REGISTERED_SPEC_MODES_ATTR, set()), ( - f"{builder.__module__}.{builder.__qualname__} is registered for " - f"{mode_name} but does not declare it" - ) - - -def test_no_builder_declares_a_mode_missing_from_the_index(): - # Decorator -> index direction: importing every indexed provider and - # walking its builders catches a mode added to an already-indexed module - # without its index entry. (A brand-new provider module is caught by the - # AST scan in tests/unittest/others/test_lazy_model_zoo.py, which needs no - # import and therefore sees modules this loop would never load.) - import importlib - - declared = set() - for module_name in set(SPEC_MODE_TO_MODULE.values()): - module = importlib.import_module(f"tensorrt_llm._torch.models.{module_name}") - for attr in vars(module).values(): - declared |= getattr(attr, _REGISTERED_SPEC_MODES_ATTR, set()) - - missing = {mode.name for mode in declared} - set(SPEC_MODE_TO_MODULE) - assert not missing, f"builders declare modes missing from _arch_index: {missing}" - - -def test_builtin_builder_does_not_override_external_registration(): - # Under lazy loading a built-in module may run its decorators *after* an - # external registration (e.g. --custom_module_dirs), so built-ins only fill - # empty slots. The reverse direction stays last-wins. - mode = SpeculativeDecodingMode.NGRAM - assert mode not in DRAFT_MODEL_BUILDER_MAPPING - - def external(model_config, draft_config, lm_head, model): - return "external" - - def builtin(model_config, draft_config, lm_head, model): - return "builtin" - - builtin.__module__ = "tensorrt_llm._torch.models.modeling_fake" - - try: - register_draft_model(mode)(external) - register_draft_model(mode)(builtin) - assert DRAFT_MODEL_BUILDER_MAPPING[mode] is external, ( - "built-in builder overrode an external registration" - ) - - del DRAFT_MODEL_BUILDER_MAPPING[mode] - register_draft_model(mode)(builtin) - register_draft_model(mode)(external) - assert DRAFT_MODEL_BUILDER_MAPPING[mode] is external - finally: - DRAFT_MODEL_BUILDER_MAPPING.pop(mode, None) - - -def test_stacked_decorators_share_one_builder(): - # Vanilla MTP and MTP_EAGLE_ONE_MODEL are one branch in - # SpeculativeDecodingMode.is_mtp_one_model(); the registry expresses that - # as two keys pointing at the same function. - mtp = get_registered_draft_model_builder(SpeculativeDecodingMode.MTP) - mtp_eagle_one = get_registered_draft_model_builder(SpeculativeDecodingMode.MTP_EAGLE_ONE_MODEL) - - assert mtp is mtp_eagle_one - declared = getattr(mtp, _REGISTERED_SPEC_MODES_ATTR, set()) - assert {SpeculativeDecodingMode.MTP, SpeculativeDecodingMode.MTP_EAGLE_ONE_MODEL} <= declared - - -def test_registry_module_placement_matches_index(): - # Builders live next to the draft model they construct, never in the - # factory file: that is what keeps get_draft_model free of concrete - # imports (and what removed the DSpark lazy import). - builder = get_registered_draft_model_builder(SpeculativeDecodingMode.DSPARK) - assert builder.__module__ == "tensorrt_llm._torch.models.modeling_dspark" - assert modeling_utils.DRAFT_MODEL_BUILDER_MAPPING is DRAFT_MODEL_BUILDER_MAPPING diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_drafter_worker_contract.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_drafter_worker_contract.py deleted file mode 100644 index d2250602d7c1..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_drafter_worker_contract.py +++ /dev/null @@ -1,403 +0,0 @@ -# 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. -"""Contract tests: a real standalone DSpark drafter driven by the real worker -that the routing selects for it. - -``test_dspark_flavour_dispatch.py`` asserts *which class* each factory returns, -using stubs. That cannot catch a worker handed a draft model whose attributes it -does not have: the two agree on paper and diverge on first contact. The routing -bug this file pins surfaced only as - - AttributeError: 'GQADSparkForCausalLM' object has no attribute 'num_stages' - -five minutes into a 16-GPU run, after the weights had loaded -- because nothing -below the factory was ever exercised. - -So these tests build the drafter for real and drive ``DFlashWorker``'s lazy -init, which is where the drafter contract actually lives: it reaches for -``fc.weight``, ``block_size``, ``_build_fused_kv_buffers``, ``_num_attn_layers``, -``_num_heads``, ``_num_kv_heads``, ``_head_dim`` and ``_get_attention_mask_args``. -A worker routed by mode instead of by flavour fails here, in seconds, on one GPU. - -The drafter is built with the TRTLLM block-decode backend, matching the K3 -serving config; the VANILLA default would pull in flash-attn. -""" - -from types import SimpleNamespace - -import pytest -import torch - -from tensorrt_llm._torch.speculative import utils as spec_utils -from tensorrt_llm._torch.speculative.dflash import DFlashWorker -from tensorrt_llm._torch.speculative.dspark import DSparkWorker, DSv4DSparkWorker -from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode -from tensorrt_llm.mapping import Mapping - -needs_gpu = pytest.mark.skipif( - not torch.cuda.is_available(), reason="the drafter and the worker buffers are CUDA-resident" -) - -VOCAB = 256 -RANK = 8 -BLOCK_SIZE = 4 -MAX_REQUESTS = 4 -MAX_SEQ_LEN = 128 - -TINY = dict( - architectures=["DSparkDraftModel"], - model_type="qwen3", - block_size=BLOCK_SIZE, - hidden_size=64, - num_hidden_layers=2, - num_attention_heads=4, - num_key_value_heads=2, - # The fused QK-norm-RoPE kernel the bf16 block decode uses rejects small - # head dims; 128 is also the real K3 drafter's head_dim. - head_dim=128, - intermediate_size=128, - hidden_act="silu", - rms_norm_eps=1e-6, - vocab_size=VOCAB, - max_position_embeddings=2048, - rope_theta=10000.0, - rope_scaling=None, - attention_bias=False, - torch_dtype="bfloat16", - num_target_layers=2, - tie_word_embeddings=False, -) - -NUM_CAPTURE = 2 - - -def _drafter_config(): - """A standalone DSpark drafter config: qwen3 backbone plus the head set.""" - from transformers import Qwen3Config - - cfg = dict(TINY) - cfg["dflash_config"] = { - "mask_token_id": VOCAB - 2, - "target_layer_ids": [0, 1], - "projector_type": "dspark", - "causal": False, - "shift_label": True, - "markov_rank": RANK, - "markov_head_type": "vanilla", - "use_confidence_head": True, - } - return Qwen3Config.from_dict(cfg) - - -def _drafter_weights(seed=11, legacy_head_keys=False): - """Synthetic drafter weights in the published key namespace. - - Head tensors are named the way both public checkpoints ship them - (``markov_head.*`` / ``confidence_head.proj.*``, after the submodules that - own them); ``legacy_head_keys`` switches to the bare spellings so the - aliasing is covered from both sides. - """ - g = torch.Generator().manual_seed(seed) - - def rnd(*shape): - return (torch.randn(*shape, generator=g) * 0.05).to(torch.bfloat16) - - h, inter = TINY["hidden_size"], TINY["intermediate_size"] - nh, nkv, hd = (TINY["num_attention_heads"], TINY["num_key_value_heads"], TINY["head_dim"]) - head = { - "markov_w1.weight": rnd(VOCAB, RANK), - "markov_w2.weight": rnd(VOCAB, RANK), - "confidence_proj.weight": rnd(1, h + RANK), - "confidence_proj.bias": rnd(1), - } - if not legacy_head_keys: - head = { - "markov_head.markov_w1.weight": head["markov_w1.weight"], - "markov_head.markov_w2.weight": head["markov_w2.weight"], - "confidence_head.proj.weight": head["confidence_proj.weight"], - "confidence_head.proj.bias": head["confidence_proj.bias"], - } - weights = { - "fc.weight": rnd(h, h * NUM_CAPTURE), - "hidden_norm.weight": rnd(h) + 1.0, - "norm.weight": rnd(h) + 1.0, - **head, - } - for i in range(TINY["num_hidden_layers"]): - p = f"layers.{i}." - weights[p + "self_attn.q_proj.weight"] = rnd(nh * hd, h) - weights[p + "self_attn.k_proj.weight"] = rnd(nkv * hd, h) - weights[p + "self_attn.v_proj.weight"] = rnd(nkv * hd, h) - weights[p + "self_attn.o_proj.weight"] = rnd(h, nh * hd) - weights[p + "self_attn.q_norm.weight"] = rnd(hd) + 1.0 - weights[p + "self_attn.k_norm.weight"] = rnd(hd) + 1.0 - weights[p + "input_layernorm.weight"] = rnd(h) + 1.0 - weights[p + "post_attention_layernorm.weight"] = rnd(h) + 1.0 - weights[p + "mlp.gate_proj.weight"] = rnd(inter, h) - weights[p + "mlp.up_proj.weight"] = rnd(inter, h) - weights[p + "mlp.down_proj.weight"] = rnd(h, inter) - return weights - - -def _drafter_config_top_level_spelling(): - """RadixArk/Kimi-K3-DSpark as published: head switches at the top level. - - ``dflash_config`` carries only mask_token_id / target_layer_ids and the - confidence flag is spelled ``enable_confidence_head``. - """ - from transformers import Qwen3Config - - cfg = dict(TINY) - # Verbatim shape of the published config: no shift_label, no - # projector_type, and dflash_config holding nothing but these two keys. - cfg["dflash_config"] = {"mask_token_id": VOCAB - 2, "target_layer_ids": [0, 1]} - cfg["markov_rank"] = RANK - cfg["markov_head_type"] = "vanilla" - cfg["enable_confidence_head"] = True - cfg["confidence_head_with_markov"] = True - return Qwen3Config.from_dict(cfg) - - -@pytest.fixture(scope="module") -def standalone_drafter(): - """The real ``GQADSparkForCausalLM``, weights loaded, on the device.""" - from tensorrt_llm._torch.model_config import ModelConfig - from tensorrt_llm._torch.models.modeling_dspark import GQADSparkForCausalLM - - model_config = ModelConfig(pretrained_config=_drafter_config(), attn_backend="TRTLLM") - drafter = GQADSparkForCausalLM(model_config, dflash_attention_backend="TRTLLM").to("cuda") - drafter.load_weights(_drafter_weights()) - return drafter - - -@needs_gpu -def test_top_level_head_spelling_activates_the_heads(): - # The published RadixArk revision spells the switches top-level. Reading - # only dflash_config resolves markov_rank to 0, which drops markov_w1/w2 - # on the floor and falls back to the DFlash slot convention -- correct - # output, lower acceptance, nothing raised. - from tensorrt_llm._torch.model_config import ModelConfig - from tensorrt_llm._torch.models.modeling_dspark import GQADSparkForCausalLM - - model_config = ModelConfig( - pretrained_config=_drafter_config_top_level_spelling(), attn_backend="TRTLLM" - ) - drafter = GQADSparkForCausalLM(model_config, dflash_attention_backend="TRTLLM").to("cuda") - drafter.load_weights(_drafter_weights()) - - assert drafter._dspark_markov_rank == RANK - assert drafter._dspark_use_confidence_head is True - assert drafter.has_markov_head, "markov weights were dropped despite being in the checkpoint" - # Declared nowhere in the published config, so it rides on the DSpark - # default. False here would run slots 1..7 on a block_size-7 drafter and - # read the next request's anchor slot. - assert drafter._dspark_shift_label is True - - -@needs_gpu -def test_published_head_weight_names_load(): - # Both public drafters nest the head tensors under the module that owns - # them. Matching only the bare names leaves markov_w1/w2 in the dict handed - # to DFlash, which drops them -- or, once the rank resolves, raises - # "missing markov_w1.weight" on a checkpoint that plainly ships it. - from tensorrt_llm._torch.model_config import ModelConfig - from tensorrt_llm._torch.models.modeling_dspark import GQADSparkForCausalLM - - for legacy in (False, True): - model_config = ModelConfig( - pretrained_config=_drafter_config_top_level_spelling(), attn_backend="TRTLLM" - ) - drafter = GQADSparkForCausalLM(model_config, dflash_attention_backend="TRTLLM").to("cuda") - drafter.load_weights(_drafter_weights(legacy_head_keys=legacy)) - assert drafter.has_markov_head, f"legacy_head_keys={legacy}" - assert drafter.confidence_proj_weight is not None, f"legacy_head_keys={legacy}" - - -@needs_gpu -def test_markov_weights_without_a_resolvable_rank_raise(): - # The inverse of the missing-weights check: weights present, rank resolved - # to 0. Loading silently would cost acceptance with no signal. - from transformers import Qwen3Config - - from tensorrt_llm._torch.model_config import ModelConfig - from tensorrt_llm._torch.models.modeling_dspark import GQADSparkForCausalLM - - cfg = dict(TINY) - cfg["dflash_config"] = {"mask_token_id": VOCAB - 2, "target_layer_ids": [0, 1]} - model_config = ModelConfig(pretrained_config=Qwen3Config.from_dict(cfg), attn_backend="TRTLLM") - drafter = GQADSparkForCausalLM(model_config, dflash_attention_backend="TRTLLM").to("cuda") - - with pytest.raises(ValueError, match="markov_rank resolved to 0"): - drafter.load_weights(_drafter_weights()) - - -def _spec_config(*, embedded): - """Only the fields the routing and the worker read.""" - return SimpleNamespace( - spec_dec_mode=SpeculativeDecodingMode.DSPARK, - draft_is_embedded_in_target=embedded, - _use_shared_kv_cache=False, - _allow_separate_draft_kv_cache=True, - # K == block_size under the dspark shift_label convention, which is - # also what DSv4DSparkWorker validates before it touches anything else. - max_draft_len=BLOCK_SIZE, - attention_backend="TRTLLM", - ) - - -def _lazy_init_args(): - spec_metadata = SimpleNamespace(max_num_requests=MAX_REQUESTS) - attn_metadata = SimpleNamespace(max_seq_len=MAX_SEQ_LEN) - return spec_metadata, attn_metadata - - -@needs_gpu -def test_routed_worker_initializes_against_a_real_standalone_drafter(standalone_drafter): - """The end-to-end contract, and the regression test for the routing bug. - - ``get_spec_worker`` picks the worker; the drafter is the real one the - builder would produce for the same config. Driving lazy init proves the two - agree on the draft-model interface. Route by mode instead of by flavour and - this raises ``AttributeError: ... has no attribute 'num_stages'``. - """ - spec_config = _spec_config(embedded=False) - worker = spec_utils.get_spec_worker( - spec_config, model_config=None, mapping=Mapping(), use_separate_draft_kv_cache=True - ) - # DFlash lineage (all the drafting plumbing), DSpark leaf (the two heads). - assert isinstance(worker, DSparkWorker) - assert isinstance(worker, DFlashWorker) - - worker.set_draft_model(standalone_drafter) - worker._lazy_init_ctx_buffers(standalone_drafter, *_lazy_init_args()) - - assert worker._ctx_buf_inited - # One scratch slot on top of the request slots, so dummy/padded writes - # cannot land on a real request's context. - assert worker._ctx_len.shape == (MAX_REQUESTS + 1,) - assert worker._dummy_slot == MAX_REQUESTS - assert worker._batch_to_slot.shape == (MAX_REQUESTS,) - assert worker._resolved_block_size == BLOCK_SIZE - assert sorted(worker._free_slots) == list(range(MAX_REQUESTS)) - - -@needs_gpu -def test_routed_worker_sees_the_dspark_heads(standalone_drafter): - """The heads moved to the DSpark subclass must stay visible to the worker. - - ``DSparkWorker`` probes them defensively (``getattr(..., False)``), - so a drafter that lost them degrades to plain DFlash silently -- lower - acceptance, no error. These are the two probes the block-draft step makes. - """ - assert getattr(standalone_drafter, "has_markov_head", False) is True - assert getattr(standalone_drafter, "_dspark_shift_label", False) is True - assert standalone_drafter.markov_w1.shape == (VOCAB, RANK) - assert standalone_drafter.markov_w2.shape == (VOCAB, RANK) - - -@needs_gpu -def test_embedded_worker_cannot_drive_a_standalone_drafter(standalone_drafter): - """Witness for why the routing has to follow the flavour. - - ``DSv4DSparkWorker`` serves the embedded DeepSeek-V4-Pro draft and reads - V4-draft-only attributes. Handing it a standalone drafter is the exact - mis-route that reached production, so pin the failure rather than trusting - the routing test alone to stay correct. - """ - worker = DSv4DSparkWorker(_spec_config(embedded=True), Mapping()) - spec_metadata, _ = _lazy_init_args() - with pytest.raises(AttributeError, match="num_stages"): - worker._lazy_init(standalone_drafter, spec_metadata) - - -@needs_gpu -def test_dspark_overrides_read_the_drafter_not_a_hardcoded_policy(standalone_drafter): - """The two DSpark policies must come off the drafter, not the class. - - ``DSparkWorker`` exists only to override the block-output slot - convention and the Markov bias. Both are declared by the checkpoint, so a - DSpark drafter trained with the legacy DFlash slot layout (or without a - Markov head) must still get the base-class behaviour. Hardcoding the - dspark answer in the subclass would pass every mode/flavour dispatch test - while silently mis-slotting such a drafter. - """ - worker = DSparkWorker(_spec_config(embedded=False), Mapping()) - worker.set_draft_model(standalone_drafter) - - # shift_label on -> slots 0..K-1 (anchor slot predicts the first draft - # token); the DFlash base would return slots 1..K for the same inputs. - ids = worker._draft_slot_ids( - standalone_drafter, num_gens=2, block_size=BLOCK_SIZE, num_draft_tokens=3 - ) - assert ids.tolist() == [0, 1, 2, BLOCK_SIZE, BLOCK_SIZE + 1, BLOCK_SIZE + 2] - base_ids = DFlashWorker._draft_slot_ids( - worker, standalone_drafter, num_gens=2, block_size=BLOCK_SIZE, num_draft_tokens=3 - ) - assert base_ids.tolist() == [1, 2, 3, BLOCK_SIZE + 1, BLOCK_SIZE + 2, BLOCK_SIZE + 3] - - # A drafter without a Markov head falls through untouched, even on the - # DSpark subclass. - logits = torch.randn(2, 3, 8, device="cuda") - object.__setattr__(standalone_drafter, "_dspark_markov_rank", 0) - standalone_drafter.markov_w1 = None - assert standalone_drafter.has_markov_head is False - out = worker._refine_block_logits(standalone_drafter, logits, {}, None) - assert out is logits - - -@needs_gpu -def test_tp_sharded_markov_chain_matches_the_unsharded_result(): - # The sharded branch is unreachable on the validated TEP16/attention-DP - # config -- there the draft head is full-vocab, so vocab_slice is None and - # no gather runs. Exercise it here with a fake two-rank split: each rank - # sees its own logit columns and a markov_w2 sliced the same way, and the - # chain must agree column-for-column with the unsharded computation. - from tensorrt_llm._torch.models.modeling_speculative import dspark_markov_chain_logits - - torch.manual_seed(0) - vocab, rank, batch, block, tp = 32, 4, 3, 5, 2 - shard = vocab // tp - w1 = torch.randn(vocab, rank, device="cuda") - w2 = torch.randn(vocab, rank, device="cuda") - base = torch.randn(batch, block, vocab, device="cuda") - anchor = torch.randint(0, vocab, (batch,), device="cuda") - - full = dspark_markov_chain_logits(base, anchor, w1, w2) - - for tp_rank in range(tp): - vocab_slice = slice(tp_rank * shard, (tp_rank + 1) * shard) - # Stands in for greedy_sample_draft_with_tp_gather: the real one - # all-gathers local (index, value) pairs, so the chain always advances - # on the global argmax rather than this rank's local one. - sharded = dspark_markov_chain_logits( - base[..., vocab_slice], - anchor, - w1, - w2[vocab_slice], - argmax_fn=_make_global_argmax(full, tp_rank, shard), - ) - torch.testing.assert_close(sharded, full[..., vocab_slice]) - - -def _make_global_argmax(full_logits, tp_rank, shard): - """Replay the global argmax the TP gather would have produced.""" - step = {"i": 0} - - def argmax_fn(_step_logits): - i = step["i"] - step["i"] += 1 - return full_logits[:, i].argmax(dim=-1) - - return argmax_fn 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 b52ba3a552e1..4a342e6ee8ef 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 @@ -223,3 +223,47 @@ def test_layer_base_not_checked_without_eplb(): validate_dspark_eplb_layer_base( _model_config(None), _model_config(None, num_hidden_layers=3) ) + + +# --------------------------------------------------------------------------- +# Deployment-form probe. Every DSpark dispatch -- draft model, worker, spec +# metadata, draft-KV decision -- reads this one flag, so a misread does not +# degrade gracefully: it routes a standalone drafter into the V4 worker. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "index_name,weight_map,model_type,expected", + [ + # A parsed index is authoritative in BOTH directions. Falling through + # to model_type on a standalone V4-shaped drafter would classify it as + # embedded, and that only surfaces inside count_dspark_stages. + ("model.safetensors.index.json", {"mtp.0.layers.0.weight": "a"}, "deepseek_v4", True), + ( + "model.safetensors.index.json", + {"layers.0.self_attn.q_proj.weight": "a"}, + "deepseek_v4", + False, + ), + # The bin index is probed too; only the safetensors one used to be. + ("pytorch_model.bin.index.json", {"mtp.1.mlp.weight": "a"}, "qwen3", True), + # No index at all -> the model_type fallback. + (None, None, "deepseek_v4", True), + (None, None, "qwen3", False), + ], + ids=["mtp_index", "standalone_index", "bin_index", "no_index_v4", "no_index_qwen3"], +) +def test_draft_form_probe_reads_the_checkpoint( + tmp_path, index_name, weight_map, model_type, expected +): + import json + + from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig + + (tmp_path / "config.json").write_text(json.dumps({"model_type": model_type})) + if index_name is not None: + (tmp_path / index_name).write_text(json.dumps({"weight_map": weight_map})) + + cfg = DSparkDecodingConfig(max_draft_len=7, speculative_model=str(tmp_path)) + + assert cfg.draft_is_embedded_in_target is expected diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_flavour_dispatch.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_flavour_dispatch.py deleted file mode 100644 index b3ee48d4eb10..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_flavour_dispatch.py +++ /dev/null @@ -1,430 +0,0 @@ -# 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. -"""Which DSpark drafter ``decoding_type: DSpark`` builds, and what DFlash refuses. - -``DSpark`` ships in two flavours -- embedded in the target checkpoint -(DeepSeek-V4-Pro's ``mtp.*``) or standalone with its own checkpoint -- and one -builder picks between them on deployment form alone. There is no second -dispatch on ``model_type``: the standalone drafter's backbone is resolved one -layer down by the model registry, and the block decode's GQA precondition is -what rejects a backbone it cannot express. - -The DFlash half matters just as much: DFlash no longer implements the DSpark -head set, so a drafter that declares it must be refused rather than served -without it. Silently dropping the Markov head does not fail anything; it lowers -the acceptance rate, which no test would attribute back to this function. - -Selection is all this file checks -- which class each factory returns, with the -classes stubbed. That is worth pinning, but it cannot catch a worker and a -drafter that agree on paper and diverge on first contact; for that see -``test_dspark_drafter_worker_contract.py``, which builds the real drafter and -drives the real worker's lazy init. The flavour probe at the bottom is the one -exception here -- it reads checkpoints written to ``tmp_path``, because it is -the single source every dispatch above consults. -""" - -import json -from types import SimpleNamespace - -import pytest - -from tensorrt_llm._torch.models import modeling_dflash, modeling_dspark -from tensorrt_llm._torch.models.modeling_dflash import declares_dspark_heads -from tensorrt_llm._torch.speculative import utils as spec_utils -from tensorrt_llm._torch.speculative.interface import ( - SpeculativeDecodingMode, - should_use_separate_draft_kv_cache, -) -from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig - -_DSV4_SENTINEL = object() -_GQA_SENTINEL = object() -_DFLASH_SENTINEL = object() -_LAGUNA_SENTINEL = object() - -_DSPARK_HEADS = { - "markov_rank": 256, - "markov_head_type": "vanilla", - "use_confidence_head": True, - "shift_label": True, - "projector_type": "dspark", -} - - -def _configs( - *, - model_type="qwen3", - dflash_config=None, - architectures=None, - attention_backend="TRTLLM", - embedded=False, - top_level=None, -): - """Duck-typed (target ModelConfig, draft ModelConfig) for dispatch-only asserts.""" - model_config = SimpleNamespace( - spec_config=SimpleNamespace( - speculative_model="/nonexistent/drafter", - block_size=7, - attention_backend=attention_backend, - draft_is_embedded_in_target=embedded, - ) - ) - draft_config = SimpleNamespace( - pretrained_config=SimpleNamespace( - model_type=model_type, - architectures=architectures, - dflash_config=dflash_config, - **(top_level or {}), - ) - ) - return model_config, draft_config - - -@pytest.fixture -def stub_dspark(monkeypatch): - """Replace the DSpark drafter classes with sentinel-returning stubs.""" - monkeypatch.setattr(modeling_dspark, "DSv4DSparkForCausalLM", lambda *a, **k: _DSV4_SENTINEL) - monkeypatch.setattr(modeling_dspark, "GQADSparkForCausalLM", lambda *a, **k: _GQA_SENTINEL) - monkeypatch.setattr(modeling_dspark, "validate_dspark_eplb_layer_base", lambda *a, **k: None) - - -@pytest.fixture -def stub_dflash(monkeypatch): - monkeypatch.setattr(modeling_dflash, "DFlashForCausalLM", lambda *a, **k: _DFLASH_SENTINEL) - monkeypatch.setattr( - modeling_dflash, "DFlashLagunaForCausalLM", lambda *a, **k: _LAGUNA_SENTINEL - ) - - -# -------------------------------------------------------------------------- -# decoding_type: DSpark -# -------------------------------------------------------------------------- - - -@pytest.mark.parametrize("model_type", ["qwen3", "llama", "gpt_oss"]) -def test_standalone_drafter_selects_gqa_dspark(monkeypatch, stub_dspark, model_type): - # No per-model_type table: every GQA backbone the registry can build gets - # the same drafter class. A new GQA family must not need an entry here. - model_config, draft_config = _configs(model_type=model_type, embedded=False) - - built = modeling_dspark._build_dspark_draft(model_config, draft_config, None, None) - - assert built is _GQA_SENTINEL - - -def test_embedded_draft_selects_dsv4_dspark(monkeypatch, stub_dspark): - monkeypatch.setattr(modeling_dspark, "count_dspark_stages", lambda _p: 3) - model_config, draft_config = _configs(model_type="deepseek_v4", embedded=True) - - built = modeling_dspark._build_dspark_draft(model_config, draft_config, None, None) - - assert built is _DSV4_SENTINEL - - -def test_non_gqa_backbone_is_rejected_at_construction(): - # The builder no longer whitelists model_type; the block decode's GQA - # precondition is what rejects an MLA-shaped drafter, and it must do so at - # construction rather than deep inside _build_fused_kv_buffers. - from types import SimpleNamespace - - drafter = modeling_dflash.DFlashForCausalLM.__new__(modeling_dflash.DFlashForCausalLM) - drafter.model = SimpleNamespace( - layers=[SimpleNamespace(self_attn=SimpleNamespace())] # no qkv_proj -> MLA-like - ) - drafter.config = SimpleNamespace() - - with pytest.raises(ValueError) as excinfo: - drafter._validate_gqa_shape() - - assert "qkv_proj" in str(excinfo.value) - - -def test_layers_with_mismatched_kv_heads_are_rejected(): - from types import SimpleNamespace - - def _layer(nkv): - return SimpleNamespace( - self_attn=SimpleNamespace(qkv_proj=object(), num_key_value_heads=nkv) - ) - - drafter = modeling_dflash.DFlashForCausalLM.__new__(modeling_dflash.DFlashForCausalLM) - drafter.model = SimpleNamespace(layers=[_layer(8), _layer(8), _layer(4)]) - drafter.config = SimpleNamespace() - - with pytest.raises(ValueError) as excinfo: - drafter._validate_gqa_shape() - - assert "[2]" in str(excinfo.value) - - -def test_standalone_drafter_receives_the_attention_backend(monkeypatch, stub_dspark): - seen = {} - - def _capture(draft_config, *, dflash_attention_backend): - seen["backend"] = dflash_attention_backend - return _GQA_SENTINEL - - monkeypatch.setattr(modeling_dspark, "GQADSparkForCausalLM", _capture) - model_config, draft_config = _configs(attention_backend="TRTLLM") - - modeling_dspark._build_dspark_draft(model_config, draft_config, None, None) - - assert seen["backend"] == "TRTLLM" - - -# -------------------------------------------------------------------------- -# decoding_type: DFlash -# -------------------------------------------------------------------------- - - -def test_dflash_refuses_a_dspark_drafter(stub_dflash): - model_config, draft_config = _configs(dflash_config=dict(_DSPARK_HEADS)) - - with pytest.raises(ValueError) as excinfo: - modeling_dflash._build_dflash_draft(model_config, draft_config, None, None) - - message = str(excinfo.value) - assert "DSpark" in message - assert "decoding_type" in message, "the error must say how to fix the config" - - -def test_dflash_refuses_a_top_level_spelling_dspark_drafter(stub_dflash): - # RadixArk/Kimi-K3-DSpark as published: the head switches sit at the top - # level and dflash_config carries only mask_token_id / target_layer_ids. - # A reader that looks in dflash_config alone misses exactly the drafter - # this guard exists to catch. - model_config, draft_config = _configs( - dflash_config={"mask_token_id": 163824, "target_layer_ids": [7, 23, 51, 67, 83]}, - top_level={ - "markov_rank": 256, - "markov_head_type": "vanilla", - "enable_confidence_head": True, - "block_size": 7, - }, - ) - - with pytest.raises(ValueError, match="DSpark"): - modeling_dflash._build_dflash_draft(model_config, draft_config, None, None) - - -def test_plain_dflash_drafter_is_unchanged(stub_dflash): - model_config, draft_config = _configs( - dflash_config={"mask_token_id": 7, "target_layer_ids": [0, 1]} - ) - - built = modeling_dflash._build_dflash_draft(model_config, draft_config, None, None) - - assert built is _DFLASH_SENTINEL - - -def test_laguna_drafter_is_unchanged(stub_dflash): - model_config, draft_config = _configs( - architectures=["DFlashLagunaForCausalLM"], - dflash_config={"mask_token_id": 7}, - ) - - built = modeling_dflash._build_dflash_draft(model_config, draft_config, None, None) - - assert built is _LAGUNA_SENTINEL - - -def test_legacy_causal_dflash_config_is_not_mistaken_for_dspark(stub_dflash): - # Laguna configs carry ``causal`` without any DSpark field; the legacy - # decode path handles it, so this must not trip the refusal. - model_config, draft_config = _configs(dflash_config={"mask_token_id": 7, "causal": True}) - - built = modeling_dflash._build_dflash_draft(model_config, draft_config, None, None) - - assert built is _DFLASH_SENTINEL - - -# -------------------------------------------------------------------------- -# The predicate itself -# -------------------------------------------------------------------------- - - -@pytest.mark.parametrize( - "dflash_config,expected", - [ - (None, False), - ({}, False), - ({"mask_token_id": 7}, False), - ({"causal": True}, False), - ({"markov_rank": 0}, False), - ({"shift_label": False}, False), - ({"markov_rank": 256}, True), - ({"shift_label": True}, True), - ({"use_confidence_head": True}, True), - ({"projector_type": "dspark"}, True), - ({"projector_type": "DSpark"}, True), - ], -) -def test_declares_dspark_heads(dflash_config, expected): - config = SimpleNamespace(dflash_config=dflash_config) - assert declares_dspark_heads(config) is expected - - -# -------------------------------------------------------------------------- -# Runtime routing -# -# The worker, the spec metadata and the separate-draft-KV-cache decision must -# follow the same flavour flag the builder follows. When they disagree, -# DSv4DSparkWorker gets a standalone drafter and dies reaching for V4-draft-only -# attributes (num_stages, write_context_windows) -- only at the first forward, -# long after the engine reported a successful build. -# -------------------------------------------------------------------------- - -_WORKER_SENTINELS = { - "DFlashWorker": object(), - "DSparkWorker": object(), - "DSv4DSparkWorker": object(), -} -# SimpleNamespace rather than object(): get_spec_metadata assigns -# ``metadata.enable_penalty`` on whatever it built, which a bare object rejects. -_METADATA_SENTINELS = { - "DFlashSpecMetadata": SimpleNamespace(), - "DSparkSpecMetadata": SimpleNamespace(), -} - - -@pytest.fixture -def stub_runtime(monkeypatch): - """Stub the worker/metadata classes: constructing the real ones needs CUDA.""" - for name, sentinel in _WORKER_SENTINELS.items(): - monkeypatch.setattr(spec_utils, name, lambda *a, _s=sentinel, **k: _s) - for name, sentinel in _METADATA_SENTINELS.items(): - monkeypatch.setattr(spec_utils, name, lambda *a, _s=sentinel, **k: _s) - - -def _spec_config(mode, *, embedded, allow_separate_kv=True): - """Duck-typed spec config carrying only what the routing functions read.""" - return SimpleNamespace( - spec_dec_mode=mode, - draft_is_embedded_in_target=embedded, - _use_shared_kv_cache=False, - _allow_separate_draft_kv_cache=allow_separate_kv, - max_draft_len=7, - max_total_draft_tokens=7, - tokens_per_gen_step=8, - target_layer_ids=[7, 23, 51, 67, 83], - advanced_sampling_mode=None, - # Read by get_spec_metadata for the occurrence-penalty workspace; the - # routing under test does not depend on it. - enable_penalty=False, - ) - - -@pytest.mark.parametrize( - "mode,embedded,expected", - [ - (SpeculativeDecodingMode.DSPARK, True, "DSv4DSparkWorker"), - (SpeculativeDecodingMode.DSPARK, False, "DSparkWorker"), - (SpeculativeDecodingMode.DFLASH, False, "DFlashWorker"), - ], -) -def test_worker_follows_the_flavour_not_the_mode(stub_runtime, mode, embedded, expected): - worker = spec_utils.get_spec_worker( - _spec_config(mode, embedded=embedded), - model_config=None, - mapping=None, - use_separate_draft_kv_cache=False, - ) - assert worker is _WORKER_SENTINELS[expected] - - -@pytest.mark.parametrize( - "mode,embedded,expected", - [ - (SpeculativeDecodingMode.DSPARK, True, "DSparkSpecMetadata"), - (SpeculativeDecodingMode.DSPARK, False, "DFlashSpecMetadata"), - (SpeculativeDecodingMode.DFLASH, False, "DFlashSpecMetadata"), - ], -) -def test_spec_metadata_follows_the_flavour_not_the_mode(stub_runtime, mode, embedded, expected): - metadata = spec_utils.get_spec_metadata( - _spec_config(mode, embedded=embedded), - SimpleNamespace(hidden_size=7168, torch_dtype=None, vocab_size=163840), - max_num_requests=8, - max_num_tokens=4096, - ) - assert metadata is _METADATA_SENTINELS[expected] - - -@pytest.mark.parametrize( - "mode,embedded,expected", - [ - # The embedded draft opts out: it owns a rolling captured-context window. - (SpeculativeDecodingMode.DSPARK, True, False), - # A standalone DSpark drafter runs on DFlashWorker's paged draft KV -- - # the path K3 used before its decoding_type moved to DSpark. - (SpeculativeDecodingMode.DSPARK, False, True), - (SpeculativeDecodingMode.DFLASH, False, True), - ], -) -def test_separate_draft_kv_cache_follows_the_flavour(mode, embedded, expected): - config = _spec_config(mode, embedded=embedded) - assert should_use_separate_draft_kv_cache(config) is expected - - -# -------------------------------------------------------------------------- -# The flavour probe itself. This is the single source every dispatch above -# reads, so it is the one place the embedded/standalone question is decided. -# -------------------------------------------------------------------------- - - -def _write_ckpt(tmp_path, *, weight_map=None, model_type=None): - if weight_map is not None: - (tmp_path / "model.safetensors.index.json").write_text( - json.dumps({"weight_map": weight_map}) - ) - if model_type is not None: - (tmp_path / "config.json").write_text(json.dumps({"model_type": model_type})) - return DSparkDecodingConfig(max_draft_len=7, speculative_model=str(tmp_path)) - - -def test_probe_reads_the_mtp_namespace_from_the_weight_index(tmp_path): - config = _write_ckpt( - tmp_path, - weight_map={ - "mtp.0.attn.wq_a.weight": "x.safetensors", - "layers.0.q.weight": "x.safetensors", - }, - model_type="deepseek_v4", - ) - assert config.draft_is_embedded_in_target is True - - -def test_probe_treats_a_standalone_drafter_index_as_standalone(tmp_path): - config = _write_ckpt( - tmp_path, - weight_map={"layers.0.self_attn.q_proj.weight": "x.safetensors"}, - model_type="qwen3", - ) - assert config.draft_is_embedded_in_target is False - - -def test_probe_falls_back_to_model_type_without_an_index(tmp_path): - # A V4 checkpoint whose index file is absent must not be read as - # standalone: the standalone lineage has no V4 drafter. - config = _write_ckpt(tmp_path, model_type="deepseek_v4") - assert config.draft_is_embedded_in_target is True - - -def test_probe_is_standalone_when_nothing_can_be_read(tmp_path): - # Fail soft: an unreadable or not-yet-downloaded checkpoint must not crash - # config validation, and standalone is the safe default (it is the flavour - # whose worker probes the DSpark heads defensively). - config = DSparkDecodingConfig(max_draft_len=7, speculative_model=str(tmp_path / "missing")) - assert config.draft_is_embedded_in_target is False 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 bd917658f9b3..d5d5ec2a0bf2 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -26,7 +26,11 @@ import pytest import torch -from tensorrt_llm._torch.speculative.dspark import DSparkSpecMetadata, DSv4DSparkWorker +from tensorrt_llm._torch.speculative.dspark import ( + DSparkSpecMetadata, + DSparkWorker, + DSv4DSparkWorker, +) from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode pytestmark = pytest.mark.skipif( @@ -735,3 +739,70 @@ def fake_sample_draft_tokens(gl, sm, bs, *, draft_step): gen_draft = nd[num_contexts:] expected_gen = torch.stack([s[num_contexts:] for s in sampled_per_step], dim=1) assert torch.equal(gen_draft, expected_gen) + + +# --------------------------------------------------------------------------- +# Routing: decoding_type DSpark serves two deployment forms, and the worker, +# the spec metadata and the draft-KV decision must all follow the same flag. +# Mis-routing is not hypothetical: handing a standalone drafter to +# DSv4DSparkWorker raises AttributeError on num_stages at lazy init. +# --------------------------------------------------------------------------- + + +def _routing_config(embedded): + return types.SimpleNamespace( + max_draft_len=5, + max_total_draft_tokens=5, + spec_dec_mode=SpeculativeDecodingMode.DSPARK, + draft_is_embedded_in_target=embedded, + attention_backend="TRTLLM", + confidence_threshold=0.5, + _use_shared_kv_cache=False, + _allow_separate_draft_kv_cache=True, + ) + + +@pytest.mark.parametrize( + "embedded,worker_cls,uses_separate_draft_kv", + [(True, DSv4DSparkWorker, False), (False, DSparkWorker, True)], + ids=["embedded", "standalone"], +) +def test_worker_and_draft_kv_follow_the_draft_form(embedded, worker_cls, uses_separate_draft_kv): + from tensorrt_llm._torch.speculative.interface import should_use_separate_draft_kv_cache + from tensorrt_llm._torch.speculative.utils import get_spec_worker + from tensorrt_llm.mapping import Mapping + + spec_config = _routing_config(embedded) + + worker = get_spec_worker(spec_config, None, Mapping()) + assert type(worker) is worker_cls + + # The embedded draft owns a rolling window and never reads the paged draft + # KV cache; the standalone one is DFlash lineage and does. + assert should_use_separate_draft_kv_cache(spec_config) is uses_separate_draft_kv + + +def test_dspark_worker_policies_come_from_the_drafter(): + """The two DSpark overrides must read the checkpoint, not hardcode dspark. + + A DSpark drafter trained with the legacy DFlash slot layout, or shipped + without a Markov head, has to get the base-class behaviour. Hardcoding the + dspark answer here would pass every routing test while silently mis-slotting + such a drafter -- and mis-slotting costs acceptance without ever failing. + """ + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + from tensorrt_llm.mapping import Mapping + + worker = DSparkWorker(_routing_config(False), Mapping()) + legacy = types.SimpleNamespace(_dspark_shift_label=False, has_markov_head=False) + + # shift_label off -> the base class' slots 1..K, not the dspark 0..K-1. + ids = worker._draft_slot_ids(legacy, num_gens=2, block_size=5, num_draft_tokens=3) + base_ids = DFlashWorker._draft_slot_ids( + worker, legacy, num_gens=2, block_size=5, num_draft_tokens=3 + ) + assert ids.tolist() == base_ids.tolist() + + # 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 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 1a061c8daa00..ff3af9a76b90 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,9 @@ confidence_proj weights load without being used. """ +import re +from types import SimpleNamespace + import pytest import torch import torch.nn.functional as F @@ -205,11 +208,29 @@ def test_swa_window_conventions(): NUM_CAPTURE = 2 -def _tiny_config(dspark: bool): +def _tiny_config(dspark: bool, *, published_spelling: bool = False): + """Tiny drafter config. + + ``published_spelling`` reproduces the two public K3 DSpark checkpoints + (RadixArk, Inferact): the head switches sit at the TOP level, the + confidence flag is ``enable_confidence_head``, and neither ``shift_label`` + nor ``projector_type`` is declared at all -- shift_label rides on the + DSpark default. Reading only ``dflash_config`` resolves markov_rank to 0 + there, which drops the heads without raising. + """ from transformers import Qwen3Config cfg = dict(TINY) dflash = {"mask_token_id": VOCAB - 2, "target_layer_ids": [0, 1]} + if dspark and published_spelling: + cfg.update( + markov_rank=RANK, + markov_head_type="vanilla", + enable_confidence_head=True, + confidence_head_with_markov=True, + ) + cfg["dflash_config"] = dflash + return Qwen3Config.from_dict(cfg) if dspark: dflash.update( projector_type="dspark", @@ -227,7 +248,13 @@ def _tiny_config(dspark: bool): return Qwen3Config.from_dict(cfg) -def _tiny_weights(seed=7): +def _tiny_weights(seed=7, *, published_head_keys=False): + """Tiny drafter weights. + + ``published_head_keys`` names the head tensors the way both public + checkpoints ship them -- after the submodules that own them -- instead of + the bare spellings the DSv4 stage weights use. + """ g = torch.Generator().manual_seed(seed) def rnd(*shape): @@ -235,15 +262,25 @@ def rnd(*shape): h, inter = TINY["hidden_size"], TINY["intermediate_size"] nh, nkv, hd = (TINY["num_attention_heads"], TINY["num_key_value_heads"], TINY["head_dim"]) - w = { - "fc.weight": rnd(h, h * NUM_CAPTURE), - "hidden_norm.weight": rnd(h) + 1.0, - "norm.weight": rnd(h) + 1.0, + head = { "markov_w1.weight": rnd(VOCAB, RANK), "markov_w2.weight": rnd(VOCAB, RANK), "confidence_proj.weight": rnd(1, h + RANK), "confidence_proj.bias": rnd(1), } + if published_head_keys: + head = { + "markov_head.markov_w1.weight": head["markov_w1.weight"], + "markov_head.markov_w2.weight": head["markov_w2.weight"], + "confidence_head.proj.weight": head["confidence_proj.weight"], + "confidence_head.proj.bias": head["confidence_proj.bias"], + } + w = { + "fc.weight": rnd(h, h * NUM_CAPTURE), + "hidden_norm.weight": rnd(h) + 1.0, + "norm.weight": rnd(h) + 1.0, + **head, + } for i in range(TINY["num_hidden_layers"]): p = f"layers.{i}." w[p + "self_attn.q_proj.weight"] = rnd(nh * hd, h) @@ -260,13 +297,24 @@ def rnd(*shape): return w -def _build_drafter(dspark: bool, weights): +def _build_drafter( + dspark: bool, + weights, + *, + published_spelling: bool = False, + dflash_attention_backend: str = "VANILLA", +): from tensorrt_llm._torch.model_config import ModelConfig - model_config = ModelConfig(pretrained_config=_tiny_config(dspark), attn_backend="TRTLLM") + model_config = ModelConfig( + pretrained_config=_tiny_config(dspark, published_spelling=published_spelling), + attn_backend="TRTLLM", + ) # The DSpark head set lives in the DSpark drafter, not in the DFlash base. drafter_cls = GQADSparkForCausalLM if dspark else DFlashForCausalLM - drafter = drafter_cls(model_config).to("cuda") + drafter = drafter_cls(model_config, dflash_attention_backend=dflash_attention_backend).to( + "cuda" + ) # Drop dspark head tensors for the plain drafter (schema without them). if not dspark: weights = {k: v for k, v in weights.items() if not k.startswith(("markov_", "confidence_"))} @@ -371,6 +419,96 @@ def test_dspark_drafter_loads_head_weights_and_parses_config(): assert drafter.confidence_proj_bias is not None +@needs_gpu +def test_published_drafter_spelling_activates_the_heads(): + """Both public K3 DSpark checkpoints load with their heads live. + + They declare the switches at the top level and name the head tensors after + the owning submodules. Reading only ``dflash_config`` and the bare tensor + names resolves markov_rank to 0 and drops markov_w1/w2 on the floor: + correct output, lower acceptance, nothing raised. + """ + weights = _tiny_weights(published_head_keys=True) + drafter = _build_drafter(True, weights, published_spelling=True) + + assert drafter.has_markov_head, "markov weights dropped despite being in the checkpoint" + assert drafter._dspark_use_confidence_head, "enable_confidence_head spelling not resolved" + # Declared nowhere in the published config, so it rides on the DSpark + # default. False would run slots 1..K on a block_size-K drafter and read + # the next request's anchor slot. + assert drafter._dspark_shift_label + torch.testing.assert_close(drafter.markov_w1.cpu(), weights["markov_head.markov_w1.weight"]) + assert drafter.confidence_proj_bias is not None + + +@needs_gpu +def test_head_weights_without_a_resolvable_rank_raise(): + """The inverse of the missing-weights check. + + A checkpoint that ships markov_w1/w2 while the rank resolves to 0 means the + switches were spelled somewhere this build cannot read. Loading it anyway + would silently cost acceptance, so it is an error. + """ + from tensorrt_llm._torch.model_config import ModelConfig + + # dspark head weights, but a config that declares no head switches at all. + model_config = ModelConfig(pretrained_config=_tiny_config(False), attn_backend="TRTLLM") + drafter = GQADSparkForCausalLM(model_config).to("cuda") + + with pytest.raises(ValueError, match="markov_rank resolved to 0"): + drafter.load_weights(_tiny_weights(published_head_keys=True)) + + +def test_dflash_refuses_a_drafter_that_declares_the_dspark_heads(): + """``decoding_type: DFlash`` must reject a DSpark drafter, not degrade it. + + DFlash does not implement the Markov / confidence / shift_label semantics, + so serving one here would lower the acceptance rate with no error and no + way to attribute it back. + """ + from tensorrt_llm._torch.models import modeling_dflash + + model_config = SimpleNamespace( + spec_config=SimpleNamespace(attention_backend="TRTLLM", speculative_model="/nonexistent") + ) + draft_config = SimpleNamespace(pretrained_config=_tiny_config(True)) + + with pytest.raises(ValueError, match="DSpark"): + modeling_dflash._build_dflash_draft(model_config, draft_config, None, None) + + +@pytest.mark.parametrize( + "layers,expected", + [ + # MLA-shaped: no per-head q/k/v projection to split at all. + ([SimpleNamespace(self_attn=SimpleNamespace())], "qkv_proj"), + # GQA-shaped but heterogeneous: the cross-layer K/V fusion needs one + # uniform num_key_value_heads. + ( + [ + SimpleNamespace(self_attn=SimpleNamespace(qkv_proj=object(), num_key_value_heads=n)) + for n in (8, 8, 4) + ], + "[2]", + ), + ], + ids=["no_fused_qkv", "mismatched_kv_heads"], +) +def test_block_decode_rejects_a_backbone_it_cannot_express(layers, expected): + """The GQA precondition replaced the old per-model_type whitelist. + + Without it the registry happily builds an unsupported backbone and the + failure surfaces much later inside _build_fused_kv_buffers, or as silently + mis-sliced weights. + """ + drafter = DFlashForCausalLM.__new__(DFlashForCausalLM) + drafter.model = SimpleNamespace(layers=layers) + drafter.config = SimpleNamespace() + + with pytest.raises(ValueError, match=re.escape(expected)): + drafter._validate_gqa_shape() + + @needs_gpu def test_plain_dflash_drafter_keeps_old_gates(): """No-regression: a config WITHOUT dspark fields resolves to the exact diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_spec_mode_gate.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_spec_mode_gate.py deleted file mode 100644 index d98fdbae4572..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_spec_mode_gate.py +++ /dev/null @@ -1,110 +0,0 @@ -# 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. -"""The Kimi K3 target's spec-dec mode gate. - -``KimiLinearForCausalLM.__init__`` whitelists the speculative-decoding modes the -target will serve. It is a plain assert on the *target* side, so nothing that -exercises the drafter, the builder or the worker can reach it -- which is how a -K3 engine configured with ``decoding_type: DSpark`` got through every unit test -and then failed at model construction with - - AssertionError: Kimi K3 supports speculative decoding only with SA or DFlash - -The gate is the second statement in ``__init__``, and everything past it builds -the real K3 model, which needs weights and 16 GPUs. So rather than stub the -framework out from under it -- the base initializer's arguments construct -``KimiLinearModel`` before the base is even called, so stubbing the base does -not help -- these tests ask a narrower question: did construction fail *at the -gate*, or did it get past it? Anything that fails later has passed the gate, -which is the whole of what is under test here. -""" - -from types import SimpleNamespace - -import pytest - -from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM -from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode - -# Admitted: SA drafts in-forward with no draft weights; DFlash and DSpark are -# the external-drafter flow, and the target side is identical for both (the -# hidden-state capture in KimiLinearModel.forward is unconditional). -ADMITTED = [ - SpeculativeDecodingMode.SA, - SpeculativeDecodingMode.DFLASH, - SpeculativeDecodingMode.DSPARK, -] -# Refused: these need draft heads that no K3 checkpoint ships. -REFUSED = [ - SpeculativeDecodingMode.MTP, - SpeculativeDecodingMode.EAGLE3_ONE_MODEL, -] - -_SPEC_GATE = "speculative decoding" -_PP_GATE = "pipeline parallelism" - - -def _model_config(mode, *, pp_size=1): - """The minimum ``__init__`` reads before it reaches the gate.""" - return SimpleNamespace( - pretrained_config=SimpleNamespace(model_type="kimi_linear", linear_attn_config={}), - mapping=SimpleNamespace(pp_size=pp_size), - spec_config=None if mode is None else SimpleNamespace(spec_dec_mode=mode), - ) - - -def _rejected_by(model_config) -> str | None: - """Which gate rejected this config, or None if construction got past them. - - Only the two guard asserts count as a rejection; construction is expected - to fail afterwards on the real model, and that failure means the config was - admitted. An AssertionError from anywhere else is a genuine problem and is - re-raised rather than silently read as a rejection. - """ - try: - KimiLinearForCausalLM(model_config) - except AssertionError as exc: - message = str(exc) - for gate in (_SPEC_GATE, _PP_GATE): - if gate in message: - return gate - raise - except Exception: - return None - return None - - -@pytest.mark.parametrize("mode", ADMITTED, ids=lambda m: m.name) -def test_admitted_modes_pass_the_gate(mode): - assert _rejected_by(_model_config(mode)) is None - - -@pytest.mark.parametrize("mode", REFUSED, ids=lambda m: m.name) -def test_refused_modes_are_rejected_at_the_spec_gate(mode): - assert _rejected_by(_model_config(mode)) == _SPEC_GATE - - -def test_the_refusal_message_names_the_admitted_modes(): - # The message is the only guidance a user gets, so it has to name the modes - # that would work -- that is what turns "not supported" into an action. - with pytest.raises(AssertionError, match="SA, DFlash or DSpark"): - KimiLinearForCausalLM(_model_config(SpeculativeDecodingMode.MTP)) - - -def test_pipeline_parallelism_is_still_rejected(): - # The pp guard sits ahead of the spec gate; pin the order so a future edit - # to the mode list cannot let a pp>1 config through. - config = _model_config(SpeculativeDecodingMode.DSPARK, pp_size=2) - assert _rejected_by(config) == _PP_GATE diff --git a/tests/unittest/others/test_lazy_model_zoo.py b/tests/unittest/others/test_lazy_model_zoo.py index 52d80a1fa460..efb50f78dffe 100644 --- a/tests/unittest/others/test_lazy_model_zoo.py +++ b/tests/unittest/others/test_lazy_model_zoo.py @@ -205,6 +205,19 @@ def _decorated_draft_model_registrations(): def test_spec_mode_index_matches_decorators(): + # SPEC_MODE_TO_MODULE is the fourth hand-maintained table in _arch_index, + # and every other one already has a drift test here -- MODEL_ARCH_TO_MODULE + # and MULTIMODAL_MODEL_TYPE_TO_MODULE in test_arch_index_matches_decorators, + # MODEL_CLASS_TO_MODULE in test_class_index_matches_package_all. Keeping the + # set complete is the point: without this, the spec-mode table would be the + # only index whose entries nothing checks. + # + # Drift is not silent, but it is misattributed. A builder added without its + # index entry (or an index entry whose module stopped declaring the mode) + # surfaces as "unsupported speculative decoding mode" to whoever next runs + # that mode -- which reads as "this algorithm is not implemented", not as + # "someone forgot a line in _arch_index". The person who sees it is rarely + # the person who caused it. from tensorrt_llm._torch.models._arch_index import SPEC_MODE_TO_MODULE mode_truth = _decorated_draft_model_registrations() @@ -221,32 +234,6 @@ def test_spec_mode_index_matches_decorators(): assert not wrong, f"index points at the wrong module: {wrong}" -def test_spec_mode_index_resolves_every_builder(): - # End-to-end check of the lazy path: every indexed mode must resolve to a - # builder through the single entry point, and that builder must itself - # declare the mode. The declaration is read off the function - # (``_registered_spec_modes``), never by scanning the mapping by identity: - # a built-in builder overridden by an external registration keeps its - # attribute but loses its mapping slot. - from tensorrt_llm._torch.models._arch_index import SPEC_MODE_TO_MODULE - from tensorrt_llm._torch.models.modeling_utils import ( - _REGISTERED_SPEC_MODES_ATTR, - get_registered_draft_model_builder, - ) - from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode - - for mode_name in SPEC_MODE_TO_MODULE: - mode = getattr(SpeculativeDecodingMode, mode_name, None) - assert mode is not None, f"{mode_name} is not a SpeculativeDecodingMode" - builder = get_registered_draft_model_builder(mode) - assert builder is not None, f"no draft-model builder resolved for {mode_name}" - declared = getattr(builder, _REGISTERED_SPEC_MODES_ATTR, set()) - assert mode in declared, ( - f"{builder.__module__}.{builder.__qualname__} is registered for " - f"{mode_name} but does not declare it" - ) - - def test_class_index_matches_package_all(): # MODEL_CLASS_TO_MODULE is the one table with no decorator to mirror: it # backs PEP 562 attribute access on the models package. Every name in the From 05dba08f55cca263b105a6427caccf10f70f8e32 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Mon, 24 Aug 2026 02:19:02 -0700 Subject: [PATCH 18/21] [None][feat] Let callers override an lm-eval task's shot count GSM8K's task yaml pins 5 shots, so a 0-shot chat evaluation was unreachable from either the CLI or a test. Forwards num_fewshot to the same task_obj.set_config lm-eval's own simple_evaluate makes. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/evaluate/lm_eval.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/evaluate/lm_eval.py b/tensorrt_llm/evaluate/lm_eval.py index d68aad84b4d7..bfc487767835 100644 --- a/tensorrt_llm/evaluate/lm_eval.py +++ b/tensorrt_llm/evaluate/lm_eval.py @@ -834,7 +834,8 @@ def __init__(self, output_path: Optional[str] = None, output_dir: Optional[str] = None, post_process_fn: Optional[Callable[[str], str]] = None, - preserve_caller_max_tokens: bool = False): + preserve_caller_max_tokens: bool = False, + num_fewshot: Optional[int] = None): try: import lm_eval except ImportError as e: @@ -889,6 +890,15 @@ def _adjust_config(task_dict, random_seed): else: # NOTE: Few-shot random seed task_obj.set_fewshot_seed(seed=random_seed) + # Caller override of the task yaml's shot count, the same + # call lm-eval's own simple_evaluate makes. Without it a + # 0-shot chat evaluation of a task whose yaml pins 5 shots + # is unreachable, which is the regime a chat-distilled + # speculative drafter has to be measured in. + if num_fewshot is not None: + task_obj.set_config(key="num_fewshot", + value=num_fewshot) + logger.info(f"num_fewshot overridden to {num_fewshot}") adjusted_task_dict[task_name] = task_obj # NOTE: Shuffle dataset @@ -1054,6 +1064,7 @@ def command_harness(cls, ctx, **kwargs): random_seed=kwargs.pop("random_seed", 0), apply_chat_template=kwargs.pop("apply_chat_template", False), fewshot_as_multiturn=kwargs.pop("fewshot_as_multiturn", False), + num_fewshot=kwargs.pop("num_fewshot", None), system_prompt=kwargs.pop("system_prompt", None), is_multimodal=kwargs.pop("is_multimodal", False), chat_template_kwargs=kwargs.pop("chat_template_kwargs", None), @@ -1129,6 +1140,12 @@ def __init__(self, **kwargs): is_flag=True, default=False, help="Apply fewshot as multiturn.") + @click.option("--num_fewshot", + type=int, + default=None, + help="Override the task yaml's shot count. Use 0 with " + "--apply_chat_template for a single-question chat " + "evaluation.") @click.option("--system_prompt", type=str, default=None, From 716d91f0994832fd626ddbc0c477234f92e48109 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Mon, 24 Aug 2026 02:19:09 -0700 Subject: [PATCH 19/21] [None][test] Add a 0-shot acceptance-length test for the standalone DSpark path Runs GSM8K on DeepSeek's published Qwen3-8B block-7 drafter, so the published head spellings stay exercised. 0-shot chat, because a DSpark drafter is distilled on the target's chat output and the harness default 5-shot completion prompt understates acceptance (4.42 vs 6.26); that regime gets its own accuracy reference under extra_acc_spec. Signed-off-by: Zhenhuan Chen --- .../references/acceptance_length.yaml | 3 + .../defs/accuracy/references/gsm8k.yaml | 9 +++ .../defs/accuracy/test_llm_api_pytorch.py | 67 +++++++++++++++++++ .../test_lists/test-db/l0_b200.yml | 1 + .../test_lists/test-db/l0_h100.yml | 1 + 5 files changed, 81 insertions(+) diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml index 9d19aa27734f..fe695e2210f6 100644 --- a/tests/integration/defs/accuracy/references/acceptance_length.yaml +++ b/tests/integration/defs/accuracy/references/acceptance_length.yaml @@ -39,3 +39,6 @@ TestQwen3_5_4B::test_dflash: TestKimiK3::test_w4a16_mxfp4: ref_al: 1.318 min_al: 1.15 +TestQwen3_8B::test_dspark: + ref_al: 6.259 + min_al: 5.946 diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index 0d98e20d649b..1a5d5adde3d5 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -192,6 +192,15 @@ Qwen3/Qwen3-8B: accuracy: 87.1114 - spec_dec_algo: DFlash accuracy: 87.1114 + # accuracy 0 means "run the task, do not gate on it" -- the same value + # TRTLLM_ACCURACY_NO_REFERENCE and INTEGRATION_TEST use, and the threshold + # (ref + z_alpha*scale) is then always cleared. test_dspark gates on + # acceptance length instead. Accuracy is not meaningful in its 0-shot chat + # regime: the model stops emitting "#### N", so strict-match is 0 and the + # reported score is the mean of that and flexible-extract (28.92 measured). + - spec_dec_algo: DSpark + extra_acc_spec: zero_shot_al_only + accuracy: 0 - quant_algo: FP8 kv_cache_quant_algo: FP8 accuracy: 87.1114 diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index b5cb009adf27..89d4ac7d590f 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -5061,6 +5061,73 @@ def test_dflash(self): task = GSM8K(self.MODEL_NAME) task.evaluate(llm) + @skip_pre_hopper + @pytest.mark.parametrize( + "attention_backend", + ["VANILLA", pytest.param("TRTLLM", marks=skip_pre_blackwell)]) + def test_dspark(self, attention_backend): + """Standalone DSpark drafter on Qwen3-8B, 0-shot chat. + + Acceptance length is the gate here. DSpark drafters are distilled on + the target's own chat-mode generations, so the harness default (5-shot + completion) is out of distribution for the drafter and understates + acceptance: 4.42 there against 6.26 here, on the same checkpoints. + GSM8K accuracy is not meaningful in this regime and is not gated -- + see the extra_acc_spec entry in references/gsm8k.yaml. + + Both block-decode backends run: TRTLLM is what deployments use, and + it is the one the acceptance number above was measured on, but it + needs SM100/SM103 so H100 only covers VANILLA. The two differ + numerically (GSM8K 28.9 vs 28.1) while landing the same acceptance + length, so they share one reference. + """ + pytorch_config = dict( + max_batch_size=8, + disable_overlap_scheduler=True, + cuda_graph_config=CudaGraphConfig(max_batch_size=8, + enable_padding=True), + ) + kv_cache_config = KvCacheConfig(enable_block_reuse=False, + free_gpu_memory_fraction=0.6) + + # DeepSeek's official DSpark head for this target (block_size 7, + # markov_rank 256, confidence head). Head tensors are named after the + # submodules that own them (markov_head.*, confidence_head.proj.*), + # which is the spelling the drafter loader has to resolve. + dspark_model_dir = ( + f"{llm_models_root()}/dspark/dspark_qwen3_8b_block7") + target_model_dir = f"{llm_models_root()}/Qwen3/Qwen3-8B" + + spec_config = DSparkDecodingConfig(max_draft_len=7, + speculative_model=dspark_model_dir, + attention_backend=attention_backend) + + with LLM(model=target_model_dir, + **pytorch_config, + kv_cache_config=kv_cache_config, + max_stats_len=-1, + enable_iter_perf_stats=True, + speculative_config=spec_config) as llm: + task = GSM8K(self.MODEL_NAME) + # 0-shot chat, not the harness default 5-shot completion: a DSpark + # drafter is distilled on the target's chat-mode output, so the + # default prompt is out of distribution for it and understates + # acceptance (4.42 vs 6.26 on these same checkpoints). Only the AL + # is gated -- extra_acc_spec selects a gsm8k.yaml entry whose + # reference accuracy is 0, i.e. run the task, do not gate on it. + task.evaluate(llm, + extra_acc_spec="zero_shot_al_only", + extra_evaluator_kwargs=dict( + num_fewshot=0, + apply_chat_template=True, + chat_template_kwargs={"enable_thinking": False}, + )) + acceptance_length = _compute_acceptance_length(llm) + print(f"[AL] test_dspark[{attention_backend}] acceptance_length " + f"= {acceptance_length:.3f}") + assert_acceptance_length("TestQwen3_8B::test_dspark", + acceptance_length) + @skip_pre_blackwell @pytest.mark.parametrize("tp_size,pp_size,ep_size,attention_dp", [(1, 1, 1, False)], diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index a4a87cd5f8c0..9e5ef4da23b1 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -52,6 +52,7 @@ l0_b200: - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_1gpu[v2_kv_cache-True-True-trtllm-fp8] - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_dummy_load_format - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_fp8_block_scales[latency] # Cover nvbugs 5461712 and 5505402 + - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_dspark[TRTLLM] # SM100+ only; l0_h100 runs [VANILLA] - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[latency_moe_trtllm-torch_compile=True] - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_w4a8_mxfp4[fp8-latency-TRTLLM] - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_w4a8_mxfp4[mxfp8-latency-TRTLLM] diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index d4da554eaffb..800c465dd6e6 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -169,6 +169,7 @@ l0_h100: - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_dummy_load_format - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_eagle3[eagle3_one_model=True-enable_chunked_prefill=False-enable_max_concurrency=False-enable_draft_len_schedule=False] - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_dflash + - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_dspark[VANILLA] - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_dflash - accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_bf16 - accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_fp8 From 05b8e73b400ee538439f832b835bd551d4002211 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Mon, 24 Aug 2026 02:19:10 -0700 Subject: [PATCH 20/21] [None][fix] Accept an unset speculative_model for DSpark, as MTP does MTP treats an unset (or target-equal) speculative_model as "load the draft from the target checkpoint" -- resolve_mtp_checkpoint_source. DSpark rejected it outright, so expressing the same intent needed different config depending on which speculative algorithm was in use. An unset speculative_model now defaults to the target, which the existing embedded-vs-standalone probe then resolves to the embedded DeepSeek-V4-Pro flavour. Pointing speculative_model at the target explicitly stays equivalent. A target that declares no mtp.* draft weights is still refused rather than defaulted: without that, a standalone drafter gets built from the target's own config and fails much later on a missing fc.weight. MTP has no equivalent refusal, so this extends the shared convention rather than diverging from it. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/llmapi/llm_args.py | 27 ++++++++++++++------------ tests/unittest/llmapi/test_llm_args.py | 10 +++++----- 2 files changed, 20 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 4350be676078..a149f7eddcce 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -6054,19 +6054,22 @@ def validate_speculative_config(self): if not spec_cfg.max_draft_len: raise ValueError("DSpark max_draft_len must be > 0; got " f"{spec_cfg.max_draft_len}") - # The DSpark draft weights live in the ``mtp.*`` namespace of a - # local checkpoint directory; without ``speculative_model`` - # neither the draft weights nor the ``dspark_*`` config - # defaults can be located, and engine construction would fail - # much later with an opaque error. + # Same convention as MTP (see resolve_mtp_checkpoint_source): + # an unset speculative_model means "the draft lives in the + # target checkpoint". For DSpark that is the embedded + # DeepSeek-V4-Pro flavour, whose draft is the target's mtp.* + # namespace. Defaulting rather than rejecting keeps the + # spec-dec API consistent across algorithms; pointing + # speculative_model at the target explicitly is equivalent. if spec_cfg.speculative_model is None: - raise ValueError( - "DSpark requires speculative_config.speculative_model " - "to point at the drafter's checkpoint directory: a " - "standalone DSpark drafter repository, or -- for the " - "embedded DeepSeek-V4-Pro flavour, whose draft weights " - "live in the mtp.* namespace -- the target checkpoint " - "directory itself.") + spec_cfg.speculative_model = self.model + if not spec_cfg.draft_is_embedded_in_target: + raise ValueError( + "speculative_config.speculative_model is unset, " + "which means 'load the draft from the target " + "checkpoint', but the target has no mtp.* draft " + "weights. Point speculative_model at a standalone " + "DSpark drafter checkpoint directory.") # Warm the embedded-vs-standalone probe here, while we are in # the main process and the checkpoint path is known local. The # flag then travels with the config, so no rank repeats the diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index b4b1569ed66d..7c7722c794ec 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -457,13 +457,13 @@ def test_dspark_target_layer_ids_order_mismatch_rejected(tmp_path): @pytest.mark.cpu_only def test_dspark_requires_speculative_model(): - # The DSpark draft weights live in the checkpoint's mtp.* namespace, so an - # unset speculative_model must fail fast at config validation instead of - # raising an opaque TypeError deep inside engine construction. + # Unset speculative_model means "load the draft from the target", as it + # does for MTP. /tmp/dummy_model carries no mtp.* draft weights, so that + # must fail fast at config validation instead of raising an opaque + # TypeError deep inside engine construction. spec_cfg = DSparkDecodingConfig(max_draft_len=5) - with pytest.raises(ValueError, - match="requires speculative_config.speculative_model"): + with pytest.raises(ValueError, match="speculative_model is unset"): TorchLlmArgs( model="/tmp/dummy_model", skip_tokenizer_init=True, From 0b5bdae48c3d9d69fea62af462ad8c952f9ad7b0 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Wed, 26 Aug 2026 02:36:05 -0700 Subject: [PATCH 21/21] [None][fix] Size the draft block by the drafter's slot convention #17935 added a block-width check assuming the plain DFlash layout, where slot 0 holds the anchor and K draft tokens need K+1 slots. DSpark's shift_label reads slots 0..K-1, so K slots suffice; the extra slot rejected both published block-7 drafters at max_draft_len=7. Width now comes from _draft_block_width, next to the _draft_slot_ids hook it has to agree with. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/speculative/dflash.py | 17 ++++++++++++++--- tensorrt_llm/_torch/speculative/dspark.py | 12 ++++++++++++ 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 40dc16711f51..b92c1a7c7ac2 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -324,11 +324,12 @@ def _lazy_init_ctx_buffers(self, draft_model, spec_metadata, attn_metadata): self.max_draft_len + 1 ) # what the draft forward actually computes - self._compute_block_size = self.max_draft_len + 1 - if self.max_draft_len + 1 > self._resolved_block_size: + self._compute_block_size = self._draft_block_width(draft_model) + if self._compute_block_size > self._resolved_block_size: + slack = self._compute_block_size - self.max_draft_len raise ValueError( f"DFlash checkpoint was trained with block_size={self._resolved_block_size}." - f"Lower max_draft_len to at most {self._resolved_block_size - 1}." + f"Lower max_draft_len to at most {self._resolved_block_size - slack}." f"Current max_draft_len={self.max_draft_len}." ) @@ -798,6 +799,16 @@ def _forward_impl( "next_new_tokens": next_new_tokens, } + def _draft_block_width(self, draft_model) -> int: + """Block slots the draft forward must compute for K draft tokens. + + Plain DFlash reads slots 1..K, so slot 0 is pure overhead and the + forward needs K+1. Families whose slot convention differs override + this alongside :meth:`_draft_slot_ids` — the two must agree, since + the width bounds the slots the gather is allowed to name. + """ + return self.max_draft_len + 1 + def _draft_slot_ids( self, draft_model, num_gens: int, block_size: int, num_draft_tokens: int ) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index e6ddbc5904e1..733b44515ea8 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -759,6 +759,18 @@ def set_draft_model(self, draft_model) -> None: f"{type(draft_model).__name__} declares one." ) + def _draft_block_width(self, draft_model) -> int: + """Block width under the dspark ``shift_label`` convention. + + shift_label reads slots 0..K-1, so K draft tokens fit in K slots and + the base class' K+1 over-demands by one -- enough to reject a block-7 + checkpoint at max_draft_len=7, which is how both published DSpark + drafters are meant to run. + """ + if getattr(draft_model, "_dspark_shift_label", False): + return self.max_draft_len + return super()._draft_block_width(draft_model) + def _draft_slot_ids( self, draft_model, num_gens: int, block_size: int, num_draft_tokens: int ) -> torch.Tensor: