feat(moe_ep): add SM90 push FP8 mega-MoE backend for Hopper

[leonardHONG]

📌 Description

This PR registers sm90_push_fp8, a whole-layer expert-parallel MoE backend for SM90 GPUs connected through single-node NVLink, as a fourth mega backend behind the existing register_mega_kernel plugin surface.

The public entry point is MoEEpLayer with Sm90PushFp8MegaMoeConfig. Existing backend behavior is unchanged.

The steady-state forward path covers dispatch, FP8 block-scale expert FFN, and combine without host-side synchronization, and supports CUDA Graph capture after initialization. The diagram below shows the path with all three optional optimizations enabled:

bump_tag
  -> wait_acks
  -> deduplicated dispatch
  -> wait_prefix
  -> compact
  -> FC1 with fused SwiGLU + 1x128 quantization
  -> FC2
  -> grouped combine with owner-side FP32 route reduction
  -> BF16 reduction of received owner rows
  -> ack

The backend provides three independently configurable optimizations. Validation and performance results below use all three unless stated otherwise.

• Deduplicated dispatch stores one payload row per token and destination rank instead of one per route. In the tested configurations, with all other options fixed, its output matched the per-route path under torch.equal. For K=6, EP=4, and random routing, it reduces dispatch payload bytes by approximately 43%.
• Grouped combine performs FP32 route reduction on the owner rank before pushing one quantized row per token. Its correctness gate compares the final output against the same reference used by the per-route path because the quantization point moves.
• Fused FC1 epilogue performs SwiGLU and 1x128 activation quantization inside the FC1 DeepGEMM epilogue, eliminating an approximately 1 GiB intermediate BF16 activation buffer per rank at the DSV3 EP8 shape. In the tested configurations, with all other options fixed, its output matched the unfused path under torch.equal. The unfused path remains the default and maintained reference; the fused path is opt-in.

Motivation

The existing Hopper EP stack composes NCCL all-to-all with a local fused-MoE operator, paying for additional launches, intermediate buffers, and quantization boundaries.

sm90_push_fp8 writes quantized payloads directly into peer-mapped symmetric memory and feeds received rows directly into the expert GEMM pipeline.

Contracts introduced by this PR

• capacity_factor bounds the GEMM, TMA, and scale buffers, rather than only the protocol window. The workspace query takes explicit expected_m and max_rows, so block-M tactics and cubin keys remain independent of the memory bound.

• Building the grouped GEMM requires CUDA Toolkit 12.8 or newer. The generator explicitly rejects older toolkits, and AOT skips this module below CUDA 12.8, allowing JIT-cache wheels to continue building on CUDA 12.6. Loading and running the generated module also requires a CUDA runtime version 12.8 or newer.

• Runner initialization elects one NVCC-capable rank per cache and tactic group to compile cold DeepGEMM cubins. Other ranks load the resulting cubins from disk.

  • A warm cache works without NVCC.
  • A cold cache without NVCC fails during initialization with an error explaining how to select the fallback.
  • TRTLLM_DG_ENABLED=0 selects the fixed-tactic CUTLASS fallback on the unfused path.

• Wait kernels use a %globaltimer deadline and publish a shared abort marker before trapping. Peers polling the same protocol window observe the marker and fail. Because a trap leaves a sticky CUDA error, the process launcher must terminate the full rank group from the CPU.

• This backend is a stateful variant of the mega-kernel contract:

  • construction binds transformed static weights;
  • stage_inputs pre-binds the caller's output tensor.

  The architecture and runbook documentation are updated accordingly, including narrowing the two-entry wording to buffer-oriented kernels. The MegaKernelBackend lifecycle itself is unchanged.

Correctness

Full kernel acceptance ran on an 8× H800 NVLink node:

Suite	Scale	Result
Single-GPU suite, 105 configurations	1 GPU	all passed
Distributed suites	EP2 / EP4 / EP8	all passed
Uneven token distributions, zero-token ranks, launch skew	EP2–EP8	all passed
CUDA Graph capture and replay	tested configurations	all passed
200-round oracle-checked soaks	every tested configuration	all passed
Eight-GPU acceptance matrix at DSV3 shapes	EP8	all passed, no skipped cases

The final PR tree was then revalidated across three Hopper variants:

Suite	Hardware	Result
Single-GPU suite, all configurations	1× H20	130 passed, 0 failed
torchrun EP2	2× H100 NVLink	15 passed, 0 failed
torchrun EP4	4× H20 NVLink	15 passed, 0 failed
compute-sanitizer memcheck / racecheck / initcheck, excluding trap and soaks	1× H20	Clean

Additional hardware-verified properties on the final PR tree:

• In the tested configurations, the fused FC1 output matched the unfused reference under torch.equal.
• Leader deduplication was verified by counting NVCC invocations against a shared cold cache: exactly one compilation per cubin across four ranks.
• Warm-cache execution without NVCC was exercised end to end.
• Cold-cache fail-fast behavior was exercised end to end.
• The CUTLASS fallback was exercised end to end.

Performance

Performance numbers were collected during kernel acceptance on the 8× H800 node. The final PR tree was functionally revalidated as listed above.

Setup

• H800 NVLink
• CUDA 12.9
• PyTorch 2.8
• barrier-aligned timing
• maximum latency across ranks
• autotuned baseline tactics
• identical correctness gates on both sides
• T denotes live tokens per rank

The tables report baseline latency / sm90_push_fp8 latency; values greater than 1× favor this backend. H, E, and K denote hidden size, expert count, and top-k routes. random distributes routes across experts, while hot1 concentrates routes on one expert.

Versus NCCL all-to-all with the same FP8 block-scale compute

Configuration	Decode T64	Throughput T2048
EP4 SMALL, H4096 E32 K6	1.1–1.5×	2.0–2.3×
EP4 DSV3, H7168 E256 K8	1.01×	2.0–2.3×

Against a wire-only NCCL transport reference that excludes compaction and combine reduction, the push transport measured 1.9–3.2× faster.

This number should be interpreted only as a comparison against that specific measured reference. It is not a claim against third-party EP libraries.

Versus NCCL + cutlass_fused_moe FP8

Configuration	Decode T64	T2048 random	T2048 hot1
EP4 DSV3	1.06×	2.3×	2.0×
EP8 DSV3	1.10×	2.2–2.4×	1.6×
EP4 SMALL	3.1×	2.1×	See limitations

At EP1 and equal FP8 precision, sm90_push_fp8 measured 2.3–3.4× faster than cutlass_fused_moe over T64–T2048.

EP8 scaling

Configuration:

DSV3
capacity = 2048
all three optimizations enabled

Results:

T64:    0.758 ms
T2048:  3.238 ms

With capacity fixed at 2048, a 32× increase in live token count increased latency by 4.27×.

Additional observations:

• no scheduling cliff near T128;
• all-remote routing adds approximately 5%;
• running 64 live tokens with buffers sized for 2048 adds approximately 0.7%.

Comparison with Pull Style

[From comment]

Ran the six suggested shapes on 8× H100 80GB SXM. These are back-to-back steady-state measurements with 20 warmups and 100 samples per path. The table reports p50 latency in milliseconds; the last column is Pull / Push, so values above 1 mean Push is faster.

Geometry	Tokens/rank	Pull	Push	Speedup
DeepSeek V3	8	0.572	0.614	0.93×
	64	0.657	0.704	0.93×
	512	1.699	1.065	1.60×
	2048	3.497	2.728	1.28×
	8192	12.935	9.886	1.31×
Kimi K2.6	8	0.716	0.749	0.96×
	64	0.913	0.949	0.96×
	512	2.038	1.163	1.75×
	2048	3.867	2.776	1.39×
	8192	13.023	10.046	1.30×
DeepSeek V4 Flash	8	0.386	0.410	0.94×
	64	0.436	0.468	0.93×
	512	0.871	0.581	1.50×
	2048	1.960	1.410	1.39×
	8192	5.752	4.609	1.25×
DeepSeek V4 Pro	8	0.951	0.966	0.98×
	64	1.301	1.294	1.01×
	512	2.660	1.479	1.80×
	2048	4.315	3.168	1.36×
	8192	14.172	10.757	1.32×

