MoE support for Qwen drafts (A1, A2, B)

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

@@ -0,0 +1,77 @@
+{
+    "architectures": [
+        "DSparkDraftModel"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "auto_map": {
+        "AutoModel": "dspark.DSparkDraftModel"
+    },
+    "block_size": 7,
+    "bos_token_id": null,
+    "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": 248070,
+        "moe_aux_loss_coeff": 0.001,
+        "moe_dispatch": "grouped_mm",
+        "moe_router_center": "sample",
+        "moe_router_center_momentum": 0.99,
+        "projector_type": "dspark",
+        "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": [
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention",
+        "full_attention"
+    ],
+    "max_position_embeddings": 262144,
+    "max_window_layers": 5,
+    "model_type": "qwen3",
+    "moe_intermediate_size": 512,
+    "moe_preset": "qwen3_5_moe",
+    "n_routed_experts": 512,
+    "n_shared_experts": 1,
+    "num_attention_heads": 32,
+    "num_experts_per_tok": 10,
+    "num_hidden_layers": 5,
+    "num_key_value_heads": 8,
+    "num_target_layers": 64,
+    "pad_token_id": 248044,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+        "beta_fast": 32.0,
+        "beta_slow": 1.0,
+        "factor": 32.0,
+        "original_max_position_embeddings": 8192,
+        "rope_theta": 10000000,
+        "rope_type": "yarn"
+    },
+    "rope_theta": 10000000,
+    "shared_expert_intermediate_size": 2048,
+    "sliding_window": null,
+    "target_hidden_size": null,
+    "tie_word_embeddings": false,
+    "use_cache": true,
+    "use_sliding_window": false,
+    "vocab_size": 248320
+}

docs/sections/advanced_features/customization.md

@@ -128,6 +128,53 @@ 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, Domino) swaps its dense MLP
+for a sparse MoE FFN when the draft JSON names a `moe_preset`. The MoE is one
+configurable layer (`specforge/modeling/draft/moe/`, see its `DESIGN.md`): the
+preset is a target family's routing recipe, registered by that family's
+module under `moe/families/`, and the JSON states the per-run sizes:
+
+```json
+{
+  "moe_preset": "<family preset>",
+  "n_routed_experts": 128,
+  "num_experts_per_tok": 8,
+  "moe_intermediate_size": 1024,
+  "dflash_config": {"moe_dispatch": "grouped_mm"}
+}
+```
+
+- **Architecture keys** at the top level override the preset for ablations
+  (`scoring_func`, `norm_topk_prob`, `routed_scaling_factor`, `balance`,
+  `n_shared_experts`, `shared_expert_intermediate_size`, `shared_expert_gate`,
+  `moe_router_bias`, ...). `moe_layout` overrides the checkpoint file naming
+  the preset selects. Expert counts must be positive integers, and
+  `n_routed_experts` without a `moe_preset` is an error, never a silent dense
+  fallback.
+- **Training-only keys** live under `dflash_config` with an `moe_` prefix and
+  never change the checkpoint's architecture: `moe_aux_loss_coeff` (the
+  load-balancing loss the router adds, on by default at 0.001; `0` turns it
+  off) and `moe_aux_loss_type` (its formula), `moe_dispatch` (`sorted_loop`
+  or `grouped_mm`), and router centering `moe_router_center` (`"sample"`)
+  with `moe_router_center_momentum`. From-scratch drafters whose routers
+  collapse onto a few experts need `moe_router_center`.
+- **Losses and metrics** need no strategy code: the trainer adds every MoE
+  layer's auxiliary loss to the objective and logs layer-averaged `moe/*`
+  load metrics for every strategy.
+- **Checkpoints and exports** name MoE tensors through the preset's layout.
+  `export --to hf` also folds training-only router state into the gate
+  (router centering becomes a fixed router logit bias, `moe_router_bias`),
+  rewrites `config.json` to describe the folded draft so the export reloads
+  as is, and adds the routing fields the family's serving engine reads.
+  Serving an MoE draft requires that engine to implement the draft
+  architecture with an MoE FFN; check the family's documentation.
+- **New families and components** register by name: a module under
+  `moe/families/` registers a family's preset and layout, and a score
+  function, router, aux loss, balance controller, experts backend or shared
+  expert registers from its own module (see `DESIGN.md`).
+
 ## Draft architectures
 
 Draft classes register through `@register_draft`. The key defaults to the

examples/configs/README.md

