Sparse MoE FFN for DFlash-family drafters

[sherlockwu]

Stacked on #798 (retarget to main as the stack merges).

Adds a DeepSeek-V4-style sparse MoE FFN that any DFlash-family draft config can opt into, plus a ready-made dense-vs-MoE ablation pair for DeepSeek-V4-Flash.

What's here

• specforge/modeling/draft/moe.py — self-contained MoE block:
  • MoEGate: sqrtsoftplus top-k routing with the aux-loss-free balancing bias. The bias is an fp32 buffer updated by a sign controller from all-reduced expert loads (never by gradients) and survives module-wide dtype casts.
  • GroupedExperts: routed experts as three stacked [E, out, in] parameters — grouped-GEMM-ready, and FSDP use_orig_params tracks 3 view tensors instead of 3*E. Loads both native stacked keys and the official per-expert naming (experts.{i}.w{1,2,3}.weight).
  • SparseMoE: sorted-segment dispatch (one argsort → contiguous per-expert segments; a per-expert torch.where loop scales overhead with active experts, ~2× step time once the balancer spreads load), optional torch._grouped_mm path with no host syncs, one shared expert.
  • Balance updates are deferred: forward N stashes routing counts, the model forward applies them before forward N+1, outside any activation-checkpoint region — a bias mutated between a forward and its recompute would route to different segment shapes and break checkpointing.
• dflash.py: a draft JSON with n_routed_experts > 0 swaps the dense MLP for SparseMoE; the bare gate Parameter is initialized via _init_weights; moe_bias_update_rate comes from dflash_config.
• backend.py: checkpoint files keep the official per-expert expert naming while modules keep stacked parameters — unstack/stack_grouped_expert_state_dict convert at the save/load boundary, so warm starts, resumes, and export bundlers are unaffected.
• Ablation pair for DSV4-Flash: configs/deepseek-v4-flash-dspark-moe.json is deepseek-v4-flash-dspark.json plus the MoE fields (64 routed + 1 shared, top-6, moe_intermediate_size 2048 — activated width ≈ the dense 12288); examples/.../deepseek-v4-flash-dspark-moe.yaml mirrors the dense recipe so the draft config is the A/B diff.

Tests

tests/test_modeling/test_moe.py: balance-update semantics (train stashes / eval doesn't / second apply is a no-op), fp32 bias through casts, stacked↔official state-dict round trip (including wrapper-level stack), grouped-GEMM vs sorted-loop parity in outputs/grads on CUDA, DFlash wiring (MoE vs dense selection, gate init, model-level round trip).

🤖 Generated with Claude Code