exp - feat(serving): SGLang MoE DSpark draft loader for Kan's MoE drafter (#811/#812), normalizer arch name and equivalence gate (stacked on #931)

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

@@ -0,0 +1,60 @@
+{
+  "architectures": ["DSparkDraftModel"],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {"AutoModel": "dspark.DSparkDraftModel"},
+  "block_size": 7,
+  "bos_token_id": 0,
+  "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": 128799,
+    "moe_bias_update_rate": 0.001,
+    "moe_dispatch": "grouped_mm",
+    "projector_type": "dspark",
+    "target_layer_ids": [1, 11, 21, 31, 41]
+  },
+  "dtype": "bfloat16",
+  "eos_token_id": 1,
+  "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": 1048576,
+  "max_window_layers": 5,
+  "model_type": "qwen3",
+  "moe_intermediate_size": 2048,
+  "moe_preset": "deepseek_v4",
+  "n_routed_experts": 64,
+  "n_shared_experts": 1,
+  "num_attention_heads": 32,
+  "num_experts_per_tok": 6,
+  "num_hidden_layers": 5,
+  "num_key_value_heads": 8,
+  "num_target_layers": 43,
+  "pad_token_id": 1,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "factor": 16.0,
+    "original_max_position_embeddings": 65536,
+    "rope_theta": 10000.0,
+    "rope_type": "yarn"
+  },
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 129280
+}

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

@@ -103,6 +103,47 @@ MI355X node over the full two epochs (1,885 optimizer steps): 6.32 s per
 128-sample step on average, 6.1-6.4 s at steady state (about 3.1 s waiting for
 capture and 3.0 s of trainer compute), 3 h 18 min end to end.
 
