[DFlash2] Add configurable MoE draft FFN

[qianlihuang]

Purpose

Refs #1057.

Make the DFlash2 feed-forward implementation configurable while preserving the existing attention, convolution, selector, and draft-block paths. draft_ffn_type=moe adds a Qwen3.5-compatible routed-plus-shared expert block; dense remains the default.

The checkpoint contract stores fused routed-expert tensors as [experts, 2 * intermediate, hidden] and [experts, hidden, intermediate], while MoE geometry remains outside the shared attention config. Invalid top-k geometry fails closed.

This first PR intentionally excludes router auxiliary loss, routing diagnostics, verifier warm start, expert parallelism, and serving-runtime changes so those concerns can be reviewed independently.

Tests

uvx ruff==0.15.20 format --check <changed files>
6 files already formatted

uvx ruff==0.15.20 check <changed files>
All checks passed!

CUDA_VISIBLE_DEVICES= python -m pytest \
  tests/unit/models/test_dflash2_model_definitions.py \
  tests/unit/train/config/test_schema.py -q
37 passed, 1 skipped

The focused suite covers dense-default compatibility, config serialization, fused checkpoint keys, non-mutating model config, fail-closed geometry, and reference MoE forward/backward. The skipped test requires CUDA.

Checklist

I have filled in:

• The purpose of the PR and the RFC it relates to.
• The test plan and results.
• Optional documentation update.
• I (a human) have written or reviewed the code in this PR to the best of my ability.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews

	Protection	Waiting on
🔴	Require approval from approved reviewers list	👀 reviews

🔴 Require approval from approved reviewers list

Waiting for any of

• approved-reviews-by = dsikka
• approved-reviews-by = fynnsu
• approved-reviews-by = orestis-z
• approved-reviews-by = rahul-tuli
• approved-reviews-by = shanjiaz

This rule is failing.

All pull requests must have at least one approving review from a member of the approved reviewers list before merging.

• any of:
  • approved-reviews-by = dsikka
  • approved-reviews-by = fynnsu
  • approved-reviews-by = orestis-z
  • approved-reviews-by = rahul-tuli
  • approved-reviews-by = shanjiaz