[Spec] Build DFlash draft layers under their checkpoint names and refuse tensors the draft would drop

[twu3202]

Motivation

A quantized DFlash draft is built from its quantization config by module name and loaded by tensor name, and the two did not see the same names. DFlashDraftModel gave every decoder layer an empty prefix unless the checkpoint was a Nemotron 3.5 draft, so quantization matching saw a bare qkv_proj, o_proj, gate_up_proj, down_proj for all layers. Name-based ignore entries such as re:.*self_attn\.q_proj$ never applied, the attention projections were built as quantized layers, and load_weights then dropped the checkpoint's BF16 q_proj/k_proj/v_proj tensors without a word because the fused module had no .weight to put them in. The packed weight went into attention uninitialized. With the W8A16 draft from the recipe below, built from the official z-lab/Qwen3.8-27B-DFlash2, main hits a device-side assert (srcIndex < srcSelectDimSize in indexSelectSmallIndex) on the server's own warm-up request, after a startup that only logged DFLASH fused KV materialization disabled: quantized qkv_proj is not supported for this path.

The second half of the same problem is fc: outside the Nemotron path it was a plain nn.Linear, so a checkpoint that quantizes fc (fc.weight_packed / weight_scale / weight_shape, or FP8 fc.weight + weight_scale_inv) had those tensors dropped the same way and fc kept its random init, so the draft would propose from a random projection of the target's hidden states. Several of the quantized DFlash2 drafts published on the Hub for Qwen3.8-27B quantize fc as well.

Modifications

python/sglang/srt/models/dflash.py only:

