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

configs/qwen3-8b-dspark-draftvocab32k.json

@@ -0,0 +1,59 @@
+{
+  "architectures": [
+    "DSparkDraftModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoModel": "dspark.DSparkDraftModel"
+  },
+  "block_size": 7,
+  "bos_token_id": 151643,
+  "dflash_config": {
+    "attention_mode": "gqa",
+    "confidence_head_alpha": 1.0,
+    "confidence_head_with_markov": true,
+    "enable_confidence_head": true,
+    "markov_head_type": "vanilla",
+    "markov_rank": 256,
+    "mask_token_id": 151669,
+    "projector_type": "dspark",
+    "target_layer_ids": [
+      1,
+      9,
+      17,
+      25,
+      33
+    ]
+  },
+  "dtype": "bfloat16",
+  "draft_vocab_size": 32000,
+  "eos_token_id": 151645,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 12288,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 40960,
+  "max_window_layers": 5,
+  "model_type": "qwen3",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 5,
+  "num_key_value_heads": 8,
+  "num_target_layers": 36,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000,
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

docs/basic_usage/training.md

@@ -288,6 +288,53 @@ assembly. In particular:
 
 There is no fallback to a removed training script.
 
+## Draft vocabulary pruning
+
+A draft can predict over a subset of the target vocabulary. This is worth doing
+when the target vocabulary is large relative to the draft: the output head and,
+for DSpark, the Markov head scale with it, while a domain corpus supervises only
+a fraction of the ids.
+
+It is off by default and enabled entirely by the draft config — add
+`draft_vocab_size` below `vocab_size` and change nothing in the run config:
+
+```json
+{
+  "architectures": ["DSparkDraftModel"],
+  "vocab_size": 151936,
+  "draft_vocab_size": 32000
+}
+```
+
+`configs/qwen3-8b-dspark-draftvocab32k.json` and
+`examples/configs/qwen3-8b-dspark-draftvocab32k-offline.yaml` are a complete
+worked pair.
+
+The mapping itself (`t2d`/`d2t`) comes from one of two places:
+
+- **Local offline runs** derive and cache it from the feature corpus when
+  `model.vocab_mapping_path` is empty, the same way EAGLE3 does.
+- **Disaggregated runs** require an explicit `model.vocab_mapping_path`:
+  producer and consumer would otherwise derive different mappings from
+  different shards.
+
+Two constraints are worth knowing before you start a long run:
+
+- A checkpoint that was trained with a pruned vocabulary records its
+  `draft_vocab_size`, and resuming it against a mapping with different contents
+  is rejected — same-sized mappings that select different tokens would silently
+  repoint every draft id at another token. Keep the mapping file next to the
+  checkpoint.
+- `draft_vocab_coverage` and, when the target's hidden states are available,
+  `teacher_kept_mass` are logged during training. They are the ceiling on
+  acceptance: coverage is how often the realized next token is even proposable,
+  and kept mass is how much of the teacher's belief survives pruning. Check them
+  on the first few hundred steps rather than at evaluation time.
+
+A run whose `draft_vocab_size` equals `vocab_size` is not pruned: the draft
+registers no mapping buffers, so `model.vocab_mapping_path` is rejected at
+config validation rather than failing later inside the model.
+
 Step limits are global optimizer updates. `training.max_steps` is a stop cap and,
 when set without `training.total_steps`, the fallback optimizer/loss schedule
 horizon. `training.total_steps` can describe a longer schedule, but does not by

examples/configs/qwen3-8b-dspark-draftvocab32k-offline.yaml

@@ -0,0 +1,42 @@
+# DSpark with a frequency-pruned draft vocabulary (32000 of 151936).
+#
+# Offline/colocated on purpose: the mapping is derived from the run's own
+# features by _ensure_offline_vocab_mapping, so no vocab_mapping_path is needed.
+# A disaggregated pruned run must instead pass model.vocab_mapping_path, because
+# producer and consumer cannot each derive the same mapping independently.
+model:
+  target_model_path: Qwen/Qwen3-8B
+  draft_model_config: configs/qwen3-8b-dspark-draftvocab32k.json
+  target_backend: sglang
+  embedding_key: model.embed_tokens.weight
+  torch_dtype: bfloat16
+
+data:
+  hidden_states_path: ./cache/hidden_states/qwen3-8b-dspark-sharegpt
+  max_length: 3072
+  chat_template: qwen
+  cache_dir: ./cache
+
+training:
+  strategy: dspark
+  num_epochs: 6
+  max_steps: 10000
+  batch_size: 1
+  learning_rate: 6.0e-4
+  warmup_ratio: 0.04
+  max_grad_norm: 1.0
+  attention_backend: flex_attention
+  num_anchors: 512
+  save_interval: 1000
+  log_interval: 50
+  dist_timeout: 30
+  seed: 42
+
+run_id: qwen3-8b-dspark-draftvocab32k-offline
+output_dir: ./outputs/qwen3-8b-dspark-draftvocab32k-offline
+
+deployment:
+  mode: local_colocated
+  trainer:
+    nnodes: 1
+    nproc_per_node: 1