[Model][DeepSeek-V4.1] Permute mega-attention wq_b / wo_a while loading

tests/kernels/test_dsv41_mega_attn_layouts.py

@@ -3,25 +3,34 @@
 """Layout plumbing for DeepSeek V4.1 mega attention.
 
 The kernel's Q and O layouts are produced by permuting ``wq_b`` rows and
-``wo_a`` columns once at load, so what has to hold is that the permuted GEMMs
-agree with a reference that permutes the activation instead.
+``wo_a`` columns in their weight loaders, so what has to hold is that the
+permuted GEMMs agree with a reference that permutes the activation instead.
 """
 
 import torch
 
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.models.deepseek_v41.common.ops.fused_layout import (
     WV_GROUP_SIZE,
     o_fused_chunk_permutation,
     o_fused_permutation,
+    permute_on_load,
     permute_q_to_fused,
-    permute_wo_a_,
-    permute_wq_b_,
     q_fused_permutation,
 )
 
 HEAD_DIM = 512
 
 
+def _load(checkpoint: torch.Tensor, perm: torch.Tensor, dim: int) -> torch.Tensor:
+    """Load ``checkpoint`` into a fresh parameter through a permuting loader."""
+    param = torch.nn.Parameter(torch.empty_like(checkpoint), requires_grad=False)
+    param.weight_loader = default_weight_loader  # type: ignore[attr-defined]
+    permute_on_load(param, perm, dim)
+    param.weight_loader(param, checkpoint)  # type: ignore[attr-defined]
+    return param.data
+
+
 def test_permuted_wq_b_gemm_matches_permuted_activation():
     """Permuting wq_b's rows makes the Q GEMM emit the fused layout directly."""
     torch.manual_seed(0)
@@ -32,18 +41,17 @@ def test_permuted_wq_b_gemm_matches_permuted_activation():
         0, 256, (num_heads * HEAD_DIM, q_lora_rank // 32), dtype=torch.uint8
     )
     qr = torch.randn(num_tokens, q_lora_rank)
-    orig_weight, orig_scale = weight.clone(), scale.clone()
 
     standard = (qr @ weight.T).view(num_tokens, num_heads, HEAD_DIM)
     expected = permute_q_to_fused(standard)
 
-    permute_wq_b_(weight, scale, num_heads)
-    fused = (qr @ weight.T).view(num_tokens, num_heads, HEAD_DIM)
+    perm = q_fused_permutation(num_heads, HEAD_DIM)
+    loaded_weight = _load(weight, perm, dim=0)
+    fused = (qr @ loaded_weight.T).view(num_tokens, num_heads, HEAD_DIM)
     torch.testing.assert_close(fused, expected)
     # The scale has to follow its row, or dequant reads another row's scale.
-    perm = q_fused_permutation(num_heads, HEAD_DIM)
-    torch.testing.assert_close(weight, orig_weight[perm])
-    torch.testing.assert_close(scale, orig_scale[perm])
+    torch.testing.assert_close(loaded_weight, weight[perm])
+    torch.testing.assert_close(_load(scale, perm, dim=0), scale[perm])
 
 
 def test_permuted_wo_a_consumes_fused_output():
