exp.[rebase] MoE FFN for DFlash-family drafters (#811/#812) on current main, to run the DeepSeek-V4-Flash MoE arm

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,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/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

scripts/warm_start_moe_drafter.py

@@ -0,0 +1,112 @@
+#!/usr/bin/env python3
+"""Build a warm-start source for an MoE DFlash-family drafter from a DeepSeek-V4 target.
+
+Constructs the draft model from its config (random init for attention, heads,
+norms, projections), seeds every MoE layer's routed experts, gate weight,
+noaux bias and shared expert from one target layer (dequantized fp4/fp8 ->
+bf16), and writes an HF-format directory usable as
+``model.draft_checkpoint_path``. Requires the draft's expert shape to match the
+target's (256 x 2048 for DeepSeek-V4-Flash); a smaller draft takes a strided
+subset of experts.
+
+Example:
+  python scripts/warm_start_moe_drafter.py \
+      --draft-config examples/configs/kan-ablations/deepseek-v4-flash-dspark-moe256-auxbal.json \
+      --target-snapshot /cluster-storage/models/models--deepseek-ai--DeepSeek-V4-Flash-0731/snapshots/<sha> \
+      --target-layers 3,11,21,31,41 --output-dir warm-starts/moe256-from-0731
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import time
+
+import torch
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+    ap.add_argument("--draft-config", required=True)
+    ap.add_argument(
+        "--target-snapshot", required=True, help="local HF snapshot dir of the target"
+    )
+    ap.add_argument(
+        "--target-layers",
+        required=True,
+        help="comma list, one target layer per draft layer",
+    )
+    ap.add_argument("--output-dir", required=True)
+    ap.add_argument("--seed", type=int, default=42)
+    ap.add_argument(
+        "--select",
+        default="strided",
+        help="expert selection when draft has fewer experts",
+    )
+    args = ap.parse_args()
+
+    from specforge.modeling.auto import AutoDraftModel, AutoDraftModelConfig
+    from specforge.modeling.draft.moe import (
+        apply_warm_start,
+        iter_moe_layers,
+        plan_warm_start,
+        resolve_moe_config,
+        to_checkpoint_state_dict,
+    )
+    from specforge.modeling.draft.moe.deepseek_v4_target import load_target_moe_layer
+
+    target_layers = [int(x) for x in args.target_layers.split(",")]
+    config = AutoDraftModelConfig.from_file(args.draft_config)
+    moe_cfg = resolve_moe_config(config)
+    if moe_cfg is None:
+        raise SystemExit("draft config is dense (n_routed_experts == 0)")
+    torch.manual_seed(args.seed)
+    t0 = time.time()
+    model = AutoDraftModel.from_config(config, torch_dtype=torch.bfloat16)
+    layers = list(iter_moe_layers(model))
+    if len(layers) != len(target_layers):
+        raise SystemExit(
+            f"{len(layers)} MoE layers but {len(target_layers)} target layers given"
+        )
+    print(
+        f"built draft ({sum(p.numel() for p in model.parameters())/1e9:.2f}B params) in {time.time()-t0:.0f}s"
+    )
+
+    target_cfg = json.load(open(os.path.join(args.target_snapshot, "config.json")))
+    n_target = int(target_cfg["n_routed_experts"])
+    for i, (layer, tl) in enumerate(zip(layers, target_layers)):
+        t1 = time.time()
+        source = load_target_moe_layer(args.target_snapshot, tl)
+        plan = plan_warm_start(
+            layer.cfg, n_target_experts=n_target, strategy=args.select
+        )
+        loaded = apply_warm_start(layer, plan, source)
+        print(
+            f"draft layer {i} <- target layer {tl}: {len(loaded)} tensors, "
+            f"experts {plan.target_expert_ids[:3]}...{plan.target_expert_ids[-1]} ({time.time()-t1:.0f}s)"
+        )
+
+    for key, value in moe_cfg.serving_fields().items():
+        setattr(model.config, key, value)
+    model.config.moe_warm_start = {
+        "target": os.path.basename(
+            os.path.dirname(os.path.dirname(args.target_snapshot.rstrip("/")))
+        ),
+        "snapshot": os.path.basename(args.target_snapshot.rstrip("/")),
+        "target_layers": target_layers,
+        "select": args.select,
+        "seed": args.seed,
+    }
+    os.makedirs(args.output_dir, exist_ok=True)
+    model.save_pretrained(
+        args.output_dir, state_dict=to_checkpoint_state_dict(model.state_dict())
+    )
+    print(f"wrote {args.output_dir} in {time.time()-t0:.0f}s total")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

specforge/export/checkpoint_io.py

@@ -159,7 +159,12 @@ def materialize_draft(
 
     draft_config = AutoDraftModelConfig.from_file(draft_config_path)
     model = AutoDraftModel.from_config(draft_config, torch_dtype=torch.bfloat16)
-    missing, unexpected = model.load_state_dict(state["draft_state_dict"], strict=False)
+    from specforge.modeling.draft.moe import from_checkpoint_state_dict
+
+    # Files use the official naming; modules may use a native MoE layout.
+    missing, unexpected = model.load_state_dict(
+        from_checkpoint_state_dict(state["draft_state_dict"]), strict=False
+    )
     if unexpected:
         raise ValueError(
             f"checkpoint carries weights the {type(model).__name__} architecture "

specforge/export/to_hf.py

@@ -31,6 +31,7 @@
     materialize_draft,
     resolve_training_state,
 )
+from specforge.modeling.draft.moe import resolve_moe_config, to_checkpoint_state_dict
 
 
 def _load_embedding_tensor(source: str, key: str) -> torch.Tensor:
@@ -87,7 +88,13 @@ def export_to_hf(
     model = materialize_draft(
         state, draft_config_path, vocab_mapping_path=vocab_mapping_path
     )
-    full_state = dict(model.state_dict())
+    full_state = dict(to_checkpoint_state_dict(model.state_dict()))
+    moe_cfg = resolve_moe_config(model.config)
+    if moe_cfg is not None:
+        # Serving engines read the routing recipe from config.json; the draft
+        # JSON only names a preset, so materialize the resolved fields.
+        for key, value in moe_cfg.serving_fields().items():
+            setattr(model.config, key, value)
     owns_embedding = hasattr(model, "embed_tokens")
     if owns_embedding and "embed_tokens.weight" not in state["draft_state_dict"]:
         if not embedding_source: