MoE drafts: DeepSeek-V4.1's own DSpark draft (stacked on #1035)

.github/CODEOWNERS

@@ -1,7 +1,7 @@
 .github @FrankLeeeee
 /specforge/core @FrankLeeeee
 /specforge/data @zyksir @sleepcoo @shuaills
-/specforge/layers @FrankLeeeee @FlamingoPg @sleepcoo @shuaills
+/specforge/modeling/draft/layers @FrankLeeeee @FlamingoPg @sleepcoo @shuaills
 /specforge/modeling @FlamingoPg @sleepcoo @shuaills @FrankLeeeee
 /tests @FrankLeeeee
 /assets @FrankLeeeee @zhyncs

configs/deepseek-v4.1-flash-moe-dspark.json

@@ -0,0 +1,68 @@
+{
+  "architectures": ["DSparkDraftModel"],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {"AutoModel": "dspark.DSparkDraftModel"},
+  "block_size": 5,
+  "bos_token_id": 0,
+  "dflash_config": {
+    "attention_mode": "deepseek_v41",
+    "confidence_head_alpha": 1.0,
+    "confidence_head_bias": false,
+    "confidence_head_input": "pre_norm",
+    "confidence_head_with_markov": true,
+    "context_window": 128,
+    "enable_confidence_head": true,
+    "markov_head_type": "vanilla",
+    "markov_rank": 256,
+    "mask_token_id": 128799,
+    "moe_router_aux_loss_coef": 0.0,
+    "moe_router_bias_update_rate": 0.0,
+    "moe_router_enable_expert_bias": true,
+    "projector_type": "dspark",
+    "target_layer_ids": [37, 38, 39]
+  },
+  "draft_residual": {
+    "kind": "mhc",
+    "num_streams": 4,
+    "sinkhorn_eps": 1e-06,
+    "sinkhorn_iters": 20
+  },
+  "dtype": "bfloat16",
+  "eos_token_id": 1,
+  "ffn_layer_types": ["moe", "moe", "moe"],
+  "head_dim": 512,
+  "hidden_act": "silu",
+  "hidden_size": 5120,
+  "initializer_range": 0.02,
+  "intermediate_size": 2304,
+  "is_causal": false,
+  "layer_types": ["full_attention", "full_attention", "full_attention"],
+  "max_position_embeddings": 1048576,
+  "model_type": "qwen3",
+  "moe_intermediate_size": 2304,
+  "moe_preset": "deepseek_v41",
+  "n_routed_experts": 128,
+  "n_shared_experts": 1,
+  "norm_topk_prob": true,
+  "num_attention_heads": 64,
+  "num_experts_per_tok": 3,
+  "num_hidden_layers": 3,
+  "num_key_value_heads": 1,
+  "num_target_layers": 40,
+  "o_groups": 8,
+  "o_lora_rank": 1024,
+  "pad_token_id": 2,
+  "q_lora_rank": 1280,
+  "qk_rope_head_dim": 64,
+  "rms_norm_eps": 1e-20,
+  "rope_parameters": {"rope_theta": 10000.0, "rope_type": "default"},
+  "routed_scaling_factor": 1.5,
+  "scoring_func": "sqrtsoftplus",
+  "sliding_window": null,
+  "swiglu_limit": 10.0,
+  "tie_word_embeddings": false,
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 129280
+}

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

@@ -0,0 +1,60 @@
+{
+  "architectures": ["DSparkDraftModel"],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {"AutoModel": "dspark.DSparkDraftModel"},
+  "block_size": 16,
+  "bos_token_id": 151643,
+  "dflash_config": {
+    "attention_mode": "gqa",
+    "confidence_head_alpha": 0.0,
+    "confidence_head_with_markov": false,
+    "enable_confidence_head": false,
+    "markov_head_type": "vanilla",
+    "markov_rank": 256,
+    "mask_token_id": 151669,
+    "moe_router_aux_loss_coef": 0.005,
+    "moe_router_aux_loss_type": "batch_wise",
+    "moe_router_bias_update_rate": 0.001,
+    "moe_router_enable_expert_bias": true,
+    "projector_type": "dspark",
+    "target_layer_ids": [1, 9, 17, 25, 33]
+  },
+  "dtype": "bfloat16",
+  "eos_token_id": 151645,
+  "ffn_layer_types": ["moe", "moe", "moe", "moe", "moe"],
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 12288,
+  "is_causal": true,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 40960,
+  "max_window_layers": 5,
+  "model_type": "qwen3",
+  "moe_intermediate_size": 1536,
+  "moe_preset": "qwen3_5_moe",
+  "norm_topk_prob": false,
+  "num_attention_heads": 32,
+  "num_experts": 8,
+  "num_experts_per_tok": 1,
+  "num_hidden_layers": 5,
+  "num_key_value_heads": 8,
+  "num_target_layers": 36,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000,
+  "shared_expert_intermediate_size": 0,
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

docs/recipes/deepseek-v4.1-flash-moe-dspark-disaggregated.md

@@ -0,0 +1,154 @@
+# DeepSeek-V4.1-Flash DSpark disaggregated training
+
+Fine-tunes DeepSeek-V4.1's own DSpark drafter, the `mtp.*` module the checkpoint
+ships (`configs/deepseek-v4.1-flash-moe-dspark.json`: three V4.1 blocks with
+single-pass mHC and a 128-expert MoE, 14.2B parameters), for
+`deepseek-ai/DeepSeek-V4.1-Flash` on thinking-mode data regenerated by the
+target. The recipe fits one 6-GPU B200 node: a TP4/EP4 capture server on
+GPUs 0-3 and a two-rank trainer on GPUs 4-5. It was validated end to end:
+capture, 30 optimizer steps from the official weights, export back to the
+official layout.
+
+The flow is the standard disaggregated setup
+(`docs/basic_usage/disaggregated_training.md`); what is DeepSeek-V4.1-specific:
+
+- **SGLang revision**: V4.1 needs SGLang 0.5.21 or later, and the capture plugin
+  needs the forward-observer and deferred-response extension points that landed
+  after the 0.5.23 release. The capture server therefore runs SGLang
+  `c757c4bb7c55d1eb69fdb2bdc6e6fb8507acb7a0` with
+  `plugins/sglang-spec-capture` installed; there is no source patch.
+- **Capture hook**: `--speculative-algorithm DSPARK --enable-return-aux-hidden-states
+  --return-aux-hidden-state-layer-ids 37 38 39` reaches the model's
+  `set_dspark_layers_to_capture` hook, which reads the mHC stream averaged over
+  its four copies at the inputs of layers 37-39, as the official drafter does.
+- **Expert parallelism**: TP alone splits the 2304-wide experts into 288
+  columns, which the fp4 MoE kernels reject (they need multiples of 128);
+  `--ep-size 4` keeps the experts whole.
+- **Memory**: 510 GB of fp4/fp8 weights on four 178 GB GPUs leave about 50 GB
+  each. `--mem-fraction-static 0.80 --max-running-requests 4` keeps the MoE
+  kernels' workspace for a 16k-token prefill; 0.85 runs out of memory there.
+- **Chat encoding**: the checkpoint ships a Python encoder; the
+  `deepseek-v4.1-thinking` template renders the regenerated rows with the
+  vendored copy (`DeepSeekV41Parser`), `reasoning_content` inside `<think>`.
+- **Mooncake**: the server, the trainer and `mooncake_master` must run the same
+  `mooncake-transfer-engine` wheel; a mismatch fails every publication with
+  `Batch RPC call failed: invalid rpc arg`.
+
+## Warm start
+
+```bash
+python scripts/convert_deepseek_v41_dspark.py \
+  --source deepseek-ai/DeepSeek-V4.1-Flash \
+  --output outputs/deepseek-v4.1-flash-dspark-draft
+```
+
+This is the recipe's `draft_checkpoint_path`: the official `mtp.*` tensors
+dequantized to bf16 (27 GB). The official routing biases are zero, which the
+draft config keeps (`moe_router_bias_update_rate: 0`).
+
+## Data
+
+Regenerate the prompts with the target in thinking mode against a plain V4.1
+server (`--reasoning-parser auto`):
+
+```bash
+python scripts/regenerate_train_data.py \
+  --model deepseek-ai/DeepSeek-V4.1-Flash \
+  --server-address localhost:31000 \
+  --concurrency 192 --max-tokens 16384 --temperature 1.0 --reasoning save \
+  --input-file-path prompts.jsonl \
+  --output-file-path cache/dataset/deepseek-v4.1-flash-thinking-regen.jsonl \
+  --resume
+```
+
+Rows whose reasoning runs past `--max-tokens` have no answer and land in the
+`_error` file (about 8% at 16k tokens on a mixed math/code/chat pool).
+
+## Capture server (GPUs 0-3)
+
+The server environment, as the training package's CI builds it: torch 2.14.1
+(cu130), `sglang @ git+https://github.com/sgl-project/sglang.git@c757c4bb7c55d1eb69fdb2bdc6e6fb8507acb7a0#subdirectory=python`
+installed with `--no-build-isolation`, the Mooncake wheel the trainer uses,
+then `pip install --no-deps plugins/sglang-spec-capture`.
+
+```bash
+mooncake_master \
+  --enable_http_metadata_server=true \
+  --http_metadata_server_host=0.0.0.0 \
+  --rpc_port=35551 \
+  --http_metadata_server_port=35880 \
+  --metrics_port=35903 \
+  --default_kv_lease_ttl=5m
+```
+
+```bash
+export MOONCAKE_MASTER_SERVER_ADDR=127.0.0.1:35551
+export MOONCAKE_METADATA_SERVER=http://127.0.0.1:35880/metadata
+export MOONCAKE_LOCAL_HOSTNAME=127.0.0.1
+export MOONCAKE_PROTOCOL=tcp
+export MC_TRANSFER_TIMEOUT=300
+export MOONCAKE_GLOBAL_SEGMENT_SIZE=$((200 << 30))
+export MOONCAKE_LOCAL_BUFFER_SIZE=$((1 << 30))
+export SPECFORGE_SPEC_CAPTURE=1
+CUDA_VISIBLE_DEVICES=0,1,2,3 python -m sglang.launch_server \
+  --host 0.0.0.0 \
+  --port 30000 \
+  --model-path deepseek-ai/DeepSeek-V4.1-Flash \
+  --trust-remote-code \
+  --skip-tokenizer-init \
+  --tp-size 4 \
+  --ep-size 4 \
+  --mem-fraction-static 0.80 \
+  --context-length 16896 \
+  --max-running-requests 4 \
+  --chunked-prefill-size -1 \
+  --max-prefill-tokens 16896 \
+  --speculative-algorithm DSPARK \
+  --enable-return-aux-hidden-states \
+  --return-aux-hidden-state-layer-ids 37 38 39 \
+  --return-hidden-states-mode full
+```
+
+The log must show `SpecForge spec capture enabled (host publication,
+writer=True)`. The first launch spends about 20 minutes compiling and
+autotuning the fp4 MoE kernels; later launches reuse the cache.
+
+## Trainer (GPUs 4-5)
+
+```bash
+CUDA_VISIBLE_DEVICES=4,5 specforge train \
+  -c examples/configs/online/disaggregated/external/deepseek-v4.1-flash-moe-dspark.yaml
+```
+
+Supply `HF_TOKEN` and `WANDB_API_KEY` through the environment, not YAML. On a
+host whose Hugging Face caches point at a read-only share, set
+`HF_DATASETS_CACHE` and `HF_MODULES_CACHE` to a writable directory. Measured on
+the validation run (global batch 8, two ranks): about 45 s per optimizer step,
+132 GB per trainer GPU with FULL_SHARD and the optimizer offloaded, no wait on
+capture.
+
+## Export
+
+```bash
+specforge export --to hf \
+  --checkpoint outputs/deepseek-v4.1-flash-moe-dspark-disaggregated/<run>-step<N> \
+  --draft-config configs/deepseek-v4.1-flash-moe-dspark.json \
+  --output-dir outputs/<run>-hf \
+  --embedding-source deepseek-ai/DeepSeek-V4.1-Flash --embedding-key embed.weight
+python scripts/convert_deepseek_v41_dspark.py --to-official \
+  --source outputs/<run>-hf --output outputs/<run>-official \
+  --official-config /path/to/DeepSeek-V4.1-Flash/config.json
+```
+
+The official directory holds the 1,224 unquantized `mtp.*` tensors (the
+image-token routing bias written as a copy of the text one) and the official
+config without its quantization fields.
+
+## Fresh attempts
+
+Delete the run's `outputs/` directory and, whenever a capture server was
+restarted, use a fresh Mooncake namespace (restart `mooncake_master`, or
+override `deployment.disaggregated.store_id=<unique-attempt-id>`): the
+producer dedups against keys already registered in the master, and keys whose
+replicas lived in a dead server's segment poison the consumer with
+`get_into failed` errors.

examples/configs/README.md

@@ -98,6 +98,14 @@ servers. Its
 covers the v0.5.18 SGLang capture patch and the bundled `deepseek-v4` chat
 template (the checkpoint ships no Jinja template).
 
+`deepseek-v4.1-flash-moe-dspark.yaml` fine-tunes DeepSeek-V4.1's own DSpark
+draft (`configs/deepseek-v4.1-flash-moe-dspark.json`, warm-started from the
+official `mtp.*` weights) for DeepSeek-V4.1-Flash on thinking-mode data the
+target regenerated, with one TP4/EP4 capture server on the capture plugin. Its
+[runbook](../../docs/recipes/deepseek-v4.1-flash-moe-dspark-disaggregated.md)
+covers the SGLang revision, the warm start and the export back to the official
+layout.
+
 `qwen3.8-27b-dflash2.yaml` (external services, two nodes) and
 its managed-local siblings `qwen3.8-27b-dflash2-4server-dp4.yaml`
 (one node, four capture servers plus a DP4 trainer) and