[Bugfix] DFlash: fuse context K/V through quant_method so a quantized draft head loads

tests/kernels/attention/test_dflash2_context_kv_quantized.py

@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Fused context K/V must survive a quantized draft checkpoint.
+
+``_build_context_kv_buffers`` slices the K/V rows out of ``qkv_proj.weight``.
+That attribute only holds the dense ``[N, K]`` matrix for an unquantized
+layer; a quantized draft head keeps packed codes there (NVFP4 packs two
+values per byte, so the tensor is ``[N, K // 2]``), and the fused matrix then
+comes out half as wide -- ``F.linear`` fails with a shape mismatch and the
+engine never finishes its profile run.
+
+These tests use a minimal fake quantization method so they stay CPU-only and
+do not need a checkpoint: what matters is that ``.weight`` is packed and that
+the dense weight is reachable only through ``quant_method.apply``.
+"""
+
+import torch
+from torch import nn
+
+from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+from vllm.model_executor.models.qwen3_dflash import DFlashQwen3Model
+
+HIDDEN = 32
+Q_SIZE = 16
+KV_SIZE = 8
+LAYERS = 3
+
+
+class _FakePackedQuantMethod:
+    """Stores the weight transposed and halved, reachable only via apply()."""
+
+    def apply(self, layer, x, bias=None):
+        return torch.nn.functional.linear(x, layer.dense_weight, bias)
+
+
+def _make_attn(dense_weight: torch.Tensor, quantized: bool) -> nn.Module:
+    qkv = nn.Module()
+    qkv.input_size_per_partition = HIDDEN
+    qkv.dense_weight = dense_weight
+    if quantized:
+        # Packed codes under the same attribute name: half as wide, uint8.
+        qkv.weight = torch.zeros(dense_weight.shape[0], HIDDEN // 2, dtype=torch.uint8)
+        qkv.quant_method = _FakePackedQuantMethod()
+    else:
+        qkv.weight = dense_weight
+        qkv.quant_method = UnquantizedLinearMethod()
+    qkv.bias = None
+    attn = nn.Module()
+    attn.qkv_proj = qkv
+    attn.q_size = Q_SIZE
+    return attn
+
+
+def _make_model(quantized: bool):
+    torch.manual_seed(1234)
+    weights = [
+        torch.randn(Q_SIZE + 2 * KV_SIZE, HIDDEN, dtype=torch.float32)
+        for _ in range(LAYERS)
+    ]
+    model = DFlashQwen3Model.__new__(DFlashQwen3Model)
+    nn.Module.__init__(model)
+    model.hidden_norm = nn.Module()
+    model.hidden_norm.weight = nn.Parameter(torch.ones(HIDDEN), requires_grad=False)
+    layers_attn = [_make_attn(w, quantized) for w in weights]
+    for attn in layers_attn:
+        attn.k_norm = nn.Module()
+        attn.k_norm.weight = nn.Parameter(torch.ones(4), requires_grad=False)
+    return model, layers_attn, weights
+
+
+def _expected(weights: list[torch.Tensor]) -> torch.Tensor:
+    return torch.cat([w[Q_SIZE:] for w in weights], dim=0)
+
+
+def test_unquantized_draft_head_fuses_eagerly():
+    """The dense path must keep fusing at build time, exactly as before."""
+    model, layers_attn, weights = _make_model(quantized=False)
+    model._build_context_kv_buffers(layers_attn, has_bias=False)
+    assert model._fused_kv_weight is not None
+    torch.testing.assert_close(model._fused_kv_weight, _expected(weights))
+
+
+def test_quantized_draft_head_defers_then_fuses_dense():
+    """The packed path defers, then rebuilds the SAME matrix via apply()."""
+    model, layers_attn, weights = _make_model(quantized=True)
+    model._build_context_kv_buffers(layers_attn, has_bias=False)
+    # Deferred: quant_method is not usable until process_weights_after_loading.
+    assert model._fused_kv_weight is None
+
+    model._fuse_dense_kv_weight(torch.float32, torch.device("cpu"))
+    torch.testing.assert_close(model._fused_kv_weight, _expected(weights))
+
+
+def test_quantized_and_unquantized_fusions_agree():
+    """Both paths must produce a bit-identical fused matrix."""
+    dense_model, dense_layers, weights = _make_model(quantized=False)
+    dense_model._build_context_kv_buffers(dense_layers, has_bias=False)
+
+    quant_model, quant_layers, _ = _make_model(quantized=True)
+    quant_model._build_context_kv_buffers(quant_layers, has_bias=False)
+    quant_model._fuse_dense_kv_weight(torch.float32, torch.device("cpu"))
+
+    assert torch.equal(dense_model._fused_kv_weight, quant_model._fused_kv_weight)
+    assert quant_model._fused_kv_weight.shape == (LAYERS * 2 * KV_SIZE, HIDDEN)

vllm/model_executor/models/qwen3_dflash.py

@@ -24,6 +24,7 @@
     QKVParallelLinear,
     ReplicatedLinear,
     RowParallelLinear,
+    UnquantizedLinearMethod,
 )
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
@@ -377,6 +378,17 @@ def forward(
         return hidden_states, residual
 
 
+def _has_dense_qkv_weight(attn: nn.Module) -> bool:
+    """True when ``attn.qkv_proj.weight`` is the plain ``[N, K]`` matrix.
+
+    Quantized checkpoints keep packed codes under the same attribute name, so
+    slicing rows out of it silently yields a differently shaped tensor.
+    """
+    return isinstance(
+        getattr(attn.qkv_proj, "quant_method", None), UnquantizedLinearMethod
+    )
+
+
 @support_torch_compile
 class DFlashQwen3Model(nn.Module):
     decoder_layer_cls = DFlashQwen3DecoderLayer
@@ -494,9 +506,20 @@ def _build_context_kv_buffers(
     ) -> None:
         self._hidden_norm_weight = self.hidden_norm.weight.data
 
-        # KV projection weights: [num_layers * 2 * kv_size, hidden_size]
-        kv_weights = [a.qkv_proj.weight[a.q_size :] for a in layers_attn]
-        self._fused_kv_weight = torch.cat(kv_weights, dim=0)
+        # KV projection weights: [num_layers * 2 * kv_size, hidden_size].
+        # Only a plain nn.Linear keeps that matrix in `.weight`; a quantized
+        # draft checkpoint stores packed codes there instead (NVFP4 packs two
+        # values per byte, so the tensor is [N, K // 2]), and the dense rows
+        # are reachable only through the quantization method. That method is
+        # not usable yet -- this builder runs at the end of load_weights,
+        # before process_weights_after_loading -- so defer the fusion to the
+        # first projection instead. See _fuse_dense_kv_weight.
+        self._context_kv_attn_layers = layers_attn
+        if all(_has_dense_qkv_weight(a) for a in layers_attn):
+            kv_weights = [a.qkv_proj.weight[a.q_size :] for a in layers_attn]
+            self._fused_kv_weight: torch.Tensor | None = torch.cat(kv_weights, dim=0)
+        else:
+            self._fused_kv_weight = None
         if has_bias:
             kv_biases = [a.qkv_proj.bias[a.q_size :] for a in layers_attn]
             self._fused_kv_bias: torch.Tensor | None = torch.cat(kv_biases, dim=0)
@@ -574,6 +597,25 @@ def _build_fused_kv_buffers(self) -> None:
         # References to inner Attention layers for direct cache writes
         self._attn_layers = [layer.self_attn.attn for layer in self.layers]
 
+    def _fuse_dense_kv_weight(self, dtype: torch.dtype, device: torch.device) -> None:
+        """Build the fused context K/V matrix for a quantized draft head.
+
+        Feeding an identity matrix through the layer's own quantization
+        method returns the dense transposed weight, whatever the packing is,
+        so this needs no knowledge of NVFP4/FP8/marlin layouts. Runs once, on
+        the first context projection, when process_weights_after_loading has
+        completed. The per-step decode path keeps using the quantized
+        weights; only this prefill-time fusion is materialized dense.
+        """
+        rows = []
+        for attn in self._context_kv_attn_layers:
+            qkv = attn.qkv_proj
+            eye = torch.eye(qkv.input_size_per_partition, dtype=dtype, device=device)
+            weight_t = qkv.quant_method.apply(qkv, eye, None)
+            rows.append(weight_t.t()[attn.q_size :].contiguous())
+            del eye, weight_t
+        self._fused_kv_weight = torch.cat(rows, dim=0)
+
     def _project_context_kv(
         self,
         context_states: torch.Tensor,
@@ -582,6 +624,8 @@ def _project_context_kv(
         num_kv_heads: int,
         head_dim: int,
     ) -> tuple[torch.Tensor, torch.Tensor]:
+        if self._fused_kv_weight is None:
+            self._fuse_dense_kv_weight(context_states.dtype, context_states.device)
         # --- Fused KV projection (one GEMM for all layers) ---
         normed_context_states = self._normalize_context_states(context_states)
         all_kv_flat = F.linear(