+## MoE-FFN arm (ablation)
+
+`examples/configs/online/disaggregated/external/deepseek-v4-flash-dspark-moe-disaggregated.yaml`
+is the same recipe with `configs/deepseek-v4-flash-dspark-moe.json`: the
+five-layer GQA decoder keeps its attention, and each layer's dense MLP becomes
+the target's MoE (`moe_preset: deepseek_v4`: sqrt-softplus scores,
+aux-loss-free top-k with the sign-controlled selection bias, combine weights
+renormalized and scaled by 1.5, one ungated shared expert, SwiGLU clamped at
+10). Sizes are per run: 64 routed experts, top-6, width 2048, so the activated
+FFN width (6 x 2048 + 2048 shared) matches the dense 12288 at ~10x the FFN
+parameters. Run it against the dense recipe with identical hparams; the two
+YAMLs differ only in the draft JSON and run names. The capture servers are
+shared by both arms unchanged.
+
+Training-only knobs live under the draft JSON's `dflash_config`:
+`moe_bias_update_rate` (0.001, the balancing controller's step) and
+`moe_dispatch` (`grouped_mm` runs the experts as grouped GEMMs with no host
+sync; `sorted_loop` is the portable fallback). The trainer logs `moe/*` load
+metrics (max/min load ratios, unused-expert fraction, balancing-bias
+magnitude) alongside the usual scalars. Checkpoints keep the official
+per-expert naming (`layers.N.mlp.experts.{i}.w{1,2,3}.weight`,
+`layers.N.mlp.gate.bias`, `layers.N.mlp.shared_experts.w{1,2,3}.weight`).
+
+Serving the MoE arm needs an MoE-capable draft class on the SGLang side:
+the stock `Qwen3DSparkModel` has a dense MLP and silently drops every expert
+weight (the server starts, but the drafter is random and acceptance length
+sits at ~1.0). Until that class is upstream, apply
+`patches/sglang/v0.5.18/dspark-moe-draft.patch` on top of the spec-capture
+patch (`cd <sglang checkout or site-packages parent> && git apply
+<SpecForge>/patches/sglang/v0.5.18/dspark-moe-draft.patch`). It adds
+`sglang/srt/models/dspark_moe.py` (`Qwen3MoEDSparkModel`: the DSpark decoder
+with the dense MLP replaced by the DeepSeek-V4 routing above, experts loaded
+as stacked grouped-GEMM weights) and makes the DFlash-family loaders reject a
+checkpoint whose weights do not match the class, instead of serving
+uninitialised modules. `scripts/gates/normalize_dflash_export.py` writes
+`architectures: ["Qwen3MoEDSparkModel"]` for an export with
+`n_routed_experts > 0` and refuses MoE DFlash/DFlash2 exports, which have no
+serving class yet. `scripts/gates/check_dspark_moe_sglang_equivalence.py`
+checks the serving FFN against SpecForge's `MoELayer` bit-for-bit on the real
+64-expert sizes (grouped_mm and loop paths).
+
 ## Fresh attempts
 
 Delete the run's `outputs/` directory and, whenever a capture server was

docs/sections/advanced_features/customization.md

@@ -128,6 +128,42 @@ drafts currently implements the GQA/MHA layout only, so plan benchmarks
 accordingly. DFlash2 otherwise follows the same mode selection; its convolution
 and selector do not change the attention projection contract.
 
+## MoE FFN for DFlash-family drafts
+
+Any DFlash-family draft (DFlash, DFlash2, DSpark) swaps its dense MLP for a
+sparse MoE FFN when the draft JSON sets `n_routed_experts > 0`. The MoE is one
+configurable layer (`specforge/modeling/draft/moe/`, see its `DESIGN.md`):
+a `moe_preset` names a target family's routing recipe, and the architecture
+keys use the target checkpoints' native HF names so they can be copied from
+the target's `config.json`:
+
+```json
+{
+  "moe_preset": "deepseek_v4",
+  "n_routed_experts": 64,
+  "num_experts_per_tok": 6,
+  "moe_intermediate_size": 2048,
+  "n_shared_experts": 1,
+  "dflash_config": {"moe_bias_update_rate": 0.001, "moe_dispatch": "grouped_mm"}
+}
+```
+
+Top-level keys override the preset (for ablations: `scoring_func`,
+`norm_topk_prob`, `routed_scaling_factor`, `balance`, `shared_expert_gate`,
+`swiglu_limit`, ...). Training-only knobs live under `dflash_config` with an
+`moe_` prefix and never change the checkpoint. Checkpoints, warm starts and
+exports keep the official per-expert naming (`experts.{i}.w{1,2,3}.weight`),
+so an exported drafter loads into SGLang unchanged. Dense drafts are
+unaffected: with no `n_routed_experts` the kernel provider's MLP is used as-is.
+
+`deepseek_v4` is the checked-in preset (DeepSeek-V4 routing:
+`sqrtsoftplus` scores, aux-loss-free `noaux_tc` balancing, combine weights
+renormalized and scaled by 1.5, one ungated shared expert, SwiGLU clamp 10);
+`configs/deepseek-v4-flash-dspark-moe.json` uses it. A new target family is a
+preset registration plus whichever components it needs (score function,
+balance controller, experts backend, shared expert); each registers by name
+from its own module.
+
 ## Draft architectures
 
 Draft classes register through `@register_draft`. The key defaults to the

examples/configs/README.md

@@ -88,6 +88,11 @@ servers. Its
 [runbook](../../docs/recipes/deepseek-v4-flash-dspark-disaggregated.md)
 covers the v0.5.18 SGLang capture patch and the bundled `deepseek-v4` chat
 template (the checkpoint ships no Jinja template).
+`deepseek-v4-flash-dspark-moe-disaggregated.yaml` is the MoE-FFN arm of the
+drafter-architecture ablation: the same recipe with
+`configs/deepseek-v4-flash-dspark-moe.json`, whose `moe_preset: deepseek_v4`
+swaps the dense MLP for the target's routing (64 routed + 1 shared experts,
+top-6, width 2048); see the runbook's MoE section.
 
 `qwen3.8-27b-dflash2-disaggregated.yaml` (external services, two nodes) and
 its managed-local siblings `qwen3.8-27b-dflash2-4server-dp4-disaggregated.yaml`

examples/configs/online/disaggregated/external/deepseek-v4-flash-dspark-moe-disaggregated.yaml

