Support R3 capture with DeepGEMM MegaMoE

[aoshen02]

Summary

Add routed-experts (R3) capture support to the deep_gemm_mega_moe path used by DeepSeek-V4 and Kimi-K3.

Why

The existing binder only recognizes MoE layers implemented through MoERunner. DeepGEMM MegaMoE bypasses that runner, even though its shared expert implementation already receives logical topk_ids before EPLB maps them to physical replicas. Enabling R3 therefore failed during model initialization with No supported MoE router found for routed-experts capture.

This does not duplicate an existing change. Upstream PR vllm-project#44214 exposes routing replay for FlashInfer TRT-LLM monolithic kernels; it does not cover DeepGEMM MegaMoE. Searches for open PRs matching deep_gemm_mega_moe R3, routing replay MegaMoE, and related terms found no implementation.

Design

• Define one small structural capture-source interface next to the existing R3 binder.
• Bind the capture callback directly during model initialization.
• Capture the existing logical topk_ids immediately before each model's existing EPLB mapping block.
• Do not change the MegaMoE kernel ABI, recompute routing, or add model-specific branches to the binder.
• When R3 is disabled, each MegaMoE forward adds only one None check; it adds no tensor operation, allocation, copy, or kernel launch.

Unsupported MoE paths continue to fail closed during initialization.

Validation

• pre-commit on all four changed files: passed, including Ruff, formatting, mypy, sign-off, forbidden-import and configuration checks.
• pytest tests/model_executor/test_routed_experts_capture.py -k public_binding -q: 2 passed.
• Added one parametrized model regression covering both DeepSeek-V4 and Kimi-K3 binding and logical-ID capture before EPLB.
• The full CPU capture file reached 15 passed / 3 skipped; its pre-existing CUDA async-output test could not import the local FlashAttention extension in this host venv.
• GB200 Kimi-K3 model E2E and output validation are pending and will be added before this draft is marked ready.

AI assistance

This change was developed with OpenAI Codex assistance. The human submitter must review and understand every changed line before submission upstream.