Push is about 4–5% slower at the smallest token counts, but becomes consistently faster from T512 onward, with gains of roughly 1.25–1.80×.

Correctness also passed across all measured cases (minimum cosine similarity 0.99816, maximum NRMSE 0.06063).

GPT-OSS 120B is not supported yet because its hidden dimensions are not 128-aligned, and Qwen3.5 397B uses top-k 10 while the current limit is 8.

Limitations

• Supports:

  • SM90;
  • single-node, peer-accessible NVLink;
  • protocol limit ep_size <= 32, with hardware validation in this PR covering up to EP8;
  • top_k in {1, 2, 4, 6, 8};
  • DeepSeek-style FP8 block scaling;
  • SwiGLU.

  Unsupported configurations raise an explicit error rather than silently falling back.

• The unfused FC1 path supports intermediate_size <= 16384; larger configurations require fuse_fc1_epilogue=True.

• At DSV3 decode shapes, the advantage over NCCL + cutlass_fused_moe is 1.06–1.10×, because computation dominates.

• In the artificial SMALL T2048 case where all K routes target one expert, the backend is 25% slower than the autotuned CUTLASS stack. Realistic hot1 routing remains faster at EP4 and EP8.

• The fused FC1 epilogue does not yet have a small-M swapAB variant and can be disabled independently below its break-even point.

• Round tags use uint32. Reuse after 2^32 forwards on one pipe is a documented limit.

• Custom raw-stream execution through bootstrap.stream is rejected explicitly rather than silently ignored.

How to run

Single-GPU and host-only suites

pytest \
  tests/moe_ep/test_sm90_push_fp8_kernel.py \
  tests/moe_ep/test_sm90_push_fp8_backend.py \
  tests/gemm/test_sm90_moe_gemm.py \
  tests/gemm/test_sm90_moe_gemm_contract.py

Multi-rank tests

Requires at least two SM90 GPUs and uses torch.distributed only.

bash tests/moe_ep/run_tests.sh sm90_push

Soak tests

SM90_PUSH_SOAK_ROUNDS=200 \
  bash tests/moe_ep/run_tests.sh sm90_push

Benchmark

torchrun \
  --standalone \
  --nproc-per-node=8 \
  benchmarks/bench_sm90_push_megamoe.py \
  --config DSV3 \
  --dedup \
  --grouped-combine \
  --fuse-fc1 \
  --assert-cos-min 0.997

🔍 Related issues

This PR contributes the Hopper FP8 block-scale milestone under #3692 and the SM90 sub-issue #3780, and follows the whole-layer integration direction described in #3704.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

🧪 Tests

• Tests added and updated.
• Single-GPU suite passes.
• EP2 distributed suite passes.
• EP4 distributed suite passes.
• CUDA Graph tests pass.
• Soak tests pass on the listed hardware.

Reviewer notes

Native NVFP4 and MXFP8 execution are out of scope because SM90 does not provide native block-scaled tensor-core instructions for those formats. A follow-up PR stacked on this one adds NVFP4 checkpoint support for SM90 on top of this backend, through online W4A8 decode kernels and a one-time requantization path that rides this FP8 backend unchanged.
The highest-value review areas are:

Entry points for the highest-risk areas:

• Symmetric-window ordering and the acknowledgement protocol:
  kernel_src/sm90_push_megamoe/src/a2a/sm90_push_a2a_ops.cu

• Grouped-combine numerical behavior, where the quantization point moves:
  The combine kernels in kernel_src/sm90_push_megamoe/src/a2a/sm90_push_a2a_ops.cu, gated by the reference comparison in tests/moe_ep/test_sm90_push_fp8_kernel.py

• DeepGEMM fused-epilogue and workspace contracts:
  src/fp8_gemm/fp8_moe_fc1_fused.cuh``shim/gemm.py

• CUDA 12.8 gating and the AOT skip:
  shim/gemm.py``flashinfer/aot.py

• Collective initialization failure handling:
  shim/protocol.py

Summary by CodeRabbit

