diff --git a/tests/kernels/test_dsv41_mega_attn_layouts.py b/tests/kernels/test_dsv41_mega_attn_layouts.py index 6f9ff56bd236..3d454183605d 100644 --- a/tests/kernels/test_dsv41_mega_attn_layouts.py +++ b/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 + ) diff --git a/vllm/models/deepseek_v41/common/ops/fused_layout.py b/vllm/models/deepseek_v41/common/ops/fused_layout.py index 071afde6bca4..50df8de3514a 100644 --- a/vllm/models/deepseek_v41/common/ops/fused_layout.py +++ b/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, + ) diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py index d1534d52483f..a09c3159c235 100644 --- a/vllm/models/deepseek_v41/nvidia/dspark.py +++ b/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. diff --git a/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py b/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py index 8efc987739b1..7a7f4f8c9016 100644 --- a/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py +++ b/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. diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index bd91e8f12b4f..28fa0a5915a8 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/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()