@@ -0,0 +1,97 @@
+# MoE-FFN arm of the DSpark drafter-architecture ablation for
+# DeepSeek-V4-Flash. Identical to deepseek-v4-flash-dspark-disaggregated.yaml
+# (the dense arm) except the draft JSON and the run/store names, so the A/B
+# diff is the FFN: 64 routed experts + 1 shared, top-6, moe_intermediate 2048
+# (activated width 6x2048 + 2048 ~= the dense 12288) with DeepSeek-V4 routing.
+# If a from-scratch MoE run destabilizes at the shared hparams, lower the LR
+# and clip on BOTH arms rather than on this file alone.
+model:
+  target_model_path: deepseek-ai/DeepSeek-V4-Flash-0731
+  draft_model_config: configs/deepseek-v4-flash-dspark-moe.json
+  target_backend: sglang
+  trust_remote_code: true
+  # DeepSeek-V4 checkpoints use the native inference weight layout.
+  embedding_key: embed.weight
+  lm_head_key: head.weight
+  mask_token_id: 128799
+  torch_dtype: bfloat16
+  sglang_mem_fraction_static: 0.85
+  # data.max_length plus headroom: capture rejects inputs at exactly context length.
+  sglang_context_length: 8704
+  sglang_max_running_requests: 8
+  # Routed experts are fp4; the default MoE path cannot run them on B200.
+  sglang_moe_runner_backend: flashinfer_mxfp4
+
+data:
+  train_data_path: ./cache/dataset/sharegpt_train.jsonl
+  max_length: 8192
+  chat_template: deepseek-v4
+  cache_dir: cache
+  build_dataset_num_proc: 64
+  # Async prefetch: overlaps per-sample Mooncake feature fetches with compute.
+  dataloader_num_workers: 8
+
+training:
+  strategy: dspark
+  num_epochs: 2
+  # 4 ranks x 32 microbatches -> global batch 128.
+  batch_size: 1
+  accumulation_steps: 32
+  learning_rate: 0.0006
+  lr_scheduler: constant
+  warmup_ratio: 0
+  max_grad_norm: 1
+  attention_backend: flex_attention
+  num_anchors: 512
+  loss_decay_gamma: 4.0
+  objective_chunk_blocks: 128
+  dspark_ce_loss_alpha: 0.1
+  dspark_l1_loss_alpha: 0.9
+  dspark_confidence_head_alpha: 1.0
+  save_interval: 128
+  log_interval: 10
+  dist_timeout: 30
+  seed: 42
+  prompt_seed: 1
+
+tracking:
+  report_to: wandb
+  wandb_project: specforge
+  wandb_name: deepseek-v4-flash-dspark-moe-disaggregated
+  wandb_dir: outputs/deepseek-v4-flash-dspark-moe-disaggregated/wandb
+
+runtime:
+  producer_lease: 8
+  producer_concurrency: 8
+  # Keep two 128-sample optimizer quanta in flight (~0.4 GiB features/sample);
+  # the low watermark is one full quantum so complete windows always dispatch.
+  in_flight_high_watermark: 256
+  in_flight_low_watermark: 128
+  resident_high_watermark_bytes: 137438953472
+  resident_low_watermark_bytes: 103079215104
+  feature_store_max_resident_bytes: 171798691840
+
+run_id: deepseek-v4-flash-dspark-moe-disaggregated
+output_dir: outputs/deepseek-v4-flash-dspark-moe-disaggregated
+
+deployment:
+  mode: disaggregated
+  trainer:
+    nnodes: 1
+    nproc_per_node: 4
+  disaggregated:
+    control_dir: outputs/deepseek-v4-flash-dspark-moe-disaggregated/control
+    consumer_state_dir: outputs/deepseek-v4-flash-dspark-moe-disaggregated/consumer-state
+    backend: mooncake
+    store_id: deepseek-v4-flash-dspark-moe-disaggregated
+    # Two TP2 servers out-produce one TP4 server (TP prefill scaling is sublinear).
+    server_urls:
+      - http://127.0.0.1:30000
+      - http://127.0.0.1:30001
+    mooncake_metadata_server: http://127.0.0.1:35880/metadata
+    mooncake_master_server_addr: 127.0.0.1:35551
+    mooncake_local_hostname: 127.0.0.1
+    mooncake_protocol: tcp
+    client_buffer_size: 1073741824
+    idle_timeout_s: 7200
+    peer_wait_timeout_s: 7200