exp - feat(moe): qwen3 preset and the Qwen3.8-27B DFlash2 MoE ablation arm (stacked on #811/#812)

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
+}

configs/qwen3.8-27b-dflash2-moe.json

@@ -0,0 +1,64 @@
+{
+  "architectures": [
+    "DFlash2DraftModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": null,
+  "is_causal": false,
+  "dflash_config": {
+    "block_size": 8,
+    "conv_group_size": 16,
+    "conv_kernel_size": 2,
+    "mask_token_id": 248070,
+    "moe_bias_update_rate": 0.001,
+    "moe_dispatch": "grouped_mm",
+    "selector_rank": 256,
+    "selector_top_k": 16,
+    "target_layer_ids": [
+      5,
+      19,
+      33,
+      47,
+      61
+    ]
+  },
+  "dtype": "bfloat16",
+  "eos_token_id": 248044,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 5120,
+  "initializer_range": 0.02,
+  "intermediate_size": 17408,
+  "layer_types": [
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention"
+  ],
+  "max_position_embeddings": 262144,
+  "max_window_layers": 5,
+  "model_type": "qwen3",
+  "moe_intermediate_size": 4352,
+  "moe_preset": "qwen3",
+  "n_shared_experts": 0,
+  "num_attention_heads": 32,
+  "num_experts": 16,
+  "num_experts_per_tok": 4,
+  "num_hidden_layers": 5,
+  "num_key_value_heads": 8,
+  "num_target_layers": 64,
+  "pad_token_id": 248044,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 10000000,
+    "rope_type": "default"
+  },
+  "sliding_window": 2048,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.12.1",
+  "use_cache": true,
+  "use_sliding_window": true,
+  "vocab_size": 248320
+}

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

@@ -103,6 +103,30 @@ 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`), so
+exports load into SGLang's DeepSeek-V4 MoE unchanged.
+
 ## Fresh attempts
 
 Delete the run's `outputs/` directory and, whenever a capture server was

docs/recipes/qwen3.8-27b-dflash2-disaggregated.md

@@ -320,6 +320,36 @@ valid together. On the two-node
 wrapper, `SERVER_EXTRA_ARGS_APPEND` adds flags after the recipe's server
 defaults, and exported variables reach every server.
 
+## MoE-FFN arm (ablation)
+
+[`managed-local/qwen3.8-27b-dflash2-moe-4server-dp4-disaggregated.yaml`](https://github.com/sgl-project/SpecForge/blob/main/examples/configs/online/disaggregated/managed-local/qwen3.8-27b-dflash2-moe-4server-dp4-disaggregated.yaml)
+is the one-node recipe with `configs/qwen3.8-27b-dflash2-moe.json`: the five
+GQA decoder layers keep their attention, convolutions and selector, and each
+layer's dense 17408-wide MLP becomes a sparse MoE with Qwen3 routing
+(`moe_preset: qwen3`: softmax scores, renormalized top-k, no shared expert).
+Sizes are per run: 16 routed experts of width 4352, top-4, so the activated
+FFN width (4 x 4352) equals the dense 17408 at 4x the FFN parameters, about
+5.9B in total. The two YAMLs differ only in the draft JSON and the run names;
+the capture servers, flow control and Mooncake settings are shared unchanged,
+so the A/B diff is the FFN.
+
+The experts are ordinary FSDP-sharded parameters (no expert parallelism);
+under `SHARD_GRAD_OP` the unsharded bf16 draft is about 12 GB per rank during
+a step, within the 4+4 layout's trainer budget on 80 GB-class GPUs.
+Balancing uses the aux-loss-free selection bias (`moe_bias_update_rate`
+0.001, applied once per training forward from all-reduced expert loads); the
+trainer logs `moe/load_max_ratio`, `moe/load_min_ratio`,
+`moe/experts_unused_frac` and `moe/bias_abs_max` with the other scalars. If
+routing collapses onto a few experts early, set
+`dflash_config.moe_aux_loss_coeff` (for example 0.001) in the draft JSON.
+`moe_dispatch: grouped_mm` runs the experts as grouped GEMMs without host
+syncs; `sorted_loop` is the portable fallback. Checkpoints and exports keep
+the per-expert naming (`layers.{l}.mlp.experts.{i}.w{1,2,3}.weight`,
+`layers.{l}.mlp.gate.{weight,bias}`). SGLang's DFlash2 draft path loads dense
+drafters only, so serving the MoE arm needs a MoE-aware draft class on the
+serving side; this recipe measures training-side acceptance through the
+trainer's metrics and `spec_generate`.
+
 ## Known limitations
 
 1. **Resolved: acknowledgement stall on partial removals.**

docs/sections/advanced_features/customization.md

@@ -128,6 +128,54 @@ 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 a positive routed-expert count
+(`n_routed_experts`, or Qwen's spelling `num_experts`). 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.
+
+`qwen3` is the Qwen3-MoE preset (`softmax` scores over all experts, top-k
+combine weights renormalized, no scaling, no shared expert, plain SwiGLU).
+`configs/qwen3.8-27b-dflash2-moe.json` uses it for the MoE arm of the
+Qwen3.8-27B DFlash2 ablation, with the Qwen key names
+(`num_experts: 16`, `num_experts_per_tok: 4`, `moe_intermediate_size: 4352`,
+so the activated width equals the dense 17408). Balancing is the same
+aux-loss-free `noaux_tc` selection bias as `deepseek_v4`;
+`dflash_config.moe_aux_loss_coeff` adds the complementary balance loss when a
+from-scratch run collapses onto a few experts.
+
+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`
@@ -100,6 +105,12 @@ Qwen3.8-27B. Their
 server/trainer splits, Mooncake lease, in-flight watermarks and throughput
 measured on B300 and H200 nodes, and the two-node launcher that starts eight
 capture servers.
