[Refactor][MoE] Separate dispatched expert compute

[michaelwithu]

Summary

• Extract communication-independent dispatched-expert MLP compute for tight expert-major batches.
• Compose routed expert execution from local dispatch, expert compute, and combine while preserving the public API.
• Add GPU-resident dynamic row counts, static capacity, and CUDA Graph replay.
• Add local dispatch and an optional DeepEP V2 adapter without introducing a runtime DeepEP dependency.
• Keep routing weights and opaque combine handles outside expert compute.

Review follow-up

• Resolve one effective grouped-GEMM kernel before constructing gate-up/down sub-ops, so an explicit override cannot undo the alignment fallback.
• Consolidate normal and device-bounded fused-gated paths around one column mapping and one pointwise epilogue while keeping distinct compiled signatures and no runtime row-policy branch.
• Remove exact-signature/dataclass-default tests and the duplicate hand-written mock dispatch/combine test; retain behavioral reference, capacity-tail, graph-replay, and real local-dispatch integration coverage.
• Consolidate benchmark timing, routing setup, expert distributions, FP32 reference checks, and TFLOPS calculation in a 167-line shared helper.
• Reduce the four benchmark entry points plus shared helper from roughly 915 to 778 lines, and remove the redundant pure/full compute benchmark while retaining the dynamic-capacity sweep and TileOps-vs-DeepGEMM comparison.

Validation

• 24 dispatched-expert and dispatch-adapter smoke tests pass.
• Normal large-row and device-bounded direct/explicit fused-gated tests pass.
• Local dispatch benchmark smoke passes.
• TileOps-vs-DeepGEMM M=1 smoke passes with FP32 correctness checks.
• DeepEP V2 2-rank fresh and cached-handle dispatch pass on 2× H200 using the non-Gin path.
• Ruff, Ruff format for benchmark files, Python compilation, and git diff --check pass.

Scope

• BF16 tight/no-pad expert compute only; FP8/FP4 and padded-direct compute remain out of scope.
• TileOps does not acquire a runtime DeepEP dependency; callers inject the DeepEP buffer.
• Capacity-tail values remain unspecified; consumers use the GPU-resident expert_offsets[-1:] valid-row view.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[Ibuki-wind]

Overall

One kernel-selection bug remains, and the draft needs a substantial size reduction before it is reviewable as a maintainable refactor.

Cross-cutting concerns

• The four new benchmark scripts repeat timing loops, tensor/model setup, output formatting, and correctness machinery across roughly 900 lines; consolidate the common harness and keep each entry point to scenario definitions plus backend adapters.

[Ibuki-wind]

The alignment fallback can be silently undone here and again for _gemm_down: after kernel_cls is changed to MoeGroupedGemmNopadKernel, a caller-supplied kernel_map['moe_grouped_gemm_kernel'] is expanded last and reinstates the misaligned 3WG kernel that the preceding check rejected. → Resolve one effective GEMM class first and either reject a conflicting override or build the sub-op maps without re-merging that key.

[Ibuki-wind]

This exact-signature introspection test (together with the adjacent dataclass-default checks) duplicates declarations instead of exercising behavior and will fail harmless parameter/default refactors. → Delete these contract-shape tests and keep the behavioral reference, capacity-tail, and graph-replay tests.

[Ibuki-wind]

This 100-line mock test never calls a dispatch or combine implementation: it manually builds the sorted batch, runs the expert op already covered above, then compares a manual weighted index_add_ against the same manual reference. → Delete it; the local-dispatch integration test in test_expert_dispatch.py covers the real composition with far less duplicate setup.

[Ibuki-wind]

Overall

The bounded-row kernel is a justified compiled specialization because valid_rows must remain device-side. Its source generation should be consolidated with the existing FusedGated strategies.

[Ibuki-wind]

These two bounded factories add another direct/explicit pair that repeats the same column mapping and gate/value/op_func/store template already owned by _make_fused_gated_direct and _make_fused_gated_explicit. The device-side row loop/guard is the real specialization; the pointwise strategy is not new. → Keep distinct compiled signatures (valid_rows must stay device-side), but use one source generator per strategy with a compile-time row policy: the normal policy emits the existing 2-D/flattened row mapping, and the bounded policy emits grid_rows, the persistent row loop, and row < valid_rows[0]. Alternatively, make one bounded generator parameterized by static num_per_thread (1 for direct). Do not add a runtime bounded/unbounded branch. Acceptance criterion: adding or changing the fused-gated epilogue or column mapping requires editing one source implementation, not four factories.

[RuneFang]

Does it now support the Megakernel MoE ?