[Spec] Support quantized target lm_head in DFlash2 selector

[LingZ315]

Summary

• Route DFlash2 candidate projection through lm_head.quant_method.apply when the target head is quantized.
• Preserve the existing dense projection and TP-local top-k behavior, including slicing padded vocabulary rows before top-k.
• Replace the unit test that rejected packed heads with a correctness test for quantization-aware candidate generation.

Motivation

DFlash2 borrows the target model's lm_head to build its selector candidates. A ModelOpt NVFP4 target stores that head as packed weights, so treating .weight as a dense FP16/BF16 matrix either fails the dense-head guard or produces a shape mismatch.

The DFlash worker already keeps a quantized selector head on the eager path instead of folding its static matrix multiplication into the draft CUDA graph. This change makes that eager fallback functional by using the same should_apply_lm_head_quant_method dispatch as the target logits path. Dense heads retain the existing sliced matrix multiplication, and TP still communicates only local top-k values and IDs.

This is a quantized-head follow-up to #35371. Related fixes use the same dispatch pattern in #30119 and #34742.

Validation

• pre-commit run --files python/sglang/srt/models/dflash.py test/registered/unit/spec/test_dflash_logits.py
• PYTHONPATH=python python -m pytest -q test/registered/unit/spec/test_dflash_logits.py (4 passed)
• An equivalent runtime override of this compute_candidates path was validated with a ModelOpt NVFP4 Qwen3.8-27B target and Qwen3.8-27B-DFlash2 draft at TP=2. The server completed target and draft CUDA graph capture, kept the quantized selector head eager as intended, served a smoke request, and completed 1000/1000, 50000/400, and 97000/3000 workloads at concurrency 1, 4, and 8.

Behavior note

This PR adds functional quantized-head support to the DFlash2 selector. It does not fold the quantized selector projection into the draft CUDA graph; that remains a separate optimization opportunity.

---

CI States

Latest PR Test (Base): ❌ Run #32231149827
Latest PR Test (Extra): ❌ Run #32231149622

[LingZ315]

@hnyls2002 @SubSir Could you please review this quantized-target follow-up to #35371?

It routes DFlash2 selector candidate projection through the target lm_head.quant_method for ModelOpt NVFP4 heads, preserves the dense and TP-local top-k paths, and intentionally keeps quantized selector projection eager. The focused unit tests and pre-commit checks pass; the equivalent runtime path also completed TP=2 startup and real workloads at concurrency 1, 4, and 8.

[abdulchotu7]

Great fix — this is exactly the crash I hit. One question on the validation target: does the "ModelOpt NVFP4 Qwen3.8-27B" target you tested include activation quantization (W4A4), or is it weights-only NVFP4 (W4A16)?

I ask because I ran the same quant_method.apply eager approach (ported from r0b0tlab/qwen38-27b-nvfp4-sm121-sglang, which validates against W4A16) against the RadixArk/Qwen3.8-27B-NVFP4 checkpoint —which is W4A4 (MLP gate/up/down and lm_head dynamic NVFP4 W4A4, group 16; attention FP8) — at TP=1, SM120/CUDA 13, and the selector still failed after applying the patch. If your target was W4A16, the W4A4 case may not be covered yet; if it was W4A4, I'll dig into why my port diverged. Happy to retest against this branch either way.

[KleinChen42]

No description provided.

[LingZ315]

@abdulchotu7 Thanks for checking. The target I validated was W4A4, not weights-only W4A16. The ModelOpt-exported checkpoint uses group-size-16 NVFP4 for the MLP gate_proj/up_proj/down_proj layers and lm_head, and it contains activation scales: 401 input_scale tensors in total, including lm_head.input_scale. The attention-side quantized layers and KV cache use FP8.

The successful runtime validation used that ModelOpt export at TP=2. Since your failure used the RadixArk checkpoint at TP=1, the difference may be checkpoint metadata/loader behavior or something in the ported patch rather than missing W4A4 coverage. Could you share the exact traceback and SGLang commit/branch? I’d be happy to compare the failing path.

[LingZ315]

Superseded by #35496, which has been merged. Thanks!