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

[Peuqui]

Purpose

DFlash never finishes its profile run when the draft checkpoint is quantized.
DFlashQwen3Model._build_context_kv_buffers fuses the per-layer K/V
projections into one matrix by slicing rows out of qkv_proj.weight:

kv_weights = [a.qkv_proj.weight[a.q_size :] for a in layers_attn]

That attribute holds the dense [N, K] matrix only for an unquantized layer.
A quantized checkpoint keeps packed codes there — NVFP4 packs two values per
byte, so the tensor is [N, K // 2] — and the fused matrix comes out half as
wide. _project_context_kv then fails:

RuntimeError: mat1 and mat2 shapes cannot be multiplied (2048x5120 and 2560x5120)

The abort happens in determine_available_memory → profile_run →
speculator.propose, before a single token is served, so DFlash is currently
unusable with any quantized draft head on any device class. Reproduced on
Qwen3.8-27B with an NVFP4 draft head derived from
incoai/Qwen3.8-27B-DFlash2; the line is unchanged on main.

Fix

Reach the dense rows through the layer's own quantization method: feeding an
identity matrix through quant_method.apply returns the transposed weight
whatever the packing is, so this needs no knowledge of NVFP4, FP8 or marlin
layouts. The unquantized path is untouched and still fuses eagerly.

The rebuild has to be deferred. _build_fused_kv_buffers runs at the end of
load_weights, before process_weights_after_loading, so quant_method is
not usable yet; the fusion is built on the first context projection instead.
It costs one dense fp16 copy of the K/V rows (~52 MiB per rank on
Qwen3.8-27B) and one GEMM per layer, once. The per-step decode path keeps
using the quantized weights, which is where the win is — the context
projection is prefill work.

Result

Qwen3.8-27B, TP2, DFlash2 k=7, greedy, 5 runs, median, fp16 draft head
against the NVFP4 one:

GPUs	fp16 draft	NVFP4 draft
2x Quadro RTX 8000	72.72	77.13 tok/s	+6.1 %
2x Tesla V100-PCIE	74.02	76.33 tok/s	+3.1 %

The answer text is byte-identical in every run (sha256 prefix
0106659946c064b1); acceptance length moves 3.353 → 3.325. That is expected
and is the safety argument for quantizing a draft head at all: only tokens
the target model would have produced anyway are accepted, so quantization can
cost acceptance rate but never correctness.

Not a duplicate

Checked on 2026-09-10 against 1CatAI/1Cat-vLLM:

gh pr list --repo 1CatAI/1Cat-vLLM --state open --search "dflash quantized draft"
gh pr list --repo 1CatAI/1Cat-vLLM --state open --search "dflash2 nvfp4"
gh pr list --repo 1CatAI/1Cat-vLLM --state open --search "context_kv"
gh pr list --repo 1CatAI/1Cat-vLLM --state open --search "_build_context_kv_buffers"
gh issue list --repo 1CatAI/1Cat-vLLM --state all --search "dflash draft quantized"
gh issue list --repo 1CatAI/1Cat-vLLM --state open --search "dflash"

No open PR or issue covers this crash. The nearest open PR, #561 (Reduce
DFlash2 weight and scale memory on SM70, draft), shares NVFP4 codes between
QPN2 and TurboMind in the target model and does not touch
qwen3_dflash.py. The nearest issue, #478 (DFlash2 acceptance=0 on V100
TP4), is an acceptance problem on a running server, not a load failure.

Test Plan

New file tests/kernels/attention/test_dflash2_context_kv_quantized.py,
CPU-only, using a minimal fake quantization method so no checkpoint is
needed: .weight is packed and the dense weight is reachable only through
quant_method.apply.

Commands (full outputs under Test Result):

python -m pytest tests/kernels/attention/test_dflash2_context_kv_quantized.py \
                 tests/kernels/attention/test_dflash2_context_pipeline.py -v
python -m pytest tests/v1/spec_decode/test_dflash2_alignment_rank.py \
                 tests/v1/spec_decode/test_dflash2_ngram_assist.py \
                 tests/v1/spec_decode/test_dflash2_structured_output.py -q
pre-commit run --files vllm/model_executor/models/qwen3_dflash.py \
                       tests/kernels/attention/test_dflash2_context_kv_quantized.py
pre-commit run mypy-3.10 --hook-stage manual --files <same files>

Test Result

$ python -m pytest tests/kernels/attention/test_dflash2_context_kv_quantized.py \
                   tests/kernels/attention/test_dflash2_context_pipeline.py -v
4 passed, 15 warnings in 2.30s

$ python -m pytest tests/v1/spec_decode/test_dflash2_alignment_rank.py \
                   tests/v1/spec_decode/test_dflash2_ngram_assist.py \
                   tests/v1/spec_decode/test_dflash2_structured_output.py -q
38 passed, 19 warnings in 28.68s

The complete existing test_dflash2_context_pipeline.py was run, not just
the new tests. With the fix reverted, the two quantized tests fail and the
unquantized one still passes, so they do catch this bug and do not constrain
the dense path.

$ pre-commit run --files vllm/model_executor/models/qwen3_dflash.py \
                         tests/kernels/attention/test_dflash2_context_kv_quantized.py
ruff check .. Passed / ruff format .. Passed / typos .. Passed / mypy .. Passed

$ pre-commit run mypy-3.10 --hook-stage manual --files <same files>
Run mypy for Python 3.10 .. Passed

AI assistance

This change was prepared with AI assistance (Claude). I reviewed every
changed line, ran the tests and linters shown above on my own hardware, and
can defend the change end-to-end.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.