• A draft whose own checkpoint is quantized, meaning its config.json carries a quantization_config, builds every decoder layer with layers.{i} (or {prefix}.layers.{i}), as Nemotron drafts and drafts whose class sets supports_quantization (LiLiCorr, [Spec] LiLiCorr: named head MLPs and quantized head linears #42057) already do. A quantization config then matches the same names the checkpoint uses.
• packed_modules_mapping on the draft model class, which the loader already injects into the quant config, so an ignore list that names q_proj/k_proj/v_proj or gate_proj/up_proj applies to the fused module. llm-compressor writes the shard names it was given, as in the recipe below. As for every other model with such a mapping, an ignore list that names a single shard now fails with the fused-layer scheme mismatch error instead of matching nothing.
• The same drafts build fc as a ReplicatedLinear with the quant config, as those drafts already do. Any other draft is built as on main, with empty layer prefixes and a dense nn.Linear fc. That includes an unquantized draft that inherits the target's --quantization: with --quantization fp8 its other linears are already quantized online on main, and a quantizable fc would have become FP8 as well. The check reads config.json only, so a draft described only by hf_quant_config.json is also built as on main, and a quantized fc tensor in it is now refused, as described below, instead of being dropped.
• The fc.weight shape check keeps the packed-aware form and encoder.fc.* vendor names alias to fc.*. load_weights refuses a checkpoint that fills the decoder layers but never fc: a quantizable fc starts uninitialized, and a dense one left at its random init gave unusable drafts before as well. A projector-only load, as Domino does, is unaffected.
• load_weights refuses a checkpoint tensor whose module the draft has, under a parameter name that module does not have (a dense tensor for a packed layer, or a packed tensor for a dense layer, under the draft's own or the vendor encoder. name), with the module's actual parameter names in the message. Tensors for modules the draft does not have are still ignored, and so are the ones that have no draft parameter by design: .bias, the .inv_freq / .cos_cached / .sin_cached rotary caches, and the .input_scale / .output_scale / .k_scale / .v_scale / .kv_scale activation and KV-cache scales.
• New CPU test test/registered/unit/spec/test_dflash_quantized_draft_layout.py: the layer names through the real decoder layer, the shard-only ignore list reaching the fused projection, fc quantized / ignored / dense, an unquantized draft under an inherited online FP8 config keeping main's names and dense fc, packed and dense fc loading (own and vendor name), the two refusals, and the tensors that must still be ignored. On main 8 of the 13 cases fail.

Accuracy Tests

Unit tests (CPU): test_dflash_quantized_draft_layout.py 13 passed (8 fail on main); test_dflash_logits.py, test_dflash_domino.py, test_oot_dflash_hooks.py, test_dflash_extra_buffer_lazy.py, test_dflash_overlap_hostsync.py, model_runner_components/test_spec_aux_hidden_state.py 51 passed. The branch is rebased onto 26a1377, after #42057, and these CPU tests were re-run there.

The quantized draft is built from the official z-lab/Qwen3.8-27B-DFlash2 with llm-compressor's model-free PTQ: INT8 weight-only for the MLP, o_proj and fc, with q/k/v, the candidate selector and the convolution kernels left BF16:

# llmcompressor==0.14.0, data-free round-to-nearest, runs on CPU
from llmcompressor import model_free_ptq

model_free_ptq(
    model_stub="z-lab/Qwen3.8-27B-DFlash2",
    save_directory="Qwen3.8-27B-DFlash2-W8A16",
    scheme="W8A16",
    ignore=["re:.*self_attn\\.(q|k|v)_proj$", "re:.*candidate_selector.*",
            "re:.*kernel_projection$", "re:.*norm.*"],
)

The target is Qwen/Qwen3.8-27B with a W4A16 body from the same tool and its lm_head left BF16, so that this change is tested on its own.

Target recipe

# llmcompressor==0.14.0, data-free round-to-nearest, runs on CPU
from compressed_tensors.quantization import QuantizationConfig, QuantizationScheme
from compressed_tensors.quantization.quant_scheme import W4A16
from llmcompressor import model_free_ptq

model_free_ptq(
    model_stub="Qwen/Qwen3.8-27B",
    save_directory="Qwen3.8-27B-W4A16",
    config=QuantizationConfig(config_groups={
        "body": QuantizationScheme(targets=["Linear"], **W4A16),
    }),
    ignore=["lm_head", "re:.*visual.*", "re:.*mtp.*", "re:.*in_proj_(a|b)$", "re:.*conv1d.*",
            "re:.*norm.*", "re:.*embed_tokens.*"],
)

One RTX 6000 Ada (48 GB, TP=1), --speculative-algorithm DFLASH --speculative-dflash-block-size 8 --disable-overlap-schedule. Accuracy is SGLang's own GSM8K eval, one request at a time so completions can be compared across runs:

python3 -m sglang.test.run_eval --eval-name gsm8k --num-examples 200 --num-threads 1 --temperature 0 --max-tokens 512

draft	branch	result	GSM8K (200)	accept length
none	main	serves	0.960
z-lab/Qwen3.8-27B-DFlash2 (BF16)	main	serves	0.970	6.035
z-lab/Qwen3.8-27B-DFlash2 (BF16)	this PR	serves	0.970	6.035
W8A16 from the recipe	main	device-side assert on the server's warm-up request
W8A16 from the recipe	this PR	serves, fused KV materialization back on	0.970	6.031

Accept length is the mean of the scheduler's accept len over every decode-log window of the run, bonus token included.

With the BF16 draft, this PR's 200 completions are identical to main's. The W8A16 draft's completions match the BF16 draft's on main in 196/200. Against the target decoding on its own, the W8A16 draft matches 165/200 and the BF16 draft 166/200.

An unquantized draft that inherits --quantization is built exactly as on main. To check that, the target is Qwen/Qwen3.8-27B-FP8 served with --quantization fp8, which the BF16 draft inherits as online FP8. An earlier revision of this PR made that draft's fc FP8 as well:

target	draft	branch	GSM8K (200)	accept length	completions identical to main
Qwen/Qwen3.8-27B-FP8, --quantization fp8	z-lab/Qwen3.8-27B-DFlash2 (BF16)	main	0.975	5.946
Qwen/Qwen3.8-27B-FP8, --quantization fp8	z-lab/Qwen3.8-27B-DFlash2 (BF16)	earlier revision, fc in FP8	0.975	5.957	168/200
Qwen/Qwen3.8-27B-FP8, --quantization fp8	z-lab/Qwen3.8-27B-DFlash2 (BF16)	this PR	0.975	5.946	200/200

All GPU runs are on 525f140. With #40883 applied as well (measured on the earlier revision, whose quantized-draft path is the same), the same W8A16 draft also serves the target from that PR, whose lm_head is W8A16 pack-quantized: GSM8K 0.965, accept length 6.025, and 165/200 completions matching that target decoding on its own.

Speed Tests and Profiling

No change for unquantized drafts: they are built as on main. The quantized drafts above go from not serving to serving; I have no clean timing for them, since the GPU these runs used was shared with other jobs.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #37295821876
Latest PR Test (Extra): ❌ Run #37295821783
Latest PR Test (AMD ROCm 10): ❌ Run #37295821835

[twu3202]

@mmangkad could you take a look when you have a moment? Quantized DFlash drafts build their layers without the checkpoint prefixes, so the quant config's ignore list never matches and the BF16 q/k/v weights get dropped at load, ending in a device-side assert. Repro and GSM8K numbers are in the description. It would also need run-ci.