• New Features

  • Added SM90 FP8 MegaMoE support with pull- and push-style execution paths.
  • Added configurable FP8 formats, scaling modes, routing, reduction, CUDA graph, and distributed execution options.
  • Added weight preprocessing and public configuration interfaces.
  • Added benchmark tools for correctness, performance, token sweeps, and peer-to-peer bandwidth.

• Documentation

  • Expanded architecture guidance, setup instructions, tuning information, and reproducibility runbooks.

• Tests

  • Added broad correctness, validation, packaging, lifecycle, distributed, and CUDA graph coverage.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

This PR adds a new SM90 push-style FP8 MegaMoE backend (protocol, GEMM, tests, benchmarks) and vendors two large CuTeDSL kernel source drops: an SM100 tree (MXFP8 and NVFP4 fused kernels) and an SM90 pull-style tree (FP8 GLU and NVFP4 kernels). Shared mega backends, docs, and packaging are updated to reference the new source layout.

Changes

SM90 Push FP8 MegaMoE Backend

Layer / File(s)	Summary
Backend contracts and lifecycle flashinfer/moe_ep/backends/mega/kernel/sm90_push_fp8/*, flashinfer/moe_ep/backends/mega/kernel/sm90_pull_fp8/*, flashinfer/moe_ep/core/..., flashinfer/moe_ep/modes/mega_layer.py	Adds the SM90 push and pull FP8 backend classes, config, weights, and staging modules. Updates core kernel/runtime/validation code and mega_layer output allocation order.
Shared mega backend updates flashinfer/moe_ep/backends/mega/kernel/{deep_gemm_mega,mxfp8_cutedsl,nvfp4_cutedsl}/*	Retargets shim imports to the new sm100 source path. Adds fused staging and output-optional compute support.
SM90 push protocol and GEMM sources flashinfer/moe_ep/kernel_src/sm90/push_style_megamoe/*	Adds the push protocol shim, runner, weight transforms, JIT generation, and CUDA A2A/GEMM kernel sources.
Correctness and contract validation tests/moe_ep/test_sm90_push_fp8_*, tests/moe_ep/test_moe_ep_*mega_multirank.py, tests/moe_ep/conftest.py	Adds CPU, distributed, graph-replay, soak, trap, and packaging tests. Adds a DeepGEMM cache isolation fixture.
Benchmarks, baseline, and reference paths benchmarks/bench_sm90_push_megamoe.py, benchmarks/sm90_push_megamoe_*.py	Adds a pure-Torch reference, a grouped-GEMM baseline, and an end-to-end benchmark with correctness gates.
Documentation and packaging wiring flashinfer/aot.py, pyproject.toml, .pre-commit-config.yaml, docs/design_docs/*	Updates AOT module generation, package data, pre-commit exclusions, and architecture/runbook documentation.

SM100 CuTeDSL MegaMoE Kernel Drop

Layer / File(s)	Summary
Shim adapters and docs flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/shim/*, *.md	Adds path bootstrap, comm, autotune, knob cache, quant staging, and tuner modules for the SM100 shim layer.
Common utilities .../src/common/*	Adds shared dtype constants, quantization helpers, and SwiGLU/PTX arithmetic used by both kernel variants.
MXFP8 GLU kernel .../src/moe_mxfp8_glu/*	Implements the MXFP8 epilogue, fused fc12 kernel, MegaMoE wrapper, reference/runner, and test scripts.
NVFP4 swap-AB kernel .../src/moe_nvfp4_swapab/*	Implements the NVFP4 contract, scheduler, dynamic mainloop, epilogue, fused kernel, and topk reduce.
Shared low-level infrastructure .../src/src/*	Adds the symmetric-heap bootstrap, cleanup kernel, config, dispatch/token-comm kernels, and PTX helpers.

SM90 Pull-Style CuTeDSL MegaMoE Kernel Drop

Layer / File(s)	Summary
Shim adapters and docs flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/{__init__.py,shim/*,*.md}	Adds path bootstrap, comm, and lazy kernel-helper exports for the SM90 pull shim.
Common utilities .../src/common/*	Adds shared dtype constants, quantization/comparison helpers, and SwiGLU/PTX arithmetic.
Hopper FP8 GLU kernel .../src/moe_hopper_fp8/*	Implements the FP8 lazy-compile frontend, epilogues, fused fc12 kernels, MegaMoE wrapper, reference/runner, and benchmark/test tooling.
NVFP4 swap-AB kernel .../src/moe_nvfp4_swapab/*	Implements the NVFP4 swap-AB epilogue, scheduler, dynamic mainloop, fused kernel, topk reduce, and reference/runner/tests.
Shared low-level infrastructure .../src/src/*	Adds the symmetric-heap bootstrap, cleanup kernel, config, dispatch/token-comm kernels, and PTX helpers.

Estimated code review effort: 5 (Critical) | ~180 minutes

Suggested labels: run-ci, op: moe

Suggested reviewers: aleozlx, anerudhan, iwakurarein

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 28.85% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the new SM90 push FP8 MegaMoE backend for Hopper and matches the primary changes.
Description check	✅ Passed	The description is detailed, covers objectives, contracts, validation, performance, limitations, related issues, and testing evidence.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.}

[gemini-code-assist]

Code Review

This pull request introduces the 'sm90_push_fp8' MegaMoE backend, which provides an optimized FP8 communication and computation path for Hopper (SM90) GPUs. The changes include new kernel sources, Python shims for JIT compilation, and comprehensive correctness tests. I have reviewed the code and identified two issues: a potential cross-device link failure in the JIT cache renaming logic and a strict aliasing violation in the tag increment kernel. Both issues have actionable suggestions provided in the comments.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

In containerized or multi-partition environments (such as Docker or cluster nodes), the temporary directory (e.g., /tmp) and the cache directory (e.g., ~/.tensorrt_llm/cache) are often located on different filesystems. In such cases, std::filesystem::rename will fail with EXDEV (Invalid cross-device link), causing the compilation to fail completely. To make this robust, add a fallback to std::filesystem::copy_file and then remove the temporary directory if rename fails.

    try {
      std::filesystem::rename(tmp_cubin_path, cubin_path);
    } catch (std::filesystem::filesystem_error const&) {
      if (!std::filesystem::exists(cubin_path)) {
        try {
          std::filesystem::copy_file(tmp_cubin_path, cubin_path, std::filesystem::copy_options::overwrite_existing);
        } catch (...) {
          std::filesystem::remove_all(tmp_path);
          throw;
        }
      }
    }

[leonardHONG]

Both directories are under the same cache root, so this does not cross filesystems. The atomic rename is intentional to avoid exposing a partially written cubin to another rank.

[gemini-code-assist]

Reinterpreting int32_t* as uint32_t* and dereferencing it violates the strict aliasing rule in C++, which is undefined behavior (UB). This can lead to compiler optimizations that misbehave or omit the write entirely. Instead of using reinterpret_cast, you can safely perform the increment by casting the dereferenced value to uint32_t, incrementing it, and casting it back to int32_t before writing it back.

  if (threadIdx.x == 0 && blockIdx.x == 0) {
    *round_ctr = static_cast<int32_t>(static_cast<uint32_t>(*round_ctr) + 1u);
  }

[leonardHONG]

Access through the matching unsigned type is valid C++. We use it intentionally so the tag wraps modulo 2^32 without signed-overflow UB.

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (2)

tests/moe_ep/run_tests.sh (1)

111-127: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Optional: run_sm90_push is gated behind require_nccl_ep inside run_multirank.

run_multirank returns early when require_nccl_ep fails (Line 112), so the SM90 push suite—which uses the push backend, not nccl_ep—won't run in the multirank flow on boxes without nccl_ep. The standalone sm90_push target still covers it, but consider running run_sm90_push before the nccl_ep gate if you want it exercised in all/multirank regardless.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/moe_ep/run_tests.sh` around lines 111 - 127, Move the run_sm90_push
invocation in run_multirank before the require_nccl_ep early-return gate so the
SM90 push suite runs regardless of NCCL EP availability. Preserve the existing
return-code aggregation by setting rc on failure, then run the NCCL-dependent
suites only after the gate succeeds.

flashinfer/moe_ep/kernel_src/sm90_push_megamoe/__init__.py (1)

19-33: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Optional: __all__ is not sorted (Ruff RUF022). Both new re-export modules declare __all__ in insertion order; apply isort-style sorting to satisfy the lint rule and keep the two surfaces consistent.

• flashinfer/moe_ep/kernel_src/sm90_push_megamoe/__init__.py#L19-L33: sort the __all__ entries alphabetically.
• flashinfer/moe_ep/kernel_src/sm90_push_megamoe/shim/__init__.py#L17-L31: sort the __all__ entries alphabetically.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@flashinfer/moe_ep/kernel_src/sm90_push_megamoe/__init__.py` around lines 19 -
33, Sort the __all__ entries alphabetically in both
flashinfer/moe_ep/kernel_src/sm90_push_megamoe/__init__.py (lines 19-33) and
flashinfer/moe_ep/kernel_src/sm90_push_megamoe/shim/__init__.py (lines 17-31),
preserving all existing exports and applying the same isort-style ordering to
both modules.

Source: Linters/SAST tools

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/moe_ep/_sm90_push_fp8_baseline.py`:
- Around line 28-63: Update quant_weights so each _weight_cache entry retains
the source w13 and w2 tensors alongside the quantized result, preventing
data_ptr reuse from returning stale weights. On cache hits, return the cached
result at the entry’s first element, and store the source tensors with the
result when populating the cache.

In `@tests/moe_ep/test_sm90_push_fp8_backend.py`:
- Around line 307-325: Update the capacity assertion in
test_public_ep1_forward_validation_and_capacity to match the token_capacity
wording emitted by validate_forward, replacing the max_tokens_per_rank
expectation while leaving the dtype checks unchanged.

---

Nitpick comments:
In `@flashinfer/moe_ep/kernel_src/sm90_push_megamoe/__init__.py`:
- Around line 19-33: Sort the __all__ entries alphabetically in both
flashinfer/moe_ep/kernel_src/sm90_push_megamoe/__init__.py (lines 19-33) and
flashinfer/moe_ep/kernel_src/sm90_push_megamoe/shim/__init__.py (lines 17-31),
preserving all existing exports and applying the same isort-style ordering to
both modules.

In `@tests/moe_ep/run_tests.sh`:
- Around line 111-127: Move the run_sm90_push invocation in run_multirank before
the require_nccl_ep early-return gate so the SM90 push suite runs regardless of
NCCL EP availability. Preserve the existing return-code aggregation by setting
rc on failure, then run the NCCL-dependent suites only after the gate succeeds.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 3c1c788f-1baf-46c3-8f7c-ba053d6f9bed

📥 Commits

Reviewing files that changed from the base of the PR and between c83607a and 0e01918ffb2b4189baf26cbb4deda5016d6987e7.

📒 Files selected for processing (47)

• benchmarks/bench_sm90_push_megamoe.py
• benchmarks/sm90_push_megamoe_baseline.py
• benchmarks/sm90_push_megamoe_reference.py
• docs/design_docs/moe_ep_architecture.md
• docs/design_docs/moe_ep_runbook.md
• flashinfer/aot.py
• flashinfer/comm/mnnvl.py
• flashinfer/moe_ep/__init__.py
• flashinfer/moe_ep/backends/mega/kernel/__init__.py
• flashinfer/moe_ep/backends/mega/kernel/deep_gemm_mega/backend.py
• flashinfer/moe_ep/backends/mega/kernel/mxfp8_cutedsl/backend.py
• flashinfer/moe_ep/backends/mega/kernel/nvfp4_cutedsl/backend.py
• flashinfer/moe_ep/backends/mega/kernel/sm90_push_fp8/__init__.py
• flashinfer/moe_ep/backends/mega/kernel/sm90_push_fp8/backend.py
• flashinfer/moe_ep/backends/mega/kernel/sm90_push_fp8/config.py
• flashinfer/moe_ep/backends/mega/kernel/sm90_push_fp8/staging.py
• flashinfer/moe_ep/backends/mega/kernel/sm90_push_fp8/weights.py
• flashinfer/moe_ep/core/kernel/base.py
• flashinfer/moe_ep/kernel_src/sm90_push_megamoe/ACKNOWLEDGEMENT.md
• flashinfer/moe_ep/kernel_src/sm90_push_megamoe/__init__.py
• flashinfer/moe_ep/kernel_src/sm90_push_megamoe/shim/__init__.py
• flashinfer/moe_ep/kernel_src/sm90_push_megamoe/shim/gemm.py
• flashinfer/moe_ep/kernel_src/sm90_push_megamoe/shim/jit.py
• flashinfer/moe_ep/kernel_src/sm90_push_megamoe/shim/protocol.py
• flashinfer/moe_ep/kernel_src/sm90_push_megamoe/shim/runner.py
• flashinfer/moe_ep/kernel_src/sm90_push_megamoe/shim/weights.py
• flashinfer/moe_ep/kernel_src/sm90_push_megamoe/src/a2a/sm90_push_a2a.cuh
• flashinfer/moe_ep/kernel_src/sm90_push_megamoe/src/a2a/sm90_push_a2a_ops.cu
• flashinfer/moe_ep/kernel_src/sm90_push_megamoe/src/fp8_gemm/fp8_moe_binding.cu
• flashinfer/moe_ep/kernel_src/sm90_push_megamoe/src/fp8_gemm/fp8_moe_fc1_fused.cuh
• flashinfer/moe_ep/kernel_src/sm90_push_megamoe/src/fp8_gemm/fp8_moe_jit.cuh
• flashinfer/moe_ep/kernel_src/sm90_push_megamoe/src/fp8_gemm/fp8_moe_launcher.cuh
• flashinfer/moe_ep/kernel_src/sm90_push_megamoe/src/fp8_gemm/fp8_moe_scheduler.cuh
• flashinfer/moe_ep/modes/mega_layer.py
• pyproject.toml
• tests/conftest.py
• tests/gemm/test_sm90_moe_gemm.py
• tests/gemm/test_sm90_moe_gemm_contract.py
• tests/moe_ep/_sm90_push_fp8_baseline.py
• tests/moe_ep/_sm90_push_fp8_reference.py
• tests/moe_ep/run_tests.sh
• tests/moe_ep/test_mega_layer_validation.py
• tests/moe_ep/test_sm90_push_fp8_backend.py
• tests/moe_ep/test_sm90_push_fp8_backend_cpu.py
• tests/moe_ep/test_sm90_push_fp8_kernel.py
• tests/moe_ep/test_sm90_push_fp8_orchestrator.py
• tests/moe_ep/test_sm90_push_fp8_packaging.py

[mhoqueanik]

@leonardHONG Thank you for the PR! let me run it on my end and get back to you!

[mhoqueanik]

@leonardHONG My apologies for the delay! I am still restructuring moe_ep to adapt new dtypes and arch (+creating vLLM PR of fi moe_ep). Tentatively will process your PR and this #4113 by mid-next week

[mhoqueanik]

@leonardHONG I was able to reproduce your perf on my end. The numbers look great! Can you resolve the conflicts, and rebase it against this PR: #4113? That way we can cleanly merge it

[leonardHONG]

Done, rebased onto the latest main and cleaned up the remaining coderabbit comments. Thanks!

[mhoqueanik]

Overall looks good to me. We can proceed with the merge once the LICENSE/ACKNOWLEDGE.md issue is clear.

[Anerudhan]

/bot run tests/moe_ep

[flashinfer-bot]

GitLab MR !1181 has been created, and the CI pipeline #62029339 is currently running. I'll report back once the pipeline job completes.

[flashinfer-bot]

[SUCCESS] Pipeline #62029339: 18/18 executed test jobs passed

[huangzhilin-hzl]

Hi @leonardHONG， Thanks for the great work! Is there a benchmark comparing the pull- and push-style backends under the same setup? Which workloads currently favor each style?

[leonardHONG]

Hi @leonardHONG， Thanks for the great work! Is there a benchmark comparing the pull- and push-style backends under the same setup? Which workloads currently favor each style?

yep, i posted the 8× h100 comparison [above]. pull is a little better for small token counts, while push wins on larger workloads. push currently needs 128-aligned dimensions and top-k ≤ 8, but i may support more shapes later.