diff --git a/tests/kernels/attention/test_dflash2_context_kv_quantized.py b/tests/kernels/attention/test_dflash2_context_kv_quantized.py new file mode 100644 index 0000000000..361821fa42 --- /dev/null +++ b/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) diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index c1b594460c..29a6f4d398 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/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(