feat(dspark): support a pruned draft vocabulary for the DFlash family

[fg11991]

Motivation

EAGLE3 can train a draft over a subset of the target vocabulary. The DFlash family cannot — its contract rejects model.vocab_mapping_path outright — even though the same argument applies more strongly there: on a 27B target, DSpark's Markov head over a full vocabulary is the dominant memory term, and most of that vocabulary never appears as a supervised target in a domain corpus.

This adds pruning to the DFlash family. Set draft_vocab_size below vocab_size in the draft config and the draft predicts over the pruned vocabulary; leave it out and nothing changes.

What's in it

• DraftVocabMappingMixin now holds the t2d/d2t contract shared by EAGLE3 and the DFlash family: registration, installation, validation, and the load_state_dict path. Buffers are registered only when the run actually prunes, so a full-vocabulary checkpoint keeps its exact state dict, and installed state is derived from the persisted buffers rather than a flag — otherwise a reloaded pruned checkpoint reports itself as uninstalled.
• Installing a different mapping of the same size is rejected. Identical shapes with different contents silently repoint every draft id at another token; no shape check catches it.
• Objective. Target ids are mapped into the draft id space; out-of-vocabulary labels leave both the numerator and the denominator of the cross entropy, so alpha_ce keeps its meaning instead of scaling with how much of the batch the vocabulary happens to cover. Acceptance uses the target's true retained probability mass rather than a distribution renormalized over the kept tokens — the renormalized form reports acceptance well above what the pruned draft can deliver.
• Two new ratio metrics make the ceiling visible during training: draft_vocab_coverage (share of supervised tokens the vocabulary can propose) and teacher_kept_mass (teacher belief that survives pruning).
• The unpruned path is unchanged, deliberately down to the arithmetic. The single-denominator expression is kept so an existing run's loss curve does not move by a few ulps for a feature it does not use, and the second all-reduce is emitted only when the run prunes, so the collective sequence of existing runs is untouched.
• keeps_vocab_buffers_when_unpruned lets validation tell the two families apart. A draft that registers the buffers unconditionally can absorb a redundant mapping; one that registers them only when pruning cannot, so that config is rejected at planning time instead of failing later inside the model. When the draft config cannot be resolved at all, _prunes_vocabulary reports None and each rule falls back to what its algorithm needs by default — config validation keeps working without the draft config on disk.

Testing

tests/test_utils/test_dspark_vocab_mapping.py runs the real DSpark draft and online model on CPU with a tiny random config, because what breaks silently here — which factor of the Markov head follows which vocabulary, which id space the labels live in — is invisible to shape-only checks. It covers t2d/d2t round trips, label id spaces, coverage, an analytic acceptance counterexample, checkpoint reload, mapping-conflict rejection, and byte-for-byte equality of the unpruned loss with the pre-change expression.

Full suite run on CPU (no CUDA available on the dev box): no test fails that does not already fail on main at the same commit.

[jiapingW]

Hi. I acknowledge this feature, but does it really deliver a substantial performance boost?

[fg11991]

Thanks for taking a look. The win isn't in the kernel — it's in what the freed memory buys you at training time, plus a small but consistent decode-side gain.

Setup. Target: Qwen3.6-27B (vocab_size 248320). Data: 600k distilled EagleChat samples. Hardware: 64x Ascend 910B3 (8 nodes x 8 cards, 64GB HBM per card). Two drafts trained on it:

	draft vocabulary	seqlen	anchor
baseline	248320 (full)	2048	512
this PR	64000 (pruned)	4096	1024

Why that table is the point. On the DFlash family the Markov head over the full vocabulary is the dominant memory term. Cutting it to ~1/4 of the vocabulary frees enough memory to double both seqlen and anchor size, so within the same 64GB per card the draft gets a configuration it simply could not have been trained at otherwise — that is the real gain, not "same run, less memory".

Quality cost. At the same number of epochs, the pruned draft's acceptance rate on general-domain datasets is 1–2% below the full-vocabulary draft.

Serving. After adapting this to vLLM, both DSpark and MTP at 7 draft steps, 4 concurrent requests: on some domains the pruned draft is faster than the full-vocabulary one by under 1ms, and it is faster than MTP on every domain we measured, by 1–2ms. (We also tried MTP at 3 steps — it was slower than at 7.)