+`qwen3.8-27b-dflash2-moe-4server-dp4-disaggregated.yaml` is the MoE-FFN arm of
+the Qwen3.8-27B drafter-architecture ablation: the one-node recipe with
+`configs/qwen3.8-27b-dflash2-moe.json`, whose `moe_preset: qwen3` swaps each
+layer's dense MLP for 16 routed experts of width 4352, top-4 (Qwen3 routing:
+softmax scores, renormalized top-k, no shared expert); see the runbook's
+MoE section.
 
 Before running a recipe, update model/data paths and create any referenced
 offline feature or vocabulary-mapping artifacts. Managed-local recipes
@@ -172,6 +183,7 @@ assume the command runs from the repository root.
 | EAGLE3 online, external services | [`online/disaggregated/external/qwen3-8b-eagle3-disaggregated.yaml`](online/disaggregated/external/qwen3-8b-eagle3-disaggregated.yaml) |
 | DFlash online, managed-local stack | [`online/disaggregated/managed-local/qwen3-8b-dflash-1server-dp7-disaggregated.yaml`](online/disaggregated/managed-local/qwen3-8b-dflash-1server-dp7-disaggregated.yaml) |
 | DFlash 2 online, managed-local stack | [`online/disaggregated/managed-local/qwen3.8-27b-dflash2-4server-dp4-disaggregated.yaml`](online/disaggregated/managed-local/qwen3.8-27b-dflash2-4server-dp4-disaggregated.yaml) |
+| DFlash 2 online, managed-local stack, MoE-FFN draft | [`online/disaggregated/managed-local/qwen3.8-27b-dflash2-moe-4server-dp4-disaggregated.yaml`](online/disaggregated/managed-local/qwen3.8-27b-dflash2-moe-4server-dp4-disaggregated.yaml) |
 | DFlash 2 online, external services on two nodes | [`online/disaggregated/external/qwen3.8-27b-dflash2-disaggregated.yaml`](online/disaggregated/external/qwen3.8-27b-dflash2-disaggregated.yaml) |
 | Domino online, managed-local stack | [`online/disaggregated/managed-local/qwen3-8b-domino-multiserver-disaggregated.yaml`](online/disaggregated/managed-local/qwen3-8b-domino-multiserver-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