[None][feat] Port DFlash weight-load checks, DSpark norm_dim and disagg tests from the Rubin branch

tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py

@@ -41,6 +41,7 @@ def is_fused_dspark_rmsnorm_rope_supported(
     freqs: torch.Tensor,
     num_heads: int,
     rope_dim: int,
+    norm_dim: int | None = None,
 ) -> bool:
     """Return whether tensors satisfy the production fused-op contract."""
     if _get_dspark_arch_str() is None or not all(t.is_cuda for t in (x, weight, freqs)):
@@ -51,10 +52,15 @@ def is_fused_dspark_rmsnorm_rope_supported(
         return False
     if x.ndim < 2 or x.shape[-1] % 32 != 0:
         return False
-    if weight.shape != (x.shape[-1],):
-        return False
     if rope_dim < 0 or rope_dim > x.shape[-1] or rope_dim % 2 != 0:
         return False
+    # norm_dim defaults to the whole row; the only other supported value is the
+    # nope prefix, where the weight spans just that prefix.
+    effective_norm_dim = x.shape[-1] if norm_dim is None else norm_dim
+    if effective_norm_dim not in (x.shape[-1], x.shape[-1] - rope_dim):
+        return False
+    if effective_norm_dim % 32 != 0 or weight.shape != (effective_norm_dim,):
+        return False
     if (x.shape[-1] - rope_dim) % 32 != 0 or (rope_dim // 2) % 32 != 0:
         return False
     rows = x.numel() // x.shape[-1]
@@ -144,14 +150,15 @@ def _compile_fused_dspark_rmsnorm_rope(
     apply_weight: bool,
     apply_rmsnorm: bool,
     inverse_rope: bool,
+    norm_dim: int,
 ):
     rows = cute.sym_int()
     freq_rows = cute.sym_int()
     x_fake = cute.runtime.make_fake_compact_tensor(
         cutlass.BFloat16, (rows, hidden_dim), stride_order=(1, 0)
     )
     weight_fake = cute.runtime.make_fake_compact_tensor(
-        cutlass.BFloat16, (hidden_dim,), stride_order=(0,)
+        cutlass.BFloat16, (norm_dim,), stride_order=(0,)
     )
     freqs_fake = cute.runtime.make_fake_compact_tensor(
         cutlass.Float32,
@@ -170,6 +177,7 @@ def _compile_fused_dspark_rmsnorm_rope(
         apply_weight,
         apply_rmsnorm,
         inverse_rope,
+        norm_dim=norm_dim,
     )
     return cute.compile(
         kernel,
@@ -296,9 +304,10 @@ def cute_dsl_dspark_rmsnorm_rope(
     apply_weight: bool,
     apply_rmsnorm: bool,
     inverse_rope: bool,
+    norm_dim: int | None = None,
 ) -> torch.Tensor:
     """Apply fused RMSNorm and adjacent-pair RoPE to contiguous BF16 rows."""
-    if not is_fused_dspark_rmsnorm_rope_supported(x, weight, freqs, num_heads, rope_dim):
+    if not is_fused_dspark_rmsnorm_rope_supported(x, weight, freqs, num_heads, rope_dim, norm_dim):
         raise ValueError(
             "cute_dsl_dspark_rmsnorm_rope requires contiguous BF16 tensors on "
             "an SM100 or SM103 GPU with a valid FP32 frequency view; "
@@ -316,6 +325,7 @@ def cute_dsl_dspark_rmsnorm_rope(
         apply_weight,
         apply_rmsnorm,
         inverse_rope,
+        x.shape[-1] if norm_dim is None else norm_dim,
     )
     compiled(x_flat, weight, freqs, output)
     return output.view(original_shape)
@@ -332,6 +342,7 @@ def _(
     apply_weight: bool,
     apply_rmsnorm: bool,
     inverse_rope: bool,
+    norm_dim: int | None = None,
 ) -> torch.Tensor:
     return torch.empty_like(x)
 

tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dspark_rmsnorm_rope.py

@@ -26,6 +26,7 @@ def __init__(
         apply_weight: bool,
         apply_rmsnorm: bool,
         inverse_rope: bool,
+        norm_dim: int | None = None,
     ):
         if hidden_dim % self.num_threads != 0:
             raise ValueError(
@@ -44,6 +45,22 @@ def __init__(
         self.apply_weight = apply_weight
         self.apply_rmsnorm = apply_rmsnorm
         self.inverse_rope = inverse_rope
+        # Which prefix the RMS is taken over, and how far the norm scale and the
+        # weight reach. DSpark normalizes the whole row (the default, bit-identical
+        # to before this knob existed); DeepSeek-style MLA normalizes only the
+        # kv_lora_rank latent and leaves k_pe raw, which is nope_dim here.
+        self.norm_dim = hidden_dim if norm_dim is None else norm_dim
+        if self.norm_dim not in (hidden_dim, self.nope_dim):
+            raise ValueError(
+                f"norm_dim must be hidden_dim ({hidden_dim}) or nope_dim "
+                f"({self.nope_dim}); got {self.norm_dim}"
+            )
+        self.norm_covers_rope = self.norm_dim == hidden_dim
+        if self.norm_dim % self.num_threads != 0:
+            raise ValueError(
+                f"norm_dim must be divisible by {self.num_threads}; got {self.norm_dim}"
+            )
+        self.norm_elements_per_thread = self.norm_dim // self.num_threads
         if self.nope_dim % self.num_threads != 0:
             raise ValueError(
                 f"nope_dim must be divisible by {self.num_threads}; got {self.nope_dim}"
@@ -85,12 +102,12 @@ def kernel(
         inverse_rms = cutlass.Float32(1.0)
         if cutlass.const_expr(self.apply_rmsnorm):
             sum_sq = cutlass.Float32(0.0)
-            for item in cutlass.range_constexpr(self.elements_per_thread):
+            for item in cutlass.range_constexpr(self.norm_elements_per_thread):
                 dim = tidx + item * self.num_threads
                 value = cutlass.Float32(x[row, dim])
                 sum_sq += value * value
             sum_sq = cute.arch.warp_reduction_sum(sum_sq)
-            inverse_rms = cute.math.rsqrt(sum_sq / self.hidden_dim + self.eps)
+            inverse_rms = cute.math.rsqrt(sum_sq / self.norm_dim + self.eps)
 
         for item in cutlass.range_constexpr(self.nope_elements_per_thread):
             dim = tidx + item * self.num_threads
@@ -105,11 +122,16 @@ def kernel(
                 pair = tidx + item * self.num_threads
                 real_dim = self.nope_dim + pair * 2
                 imag_dim = real_dim + 1
-                real = cutlass.Float32(x[row, real_dim]) * inverse_rms
-                imag = cutlass.Float32(x[row, imag_dim]) * inverse_rms
-                if cutlass.const_expr(self.apply_weight):
-                    real *= cutlass.Float32(weight[real_dim])
-                    imag *= cutlass.Float32(weight[imag_dim])
+                real = cutlass.Float32(x[row, real_dim])
+                imag = cutlass.Float32(x[row, imag_dim])
+                # Outside norm_dim the rope lanes are passed through raw: no RMS
+                # scale, no weight. weight is only norm_dim long in that case.
+                if cutlass.const_expr(self.norm_covers_rope):
+                    real *= inverse_rms
+                    imag *= inverse_rms
+                    if cutlass.const_expr(self.apply_weight):
+                        real *= cutlass.Float32(weight[real_dim])
+                        imag *= cutlass.Float32(weight[imag_dim])
                 cos = cutlass.Float32(freqs[freq_row, pair, 0])
                 sin = cutlass.Float32(freqs[freq_row, pair, 1])
                 if cutlass.const_expr(self.inverse_rope):

tensorrt_llm/_torch/models/modeling_dflash.py

@@ -26,7 +26,7 @@
     get_dflash_trtllm_gen_ops,
 )
 from ..speculative.interface import SpeculativeDecodingMode
-from .modeling_utils import get_model_architecture, register_draft_model
+from .modeling_utils import FUSED_MODULE_COMPONENTS, get_model_architecture, register_draft_model
 
 
 def dspark_layer_window_size(
@@ -323,6 +323,12 @@ class DFlashForCausalLM(nn.Module):
     Reference: https://arxiv.org/pdf/2602.06036
     """
 
+    # Whether ``dflash_attention_backend`` drives this drafter's block decode.
+    # Subclasses that bring their own attention set this False: neither backend
+    # can express every drafter shape, the ops behind them are optional
+    # dependencies, and the worker's per-backend shape checks do not apply.
+    _uses_worker_attention_backend = True
+
     def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"):
         """Build the draft model, resolving its architecture from the draft config
         (falling back to a model_type-derived name when the checkpoint uses a
@@ -374,23 +380,32 @@ def __init__(self, draft_config, *, dflash_attention_backend: str = "VANILLA"):
             "block_size", getattr(pretrained_config, "block_size", None)
         )
         self.dflash_attention_backend = dflash_attention_backend
+        if self.dflash_attention_backend not in ("VANILLA", "TRTLLM", "FA4"):
+            raise ValueError(
+                "DFlash attention backend must be VANILLA, TRTLLM or FA4, got "
+                f"{self.dflash_attention_backend!r}."
+            )
         # Each backend loads only its own ops; the rest stay None so the
         # shared paged prologue can read them unconditionally.
+        self._dflash_flash_attention = None
         self._dflash_trtllm_gen_ops = None
         self._dflash_fa4_fwd = None
         self._dflash_paged_append = None
-        if self.dflash_attention_backend == "VANILLA":
+        if not self._uses_worker_attention_backend:
+            # Still validated above so a typo fails here rather than silently,
+            # but no op set is loaded: this drafter calls none of them.
+            logger.info_once(
+                f"{type(self).__name__} brings its own block decode; "
+                f"attention_backend={self.dflash_attention_backend!r} is not used.",
+                key=f"dflash_own_attention_{type(self).__name__}",
+            )
+        elif self.dflash_attention_backend == "VANILLA":
             self._dflash_flash_attention = get_dflash_flash_attention()
         elif self.dflash_attention_backend == "TRTLLM":
             self._dflash_trtllm_gen_ops = get_dflash_trtllm_gen_ops()
-        elif self.dflash_attention_backend == "FA4":
+        else:
             self._dflash_fa4_fwd = get_dflash_fa4_fwd()
             self._dflash_paged_append = get_dflash_paged_append()
-        else:
-            raise ValueError(
-                "DFlash attention backend must be VANILLA, TRTLLM or FA4, got "
-                f"{self.dflash_attention_backend!r}."
-            )
         self._dflash_trtllm_gen_workspace = None
         self._dflash_trtllm_gen_counters = None
         self.register_buffer("_dflash_batch_indices", None, persistent=False)
@@ -727,6 +742,21 @@ def load_weights(self, weights: Dict, weight_mapper=None, **kwargs):
             else:
                 remapped[key] = value
 
+        # Wrapper-owned and built FROM the checkpoint, so they are never in
+        # draft_model_full's module tree and _assert_backbone_complete cannot
+        # see them: a module conjured from data cannot be reported missing by
+        # walking modules. One tuple drives both the extraction below and this
+        # check, so the two cannot drift. Without fc the drafter has no capture
+        # projection, has_target_features stays False, _ctx_len never advances
+        # and it drafts from an empty context forever (the `hasattr` guards on
+        # that path are degradation, not a supported mode).
+        wrapper_missing = [k for k in self.WRAPPER_OWNED_WEIGHTS if k not in remapped]
+        if wrapper_missing:
+            raise ValueError(
+                f"{type(self).__name__}: checkpoint is missing {wrapper_missing}, "
+                "which this wrapper owns and builds from the checkpoint."
+            )
+
         # Load DFlash-specific weights directly
         if "fc.weight" in remapped:
             self.fc = nn.Linear(
@@ -751,9 +781,10 @@ def load_weights(self, weights: Dict, weight_mapper=None, **kwargs):
             self.hidden_norm.weight.data.copy_(remapped["hidden_norm.weight"])
             del remapped["hidden_norm.weight"]
 
-        # Load remaining weights into the draft model.
-        # DFlash checkpoints don't include embed_tokens or lm_head, so allow partial loading
-        # since those modules won't find matching weights.
+        # allow_partial_loading is what lets the shared modules below be absent.
+        # It is also why a truncated checkpoint loads clean, so gate it on the
+        # subclass's own declaration of what may legitimately be missing.
+        self._assert_backbone_complete(remapped, weight_mapper)
         self.draft_model_full.load_weights(
             weights=remapped, weight_mapper=weight_mapper, allow_partial_loading=True
         )
@@ -805,6 +836,78 @@ def take(key: str) -> torch.Tensor:
         self.candidate_selector = selector
         return {k: v for k, v in weights.items() if k not in consumed}
 
+    #: Tensors the wrapper itself owns: not in draft_model_full, built from the
+    #: checkpoint in load_weights, and required.
+    WRAPPER_OWNED_WEIGHTS = ("fc.weight", "hidden_norm.weight")
+
+    #: Parameter-name prefixes this drafter takes from the target instead of its
+    #: own checkpoint. Everything else in ``draft_model_full`` must be provided.
+    #: GQA DFlash checkpoints ship neither embedding nor head; an MLA drafter
+    #: ships its own embedding and overrides this.
+    WEIGHTS_SHARED_WITH_TARGET = ("embed_tokens", "lm_head")
+
+    def _assert_backbone_complete(self, weights: Dict, weight_mapper=None) -> None:
+        """Fail on a checkpoint missing weights the drafter does not share.
+
+        The truth is the constructed module tree, not a hand-kept list that
+        rots as the backbone changes.
+
+        Checked at module granularity. For a PLAIN module that is the hole the
+        flag cannot close: the loader skips one whose subtree filters to
+        nothing (modeling_utils.py `if module_weights:`) whatever the flag
+        says. For a FUSED module `allow_partial_loading=False` would catch it
+        (linear.py asserts all three shards) -- but the flag has to stay True
+        for the target-shared modules, so the check covers that case here
+        instead, and requires every component rather than any.
+
+        Missing parameters INSIDE a present component stay tolerated: a
+        checkpoint with all three weights but only `q_proj.bias` takes the same
+        per-shard copy and leaves the rest at `torch.empty`. Module-granular
+        checking cannot see that and does not pretend to.
+        """
+        provided = set(weights)
+
+        # Whichever fusion table the load below will actually use, not a third
+        # copy: with a mapper modeling_utils dispatches to _load_weights_impl_v2
+        # and the mapper's own table applies; without one it falls back to
+        # _load_weights_impl, whose table is FUSED_MODULE_COMPONENTS. An empty
+        # `mapping` means init_model_and_config has not run, so the mapper has
+        # no table to offer yet and the constant is still the right answer.
+        fusion = dict(getattr(weight_mapper, "mapping", None) or FUSED_MODULE_COMPONENTS)
+
+        def _has(prefix: str) -> bool:
+            return any(k == prefix or k.startswith(prefix + ".") for k in provided)
+
+        def _supplied(module_name: str) -> bool:
+            # A fused module is named once here and stored unfused in the
+            # checkpoint. ALL components must be present, not any: the fused
+            # load path is happy with a subset under allow_partial_loading
+            # (linear.py load_weights_fused_qkv_helper) and leaves the absent
+            # shards at torch.empty -- uninitialised device memory, not zeros.
+            if _has(module_name):
+                return True
+            for fused, parts in fusion.items():
+                if fused in module_name:
+                    return all(_has(module_name.replace(fused, p)) for p in parts)
+            return False
+
+        missing = sorted(
+            {
+                name.rsplit(".", 1)[0]
+                for name, _ in self.draft_model_full.named_parameters()
+                if not any(part in name for part in self.WEIGHTS_SHARED_WITH_TARGET)
+                and not _supplied(name.rsplit(".", 1)[0])
+            }
+        )
+        if missing:
+            raise ValueError(
+                f"{type(self).__name__}: checkpoint provides no weights for "
+                f"{missing[:8]}{' ...' if len(missing) > 8 else ''}. These are "
+                f"not in WEIGHTS_SHARED_WITH_TARGET "
+                f"({', '.join(self.WEIGHTS_SHARED_WITH_TARGET)}), so loading "
+                "would leave them randomly initialized."
+            )
+
     def load_weights_from_target_model(self, target_model: torch.nn.Module) -> None:
         """Share embed_tokens and lm_head from the target model."""
         self.draft_model_full.model.embed_tokens = target_model.model.embed_tokens

tensorrt_llm/_torch/models/modeling_speculative.py

@@ -1405,6 +1405,17 @@ def external_drafter_config_kwargs(model_config, spec_config) -> dict:
         spec_config=None,  # Avoid recursive spec-dec
         max_num_tokens=model_config.max_num_tokens,
         moe_max_num_tokens=model_config.moe_max_num_tokens,
+        # Bounds the drafter's position tables. Without it the field stays
+        # None and they fall back to the checkpoint's advertised
+        # max_position_embeddings -- 1,048,576 for K3, a ~256 MiB complex64
+        # table per rank for a context the runtime bounds far below that.
+        #
+        # The user's value, NOT the engine's. py_executor_creator raises
+        # model_engine_max_seq_len past this and never writes it back, so a
+        # drafter that indexes absolute positions must read the raised value at
+        # runtime rather than have this line predict it -- reproducing that
+        # arithmetic here is what let the two drift apart in the first place.
+        max_seq_len=model_config.max_seq_len,
     )
     # Only the embedded DSpark draft shares the target's EPLB namespace (its
     # stages are target decoder blocks registered into the target's balancer).

tensorrt_llm/_torch/models/modeling_utils.py

@@ -1512,6 +1512,17 @@ def run_concurrently(func,
                 raise
 
 
+#: Modules stored unfused in a checkpoint and fused in the module tree, mapped
+#: to the components they are built from. This is the table the mapper-less
+#: load path below uses; HfWeightMapper.map_weights installs the same pairs for
+#: the mapper path. Anything reasoning about what a checkpoint must provide has
+#: to read one of the two rather than keep a third copy.
+FUSED_MODULE_COMPONENTS = {
+    'qkv_proj': ['q_proj', 'k_proj', 'v_proj'],
+    'gate_up_proj': ['gate_proj', 'up_proj'],
+}
+
+
 def _load_weights_impl(model: Union[nn.Module, DecoderModelForCausalLM],
                        weights: Dict,
                        skip_modules: List[str] = [],
@@ -1538,10 +1549,7 @@ def _load_weights_impl(model: Union[nn.Module, DecoderModelForCausalLM],
         model.config, 'num_key_value_heads'
     ) and model.config.num_key_value_heads is not None else model.config.num_attention_heads
 
-    params_map = {
-        'qkv_proj': ['q_proj', 'k_proj', 'v_proj'],
-        'gate_up_proj': ['gate_proj', 'up_proj']
-    }
+    params_map = dict(FUSED_MODULE_COMPONENTS)
     device_id = local_mpi_rank()
 
     def load_single_module(name, module):