@@ -54,18 +62,18 @@ def test_permuted_wo_a_consumes_fused_output():
     weight = torch.randn(out_features, in_features)
     scale = torch.randint(0, 256, (out_features, in_features // 32), dtype=torch.uint8)
     standard_o = torch.randn(num_tokens, in_features)
-    orig_weight, orig_scale = weight.clone(), scale.clone()
 
     expected = standard_o @ weight.T
     # The kernel emits the same values in the fused chunk order.
     perm = o_fused_permutation(WV_GROUP_SIZE, HEAD_DIM)
     fused_o = standard_o[:, perm]
 
-    permute_wo_a_(weight, scale, WV_GROUP_SIZE)
-    torch.testing.assert_close(weight, orig_weight[:, perm])
-    torch.testing.assert_close(
-        scale, orig_scale[:, o_fused_chunk_permutation(WV_GROUP_SIZE, HEAD_DIM)]
-    )
+    loaded_weight = _load(weight, perm, dim=1)
+    torch.testing.assert_close(loaded_weight, weight[:, perm])
+    chunk_perm = o_fused_chunk_permutation(WV_GROUP_SIZE, HEAD_DIM)
+    torch.testing.assert_close(_load(scale, chunk_perm, dim=1), scale[:, chunk_perm])
     # Unlike wq_b, this permutation sits inside the 4096-term reduction, so the
     # summation order changes and the result differs in the last few ulps.
-    torch.testing.assert_close(fused_o @ weight.T, expected, rtol=1e-3, atol=1e-3)
+    torch.testing.assert_close(
+        fused_o @ loaded_weight.T, expected, rtol=1e-3, atol=1e-3
+    )

vllm/models/deepseek_v41/common/ops/fused_layout.py

@@ -12,6 +12,8 @@
 
 import torch
 
+from vllm.model_executor.model_loader.weight_utils import composed_weight_loader
+
 HEAD_DIM = 512
 Q_CHUNK = 16
 O_CHUNK = 32
@@ -60,27 +62,13 @@ def _bytes_view(t: torch.Tensor) -> torch.Tensor:
     return t.view(torch.uint8) if t.element_size() == 1 else t
 
 
-def permute_wq_b_(
-    weight: torch.Tensor, weight_scale: torch.Tensor, num_local_heads: int
-) -> None:
-    """Permute an MXFP8 ``wq_b`` shard's rows and per-row scales, in place."""
-    head_dim = weight.shape[0] // num_local_heads
-    perm = q_fused_permutation(num_local_heads, head_dim).to(weight.device)
-    for t in (weight, weight_scale):
-        b = _bytes_view(t)
-        b.copy_(b[perm])
-
-
-def permute_wo_a_(
-    weight: torch.Tensor,
-    weight_scale: torch.Tensor,
-    heads_per_group: int = WV_GROUP_SIZE,
-) -> None:
-    """Permute an MXFP8 ``wo_a`` shard's input columns and scales, in place."""
-    head_dim = weight.shape[1] // heads_per_group
-    perm = o_fused_permutation(heads_per_group, head_dim).to(weight.device)
-    w = _bytes_view(weight)
-    w.copy_(w[:, perm])
-    chunk_perm = o_fused_chunk_permutation(heads_per_group, head_dim)
-    s = _bytes_view(weight_scale)
-    s.copy_(s[:, chunk_perm.to(weight.device)])
+def permute_on_load(param: torch.Tensor, perm: torch.Tensor, dim: int) -> None:
+    """Make ``param``'s weight loader gather the loaded shard by ``perm``."""
+
+    def gather(t: torch.Tensor) -> torch.Tensor:
+        return _bytes_view(t).index_select(dim, perm.to(t.device)).view(t.dtype)
+
+    param.weight_loader = composed_weight_loader(  # type: ignore[attr-defined]
+        param.weight_loader,  # type: ignore[attr-defined]
+        gather,
+    )

vllm/models/deepseek_v41/nvidia/dspark.py

@@ -500,22 +500,8 @@ def _finalize_moe(self) -> None:
         for layer in self.model.layers:
             layer.ffn.finalize_mega_moe_weights()
 
-    def _finalize_attn(self) -> None:
-        """Run the attention backend's post-load weight step on the draft layers.
-
-        They are ordinary DeepseekV4DecoderLayers, so they take the engine-wide
-        attention backend and owe it whatever the target model owes it -- mega
-        attention permutes wq_b / wo_a here and refuses to run without it.
-        Idempotent, like the target's.
-        """
-        for layer in self.model.layers:
-            finalize = getattr(layer.attn, "finalize_loaded_weights", None)
-            if finalize is not None:
-                finalize()
-
     def process_weights_after_loading(self) -> None:
         self._finalize_moe()
-        self._finalize_attn()
 
     def _remap_dspark_name(self, name: str) -> str | None:
         """Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path.

vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py

@@ -32,8 +32,10 @@
 )
 from vllm.models.deepseek_v41.common.ops.fused_layout import (
     WV_GROUP_SIZE,
-    permute_wo_a_,
-    permute_wq_b_,
+    o_fused_chunk_permutation,
+    o_fused_permutation,
+    permute_on_load,
+    q_fused_permutation,
 )
 from vllm.models.deepseek_v41.nvidia.flashmla import DeepseekV4FlashMLAAttention
 from vllm.models.deepseek_v41.sparse_mla import (
@@ -182,7 +184,13 @@ def __init__(self, vllm_config: VllmConfig, *args, **kwargs) -> None:
                 "wo_a group."
             )
         self.n_wv_group = self.padded_heads // WV_GROUP_SIZE
-        self._fused_layouts_ready = False
+        q_perm = q_fused_permutation(self.n_local_heads, self.head_dim)
+        permute_on_load(self.wq_b.weight, q_perm, dim=0)
+        permute_on_load(self.wq_b.weight_scale, q_perm, dim=0)
+        o_perm = o_fused_permutation(WV_GROUP_SIZE, self.head_dim)
+        permute_on_load(self.wo_a.weight, o_perm, dim=1)
+        o_chunk_perm = o_fused_chunk_permutation(WV_GROUP_SIZE, self.head_dim)
+        permute_on_load(self.wo_a.weight_scale, o_chunk_perm, dim=1)
 
     # ---- interface contract ------------------------------------------------
 
@@ -227,25 +235,6 @@ def _o_proj(
         )
         return self._wo_b_proj(z.flatten(1))
 
-    # ---- weights -----------------------------------------------------------
-
-    def finalize_loaded_weights(self) -> None:
-        """Permute wq_b rows / wo_a columns into the kernel's layouts.
-
-        Idempotent: a second post-load pass must not permute twice.
-        """
-        if self._fused_layouts_ready:
-            return
-        permute_wq_b_(
-            self.wq_b.weight.data, self.wq_b.weight_scale.data, self.n_local_heads
-        )
-        permute_wo_a_(
-            self.wo_a.weight.data,
-            self.wo_a.weight_scale.data,
-            self.n_local_heads // self.n_local_groups,
-        )
-        self._fused_layouts_ready = True
-
     # ---- forward -----------------------------------------------------------
 
     def forward_mqa(
@@ -255,11 +244,6 @@ def forward_mqa(
         positions: torch.Tensor,
         output: QuantizedActivation,
     ) -> None:
-        if not self._fused_layouts_ready:
-            raise RuntimeError(
-                f"{self.prefix}: wq_b / wo_a were never permuted for the mega "
-                "attention kernel; refusing to run with mismatched layouts."
-            )
         attn_metadata = get_forward_context().attn_metadata
         if attn_metadata is None:
             # Warmup dummy run: reserve the prefill workspace, produce zeros.

vllm/models/deepseek_v41/nvidia/model.py

@@ -1354,17 +1354,6 @@ def finalize_mega_moe_weights(self) -> None:
         for layer in islice(self.layers, self.start_layer, self.end_layer):
             layer.ffn.finalize_mega_moe_weights()
 
-    def finalize_mega_attn_weights(self) -> None:
-        """Permute wq_b / wo_a into FlashMLA's mega-attention layouts.
-
-        A no-op for every other attention layer, and idempotent, so a second
-        post-load pass cannot permute twice.
-        """
-        for layer in islice(self.layers, self.start_layer, self.end_layer):
-            finalize = getattr(layer.attn, "finalize_loaded_weights", None)
-            if finalize is not None:
-                finalize()
-
     def finalize_mhc_broadcast_weights(self) -> None:
         if not get_pp_group().is_first_rank or self.start_layer >= self.end_layer:
             return
@@ -1637,7 +1626,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
     def process_weights_after_loading(self) -> None:
         self.model.finalize_mega_moe_weights()
         self.model.finalize_mhc_broadcast_weights()
-        self.model.finalize_mega_attn_weights()
 
     def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
         return self.model.get_expert_mapping()