@@ -110,6 +110,19 @@ 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-dspark-moe-h200-disaggregated.yaml` (one 8x H200 node, four FP8
+capture servers plus a 4-rank FULL_SHARD trainer with optimizer offload)
+trains the DSpark drafter with a Qwen3.5-MoE FFN in
+`configs/qwen3.8-27b-dspark-moe.json` (`moe_preset: qwen3_5_moe`, 512 experts
+of width 512, top-10, router centering) from scratch. See "MoE FFN for
+DFlash-family drafts" in the
+[customization guide](../../docs/sections/advanced_features/customization.md)
+for the MoE keys; serving this drafter needs an SGLang `Qwen3MoeDSparkModel`.
+Launch it with `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`: its trainer
+ranks run at the H200 memory limit, where the caching allocator otherwise
+fragments (see "Late OOM or non-finite hidden states on online runs" in the
+[training guide](../../docs/sections/basic_usage/training.md)).
+
 Before running a recipe, update model/data paths and create any referenced
 offline feature or vocabulary-mapping artifacts. Managed-local recipes
 intentionally record their GPU allocation and loopback services. External

examples/configs/online/disaggregated/managed-local/qwen3.8-27b-dspark-moe-h200-disaggregated.yaml

@@ -0,0 +1,112 @@
+# Qwen3.8-27B DSpark with a Qwen3.5-MoE FFN (512 experts x 512, top-10, one
+# sigmoid-gated 2048-wide shared expert; configs/qwen3.8-27b-dspark-moe.json)
+# on one 8x H200 node: four TP=1 FP8 capture servers on GPUs 0-3 and a 4-rank
+# FSDP trainer on GPUs 4-7, all owned by `specforge train`.
+#
+# The drafter has about 21B parameters, 20B of them routed experts (about 1B
+# are active per token), so the trainer shards everything (FULL_SHARD) and
+# offloads optimizer state to the host. The draft trains from scratch: the target is dense.
+# Router centering (moe_router_center: sample) and the aux loss (0.001) keep
+# the router from collapsing onto a few experts; watch moe/experts_unused_frac
+# and moe/route_entropy_frac.
+#
+# Launch with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True. The trainer
+# ranks run at the H200 memory limit (each micro-step all-gathers a layer's
+# 2.5 GiB stacked experts while no_sync() holds the 41.6 GB unsharded
+# gradient), and without expandable segments the caching allocator fragments
+# and retries failed allocations on every micro-step.
+model:
+  target_model_path: Qwen/Qwen3.8-27B-FP8
+  draft_model_config: configs/qwen3.8-27b-dspark-moe.json
+  target_backend: sglang
+  trust_remote_code: true
+  embedding_key: model.language_model.embed_tokens.weight
+  lm_head_key: lm_head.weight
+  load_target_embedding: true
+  torch_dtype: bfloat16
+  mask_token_id: 248070
+  # FA3 for the hybrid target's full-attention layers on Hopper; the GDN
+  # layers use SGLang's default linear-attention kernels.
+  sglang_attention_backend: fa3
+  sglang_disable_radix_cache: true
+  # data.max_length + 7 tokens of request headroom for the capture endpoint.
+  sglang_context_length: 8199
+data:
+  train_data_path: ./cache/dataset/qwen3.8-27b-regen-train.jsonl
+  max_length: 8192
+  chat_template: qwen3.5
+  build_dataset_num_proc: 32
+  dataloader_num_workers: 8
+  cache_dir: cache/qwen3.8-27b-dspark-moe
+training:
+  strategy: dspark
+  num_epochs: 1
+  # 4 ranks x 2 x 64 = 512 samples per optimizer step. On B200 the same run
+  # used batch_size 4 x accumulation 32; H200 ranks have less memory, so this
+  # recipe halves the micro-batch and doubles the accumulation.
+  batch_size: 2
+  accumulation_steps: 64
+  fsdp_sharding: FULL_SHARD
+  optimizer_cpu_offload: true
+  learning_rate: 0.0005
+  lr_scheduler: constant
+  warmup_ratio: 0.04
+  max_grad_norm: 1.0
+  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: 500
+  max_checkpoints: 3
+  log_interval: 10
+  dist_timeout: 15
+  seed: 42
+  prompt_seed: 42
+tracking:
+  report_to: none
+runtime:
+  producer_lease: 8
+  producer_concurrency: 8
+  in_flight_high_watermark: 768
+  in_flight_low_watermark: 512
+run_id: qwen3.8-27b-dspark-moe-h200-disaggregated
+output_dir: outputs/qwen3.8-27b-dspark-moe-h200-disaggregated
+deployment:
+  mode: disaggregated
+  trainer:
+    nnodes: 1
+    nproc_per_node: 4
+  disaggregated:
+    control_dir: outputs/qwen3.8-27b-dspark-moe-h200-disaggregated/control
+    consumer_state_dir: outputs/qwen3.8-27b-dspark-moe-h200-disaggregated/consumer-state
+    backend: mooncake
+    client_buffer_size: 1073741824
+    managed_local:
+      trainer_cuda_visible_devices: ["4", "5", "6", "7"]
+      mooncake:
+        # TCP works everywhere; RDMA loopback (protocol: rdma plus
+        # rdma_devices) removes the CPU copies of every captured byte.
+        protocol: tcp
+        global_segment_size_bytes: 107374182400
+        local_buffer_size_bytes: 1073741824
+        default_kv_lease_ttl_ms: 300000
+      capture_servers:
+        - port: 30000
+          cuda_visible_devices: ["0"]
+          tp_size: 1
+          mem_fraction_static: 0.85
+        - port: 30001
+          cuda_visible_devices: ["1"]
+          tp_size: 1
+          mem_fraction_static: 0.85
+        - port: 30002
+          cuda_visible_devices: ["2"]
+          tp_size: 1
+          mem_fraction_static: 0.85
+        - port: 30003
+          cuda_visible_devices: ["3"]
+          tp_size: 1
+          mem_fraction_static: 0.85

scripts/gates/normalize_dflash_export.py

@@ -86,6 +86,16 @@ def _normalize_dspark(config: Dict[str, Any], method_config: Dict[str, Any]) ->
     config["architectures"] = ["Qwen3DSparkModel"]
 
 
+def _apply_moe_serving(config: Dict[str, Any]) -> None:
+    """MoE drafts: the family's serving class and the routing fields its engine reads."""
+    if not config.get("moe_preset"):
+        return
+    from specforge.modeling.draft.moe import serving_architecture, serving_config
+
+    config["architectures"] = [serving_architecture(config, config["architectures"][0])]
+    config.update(serving_config(config))
+
+
 def normalize_export(config_path: str, expected_block_size: int) -> Dict[str, Any]:
     path = Path(config_path)
     with path.open(encoding="utf-8") as handle:
@@ -132,6 +142,7 @@ def normalize_export(config_path: str, expected_block_size: int) -> Dict[str, An
         _normalize_dflash2(config, method_config)
     else:
         config["architectures"] = ["DFlashDraftModel"]
+    _apply_moe_serving(config)
     config.pop("auto_map", None)
     with path.open("w", encoding="utf-8") as handle:
         json.dump(config, handle, indent=2)

specforge/export/checkpoint_io.py

@@ -177,7 +177,12 @@ def materialize_draft(
     draft_config = AutoDraftModelConfig.from_file(draft_config_path)
     torch_dtype = checkpoint_dtype(state["draft_state_dict"]) or torch.bfloat16
     model = AutoDraftModel.from_config(draft_config, torch_dtype=torch_dtype)
-    missing, unexpected = model.load_state_dict(state["draft_state_dict"], strict=False)
+    from specforge.modeling.draft.moe import from_checkpoint_state_dict
+
+    # draft_state_dict uses the draft's checkpoint layout (MoE file naming).
+    missing, unexpected = model.load_state_dict(
+        from_checkpoint_state_dict(state["draft_state_dict"], model), strict=False
+    )
     if unexpected:
         raise ValueError(
             f"checkpoint carries weights the {type(model).__name__} architecture "

specforge/export/to_hf.py

@@ -24,6 +24,11 @@
     materialize_draft,
     resolve_training_state,
 )
+from specforge.modeling.draft.moe import (
+    exported_config_updates,
+    from_checkpoint_state_dict,
+    to_serving_state_dict,
+)
 from specforge.modeling.target.checkpoint import load_checkpoint_tensors
 
 
@@ -75,15 +80,23 @@ def export_to_hf(
             embedding_source,
             keys=[embedding_key],
         )[embedding_key]
-    # Preserve checkpoint weights at their original precision, but keep the
-    # mapping explicitly loaded by materialize_draft when one was requested.
+    # Trained keys win at their original precision; keep a requested vocab mapping.
     full_state.update(
         {
             key: value
-            for key, value in state["draft_state_dict"].items()
+            for key, value in from_checkpoint_state_dict(
+                state["draft_state_dict"], model
+            ).items()
             if not (vocab_mapping_path and key in {"t2d", "d2t"})
         }
     )
+    # MoE drafts: fold centering, use the family's naming, update config.json to match.
+    updates = exported_config_updates(model)
+    full_state = to_serving_state_dict(full_state, model)
+    for key, value in updates.items():
+        if key == "dflash_config":
+            value = {**(getattr(model.config, key, None) or {}), **value}
+        setattr(model.config, key, value)
     model.save_pretrained(output_dir, state_dict=full_state)
     apply_legacy_rope_scaling(output_dir)
     return output_dir