feat(moe_ep): SM90 (Hopper) pull-style FP8 mega-MoE backend - #4113
Merged
Anerudhan merged 58 commits intoAug 8, 2026
Merged
Conversation
mhoqueanik
requested review from
Anerudhan,
IwakuraRein,
aleozlx,
bkryu,
cyx-6,
dhiraj113,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
qiching,
qsang-nv,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
July 23, 2026 20:28
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Contributor
|
Note Reviews pausedIt 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 Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR vendors two large CuTeDSL MegaMoE fused-kernel drops under new architecture-scoped paths ( ChangesCore MoE EP runtime and backend wiring
Estimated code review effort: 5 (Critical) | ~150+ minutes SM100 CuTeDSL MegaMoE vendored kernel drop
SM90 pull-style CuTeDSL MegaMoE vendored kernel drop
Test suite and CI updates
Estimated code review effort: 5 (Critical) | ~150+ minutes Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 15
🧹 Nitpick comments (6)
flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/iket_compat.py (1)
15-16: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAvoid an unconditional
By the surrounding comment's own reasoning, the public-release / CTK-12.9 CI wheels fall into this branch, so this line emits
!!!! Iket is not enabled !!!!to stdout on every import in the common configuration. Preferwarnings.warn(...)(or a module logger), which is filterable and doesn't pollute stdout of programs importing the kernel tree.♻️ Suggested change
try: from cutlass.cute import iket # type: ignore except (ImportError, NotImplementedError): - print("!!!! Iket is not enabled !!!!") + import warnings + + warnings.warn( + "cutlass IKET dialect unavailable; using no-op IKET markers.", + stacklevel=2, + )🤖 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/pull_style_cutedsl_megakernel/src/src/iket_compat.py` around lines 15 - 16, Replace the unconditional print in the Iket import-exception handler with a filterable warnings.warn call or module logger, preserving the existing ImportError and NotImplementedError handling without writing to stdout during module import.flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/contract.py (1)
311-374: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueConsider caching the normalized table to avoid repeated re-enumeration.
tablerecomputesmapping.normalize(...)on every access. ForFunctionMapping-backed contracts each call re-enumerates the full domain and re-invokes the user function, andassert_equivalent_toreads.tableseveral times (viais_equivalent_toplus the details block), multiplying that cost. SinceContractis frozen and already validated in__post_init__, caching the result once avoids the redundant passes.♻️ Optional: cache the normalized table
def __post_init__(self) -> None: - # Validate eagerly so malformed contracts fail at construction time. - self.mapping.normalize(domain=self.domain, codomain=self.codomain) + # Validate eagerly so malformed contracts fail at construction time, + # and cache the canonical table (Contract is frozen/immutable). + object.__setattr__( + self, + "_table", + self.mapping.normalize(domain=self.domain, codomain=self.codomain), + ) `@property` def table(self) -> tuple[int, ...]: - return self.mapping.normalize(domain=self.domain, codomain=self.codomain) + return self._table🤖 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/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/contract.py` around lines 311 - 374, Cache the normalized mapping table during Contract initialization and have the table property return that cached value instead of calling mapping.normalize repeatedly. Update the frozen dataclass initialization in __post_init__ safely, preserving eager validation and the existing behavior of rename_domain, rename_codomain, is_equivalent_to, and assert_equivalent_to.flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_mxfp8_glu/kernel_mxfp8_glu_fc12.py (2)
572-581: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winLeftover debug
_compute_stages.The
print(...)runs on every kernel setup and is a debug artifact; consider dropping it or gating behind a debug flag. Separately, the signature is annotated-> Tuple[int, int, int]but the function returns a 4-tuple (num_acc_stage, num_a_stage, num_b_stage, num_sched_stages).🤖 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/sm100/cutedsl_megamoe/src/moe_mxfp8_glu/kernel_mxfp8_glu_fc12.py` around lines 572 - 581, Remove the unconditional debug print from _compute_stages, and update its return annotation to describe the four integers it returns: num_acc_stage, num_a_stage, num_b_stage, and num_sched_stages.
1768-1801: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDead debug block still issues GMEM loads.
The
cute.printf(...)is commented out, butcounter_val_postandfc1_first_i32are still loaded from GMEM on the fc2 A-side path (fortile_n_idx == 0) and never consumed. Remove the whole guarded block so the probe loads don't linger in the compiled TMA-A warp.🤖 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/sm100/cutedsl_megamoe/src/moe_mxfp8_glu/kernel_mxfp8_glu_fc12.py` around lines 1768 - 1801, Remove the entire guarded debug block keyed by tma_a_warp_id and tile_n_idx == 0, including the counter_val_post and fc1_first_i32 GMEM loads and related offset/pointer calculations. Leave the surrounding fc2 A-side spin logic unchanged.flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/ptx_helpers.py (1)
530-539: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
_fence_rel_gpudocstring says "system scope" but the fence uses device scope. Copy-paste from_fence_rel_sys; update to "device scope" to avoid confusion.🤖 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/sm100/cutedsl_megamoe/src/src/ptx_helpers.py` around lines 530 - 539, Update the docstring of _fence_rel_gpu to describe the fence as having acquire-release semantics at device scope, matching its llvm.fence syncscope="device" implementation.flashinfer/moe_ep/core/kernel/base.py (1)
224-243: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winHardcoded sm100-specific module path in the shared base-class
destroy().
destroy()only knows how to evictkernel_src.sm100.cutedsl_megamoe.shim.quant_stage's staged-token memo. If sm90_pull_fp8 (or a future sm120 tree) adds an analogous tensor-data_ptr()-keyed staging memo, this generic hook silently no-ops for it (sys.modules.getreturnsNone), reopening the stale-pointer-aliasing hazard on symmetric-heap address reuse that this exact code was written to avoid for sm100. Consider extracting this into a small overridable hook so each backend owns its own memo-eviction logic instead of the shared base class special-casing one architecture's shim path.Based on learnings, "ensure that every tensor whose
data_ptr()is included in the key also participates in cache eviction/invalidation... Verify the cache invalidation triggers when those key tensors go out of scope".♻️ Proposed hook-based refactor
- def destroy(self, workspace: Any) -> None: - """Release durable workspace resources (pool-aware, refcounted).""" - if workspace is None: - return - from .workspace_pool import release_workspace - - if release_workspace(workspace): - # The fused-stage memos key on topk_idx.data_ptr(); the symmetric - # heap reuses freed addresses, so evict before the buffer dies. - # sys.modules lookup (not an import): if the shim was never - # loaded, no memo exists and the heavy import must not happen. - import sys - - quant_stage = sys.modules.get( - "flashinfer.moe_ep.kernel_src.sm100.cutedsl_megamoe.shim.quant_stage" - ) - topk_idx = getattr(workspace, "topk_idx", None) - if quant_stage is not None and topk_idx is not None: - quant_stage.forget_staged_tokens(topk_idx) - workspace.destroy() + def destroy(self, workspace: Any) -> None: + """Release durable workspace resources (pool-aware, refcounted).""" + if workspace is None: + return + from .workspace_pool import release_workspace + + if release_workspace(workspace): + self._evict_staged_memo(workspace) + workspace.destroy() + + def _evict_staged_memo(self, workspace: Any) -> None: + """Backend hook: evict any tensor-keyed staging memo before the + symmetric heap reuses this workspace's addresses. No-op by default. + """🤖 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/core/kernel/base.py` around lines 224 - 243, Replace the sm100-specific module lookup in the shared destroy method with an overridable backend hook, invoked after release_workspace succeeds and before workspace.destroy(). Move sm100 staged-token eviction into that backend-specific override, ensuring every tensor used in the memo key is invalidated when the workspace is released; leave the base implementation architecture-agnostic and allow future backends such as sm90 or sm120 to provide their own eviction logic.Source: Learnings
🤖 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 `@benchmarks/bench_moe_ep_sm90_mega.py`:
- Around line 428-437: Replace both cleanup try/except/pass blocks in the
benchmark teardown with contextlib.suppress, covering
bench_backend.destroy(bench_workspace) and layer.destroy(). Add the contextlib
import alongside the existing standard-library imports.
In `@flashinfer/moe_ep/backends/mega/kernel/sm90_pull_fp8/weights.py`:
- Line 23: Remove the unused Optional symbol from the typing import in
weights.py, while preserving the remaining TYPE_CHECKING, Literal, and Tuple
imports.
- Around line 319-334: Annotate fc1_sf_shape and fc2_sf_shape before the
blockwise conditional as variable-length integer shape tuples, such as
tuple[int, ...], so both the 3-dimensional blockwise assignments and
2-dimensional non-blockwise assignments type-check without changing their
values.
In
`@flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/common/megamoe_constants.py`:
- Around line 19-25: Remove the duplicate second definitions of
Nvfp4E2M1RcpLimit, Fp8E4M3RcpLimit, and Fp8E5M2RcpLimit in megamoe_constants.py,
preserving the first reciprocal-limit block unchanged.
In
`@flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/mega_reference.py`:
- Line 97: Update the return annotation of reference_expert_fc12 to declare a
four-element tuple matching its four returned tensors: fc2_fp32, fc1_q,
fc1_sf_out, and fc1_fp32. Keep the return values and caller unpacking unchanged.
In
`@flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/mega_runner.py`:
- Around line 2758-2761: Replace the ValueError message in the num_topk
validation with a clear, professional diagnostic stating that values above 32
are unsupported by the current implementation; remove the profanity and
incomplete wording while preserving the existing validation behavior.
In
`@flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_common.py`:
- Around line 567-582: Document and enforce the supported 32-bit offset range in
_rcp_approx_kernel and _swiglu_pair_kernel. Add a small guard or assertion that
rejects n_elements at or beyond the signed 32-bit limit, and clearly comment
that these kernels intentionally use 32-bit offsets and do not support
~2^31-element tensors.
In
`@flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_fc12_common.py`:
- Around line 111-130: The hidden-size validation in ProblemDesc.__post_init__
uses Nvfp4BlockSize for every kind. Select the activation block size based on
self.kind, using Nvfp4BlockSize for nvfp4 and Mxfp8BlockSize for MXFP8 kinds,
then validate hidden against that selected size while preserving the existing
positive-multiple error behavior.
In `@flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/iket_compat.py`:
- Line 19: Replace the import-time print in the Iket compatibility fallback with
warnings.warn or remove it entirely, preserving the no-op shim behavior and
avoiding stdout output on common imports. Update only the fallback notification
around the “Iket is not enabled” message.
In `@flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/reference.py`:
- Around line 15-18: Update the module header metadata-packing documentation to
describe the 8-byte i64 layout implemented by _pack_metadata and
TokenSrcMetadata: low 32 bits are src_token, while high 32 bits contain
(src_rank << 16) | src_topk. Replace the stale 12-byte field description and
ensure the padding sentinel behavior remains documented.
In
`@flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/shim/hopper_fp8.py`:
- Around line 924-933: Rename the `shape` loop variable in the `scale_checks`
iteration to a distinct name, and update the associated shape comparison and
error message in that block. Leave the earlier `weight_checks` loop binding
unchanged.
- Line 1720: Update the FP8_SCALE_MODE initialization and the call to
create_dummy_inputs so the environment value is validated or narrowed to the
accepted Literal["per_tensor", "blockwise"] values before being passed, while
preserving "per_tensor" as the default and rejecting or handling unsupported
values explicitly.
In
`@flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/epilogue_refactor.py`:
- Around line 464-475: Update _amax_lane for sf_vec_size < 32 so each 16-lane
half-warp computes its own amax: use a converged warp reduction with lane-masked
values for the second half rather than calling warp_redux_sync only inside if
not first_half. Preserve the existing full-warp reduction when sf_vec_size == 32
and ensure both half-warp reductions execute uniformly across participating
lanes.
In
`@flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/mega_reference.py`:
- Line 97: Update the return annotation of reference_expert_fc12 to describe a
4-tensor tuple, matching its four returned values and both call-site unpacking
patterns.
In `@tests/moe_ep/test_moe_ep_sm90_pull_fp8_mega_multirank.py`:
- Around line 578-579: Add the gpu_4 pytest marker to
test_sm90_pull_fp8_preprocess_mega_weights_from_bf16 so it is selected by the
mega_sm90 multirank filter alongside arch_hopper. For
test_sm90_pull_fp8_mega_kernel_is_registered, either add gpu_4 and arch_hopper
markers or remove the redundant test, since registry coverage already exists in
test_registry_resolves_backend.
---
Nitpick comments:
In `@flashinfer/moe_ep/core/kernel/base.py`:
- Around line 224-243: Replace the sm100-specific module lookup in the shared
destroy method with an overridable backend hook, invoked after release_workspace
succeeds and before workspace.destroy(). Move sm100 staged-token eviction into
that backend-specific override, ensuring every tensor used in the memo key is
invalidated when the workspace is released; leave the base implementation
architecture-agnostic and allow future backends such as sm90 or sm120 to provide
their own eviction logic.
In
`@flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_mxfp8_glu/kernel_mxfp8_glu_fc12.py`:
- Around line 572-581: Remove the unconditional debug print from
_compute_stages, and update its return annotation to describe the four integers
it returns: num_acc_stage, num_a_stage, num_b_stage, and num_sched_stages.
- Around line 1768-1801: Remove the entire guarded debug block keyed by
tma_a_warp_id and tile_n_idx == 0, including the counter_val_post and
fc1_first_i32 GMEM loads and related offset/pointer calculations. Leave the
surrounding fc2 A-side spin logic unchanged.
In
`@flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/contract.py`:
- Around line 311-374: Cache the normalized mapping table during Contract
initialization and have the table property return that cached value instead of
calling mapping.normalize repeatedly. Update the frozen dataclass initialization
in __post_init__ safely, preserving eager validation and the existing behavior
of rename_domain, rename_codomain, is_equivalent_to, and assert_equivalent_to.
In `@flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/ptx_helpers.py`:
- Around line 530-539: Update the docstring of _fence_rel_gpu to describe the
fence as having acquire-release semantics at device scope, matching its
llvm.fence syncscope="device" implementation.
In
`@flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/iket_compat.py`:
- Around line 15-16: Replace the unconditional print in the Iket
import-exception handler with a filterable warnings.warn call or module logger,
preserving the existing ImportError and NotImplementedError handling without
writing to stdout during module import.
🪄 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 Plus
Run ID: d01bd6f9-7d64-4039-91b5-cb6f1c095325
📥 Commits
Reviewing files that changed from the base of the PR and between 0f5bb82 and b6d6858caf72815de998b4ac4cac57bb83790950.
📒 Files selected for processing (187)
.pre-commit-config.yamlCLAUDE.mdbenchmarks/bench_moe_ep_sm90_mega.pydocs/design_docs/moe_ep_architecture.mddocs/design_docs/moe_ep_runbook.mdflashinfer/moe_ep/__init__.pyflashinfer/moe_ep/backends/mega/kernel/__init__.pyflashinfer/moe_ep/backends/mega/kernel/deep_gemm_mega/backend.pyflashinfer/moe_ep/backends/mega/kernel/deep_gemm_mega/staging.pyflashinfer/moe_ep/backends/mega/kernel/deep_gemm_mega/weights.pyflashinfer/moe_ep/backends/mega/kernel/mxfp8_cutedsl/backend.pyflashinfer/moe_ep/backends/mega/kernel/mxfp8_cutedsl/config.pyflashinfer/moe_ep/backends/mega/kernel/mxfp8_cutedsl/staging.pyflashinfer/moe_ep/backends/mega/kernel/mxfp8_cutedsl/weights.pyflashinfer/moe_ep/backends/mega/kernel/nvfp4_cutedsl/backend.pyflashinfer/moe_ep/backends/mega/kernel/nvfp4_cutedsl/config.pyflashinfer/moe_ep/backends/mega/kernel/nvfp4_cutedsl/staging.pyflashinfer/moe_ep/backends/mega/kernel/nvfp4_cutedsl/weights.pyflashinfer/moe_ep/backends/mega/kernel/sm90_pull_fp8/__init__.pyflashinfer/moe_ep/backends/mega/kernel/sm90_pull_fp8/backend.pyflashinfer/moe_ep/backends/mega/kernel/sm90_pull_fp8/config.pyflashinfer/moe_ep/backends/mega/kernel/sm90_pull_fp8/staging.pyflashinfer/moe_ep/backends/mega/kernel/sm90_pull_fp8/weights.pyflashinfer/moe_ep/core/kernel/base.pyflashinfer/moe_ep/core/kernel/workspace_pool.pyflashinfer/moe_ep/core/runtime/__init__.pyflashinfer/moe_ep/core/runtime/bootstrap.pyflashinfer/moe_ep/core/validation/common.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/_paths.pyflashinfer/moe_ep/kernel_src/sm100/__init__.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/ACKNOWLEDGEMENT.mdflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/SKILL.mdflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/TUNING.mdflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/__init__.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/shim/__init__.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/shim/_paths.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/shim/autotune.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/shim/comm.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/shim/correctness.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/shim/kernel_helpers.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/shim/knob_cache.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/shim/mxfp8.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/shim/nvfp4.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/shim/quant_stage.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/shim/tuner.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/common/__init__.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/common/host_utils.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/common/megamoe_constants.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/common/moe_utils.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_mxfp8_glu/__init__.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_mxfp8_glu/epilogue_mxfp8.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_mxfp8_glu/kernel_mxfp8_glu_fc12.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_mxfp8_glu/mega_reference_mxfp8.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_mxfp8_glu/mega_runner.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_mxfp8_glu/megamoe_kernel_mxfp8.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_mxfp8_glu/run_functional_tests.shflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_mxfp8_glu/run_mega_tests.shflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_mxfp8_glu/runner_common.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_mxfp8_glu/runner_fc12.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/__init__.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/benchmark_p2p.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/contract.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/custom_ext.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/dynamic_mainloop.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/epilogue.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/epilogue_refactor.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/fc1_fc2_fuse_sched.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/kernel_fc12.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/mega_reference.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/mega_runner.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/megamoe_kernel.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/moe_persistent_scheduler.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/moe_utils.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/run_functional_tests.shflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/run_mega_tests.shflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_common.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_fc12.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_fc12_common.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/simulate_fc1_fc2_sched.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/topk_reduce.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/__init__.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/bootstrap.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/cleanup_kernel.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/config.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/dispatch_kernel.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/flag_batch.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/grid_sync.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/iket_compat.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/inputs_process.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/ptx_helpers.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/reference.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/sf_swizzle.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/sym_buffer.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/token_comm.pyflashinfer/moe_ep/kernel_src/sm90/__init__.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/SKILL.mdflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/TUNING.mdflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/__init__.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/shim/__init__.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/shim/_paths.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/shim/comm.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/shim/hopper_fp8.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/shim/kernel_helpers.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/common/__init__.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/common/host_utils.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/common/megamoe_constants.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/common/moe_utils.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/__init__.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/benchmark_requirements.txtflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/epilogue_fp8.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/epilogue_fp8_common.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/epilogue_fp8_swapab.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/hopper_moe_utils.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/kernel_fp8_glu_fc12.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/kernel_fp8_glu_fc12_swapab.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/mega_reference_fp8.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/mega_runner.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/megamoe_kernel_fp8.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/plot_token_sweep.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/run_functional_tests.shflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/run_mega_tests.shflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/run_perf_test.shflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/run_token_sweep_benchmark.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/run_token_sweep_benchmark.shflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/runner_fc12.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/summarize_token_sweep.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/__init__.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/benchmark_p2p.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/contract.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/custom_ext.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/cute_ref_ops.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/dynamic_mainloop.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/epilogue.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/epilogue_refactor.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/fc1_fc2_fuse_sched.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/kernel_fc12.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/mega_reference.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/mega_runner.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/megamoe_kernel.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/moe_persistent_scheduler.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/moe_utils.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/run_functional_tests.shflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/run_mega_tests.shflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/runner_common.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/runner_fc12.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/runner_fc12_common.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/simulate_fc1_fc2_sched.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/topk_reduce.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/__init__.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/bootstrap.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/cleanup_kernel.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/config.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/dispatch_kernel.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/flag_batch.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/grid_sync.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/iket_compat.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/inputs_process.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/ptx_helpers.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/reference.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/sf_swizzle.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/sym_buffer.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/token_comm.pyflashinfer/moe_ep/layer.pyflashinfer/moe_ep/modes/mega_layer.pyflashinfer/moe_ep/modes/split_layer.pyflashinfer/moe_ep/tune.pyflashinfer/moe_ep/weights.pypyproject.tomltests/conftest.pytests/moe_ep/run_tests.shtests/moe_ep/smoke_nixl_ep.pytests/moe_ep/test_fused_quant_stage.pytests/moe_ep/test_knob_cache.pytests/moe_ep/test_layer_single_gpu.pytests/moe_ep/test_mega_cuda_graph.pytests/moe_ep/test_mega_cuda_graph_multirank.pytests/moe_ep/test_mega_layer_validation.pytests/moe_ep/test_moe_ep_deep_gemm_skew_determinism.pytests/moe_ep/test_moe_ep_mxfp8_cutedsl_mega_multirank.pytests/moe_ep/test_moe_ep_nvfp4_cutedsl_mega_multirank.pytests/moe_ep/test_moe_ep_sm90_pull_fp8_mega_multirank.pytests/moe_ep/test_mxfp8_cutedsl_preprocess_vs_reference.pytests/moe_ep/test_nvfp4_cutedsl_kernel_vs_reference.pytests/moe_ep/test_sm90_pull_fp8_config.pytests/moe_ep/test_sm90_pull_fp8_kernel_vs_reference.pytests/moe_ep/test_weight_pack_union.pytests/moe_ep/test_workspace_pool.py
💤 Files with no reviewable changes (1)
- flashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/_paths.py
Contributor
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's internal server error or limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
🛑 Comments failed to post (15)
benchmarks/bench_moe_ep_sm90_mega.py (1)
428-437: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
pre-commit failure: replace
try/except/passwithcontextlib.suppress.ruff
SIM105fails on both blocks. Addimport contextlibat the top and suppress.Proposed fix
- if bench_backend is not None and bench_workspace is not None: - try: - bench_backend.destroy(bench_workspace) - except Exception: # noqa: BLE001 - pass - if layer is not None: - try: - layer.destroy() - except Exception: # noqa: BLE001 - pass + if bench_backend is not None and bench_workspace is not None: + with contextlib.suppress(Exception): + bench_backend.destroy(bench_workspace) + if layer is not None: + with contextlib.suppress(Exception): + layer.destroy()Add the import near the other stdlib imports:
import argparse +import contextlib import gc📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.if bench_backend is not None and bench_workspace is not None: with contextlib.suppress(Exception): bench_backend.destroy(bench_workspace) if layer is not None: with contextlib.suppress(Exception): layer.destroy()🧰 Tools
🪛 GitHub Actions: pre-commit / 0_pre-commit.txt
[error] 429-429: ruff check (SIM105): Use
contextlib.suppress(Exception)instead oftry-except-pass
[error] 434-434: ruff check (SIM105): Use
contextlib.suppress(Exception)instead oftry-except-pass🪛 GitHub Actions: pre-commit / pre-commit
[error] 429-429: ruff check (SIM105): Use
contextlib.suppress(Exception)instead oftry-except-pass
[error] 434-434: ruff check (SIM105): Use
contextlib.suppress(Exception)instead oftry-except-pass🪛 Ruff (0.15.21)
[error] 431-432:
try-except-passdetected, consider logging the exception(S110)
[error] 436-437:
try-except-passdetected, consider logging the exception(S110)
🤖 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 `@benchmarks/bench_moe_ep_sm90_mega.py` around lines 428 - 437, Replace both cleanup try/except/pass blocks in the benchmark teardown with contextlib.suppress, covering bench_backend.destroy(bench_workspace) and layer.destroy(). Add the contextlib import alongside the existing standard-library imports.Source: Pipeline failures
flashinfer/moe_ep/backends/mega/kernel/sm90_pull_fp8/weights.py (2)
23-23: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Remove unused
Optionalimport (ruff F401, failing pre-commit).
Optionalis imported but never referenced in this module.Proposed fix
-from typing import TYPE_CHECKING, Literal, Optional, Tuple +from typing import TYPE_CHECKING, Literal, Tuple📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.from typing import TYPE_CHECKING, Literal, Tuple🧰 Tools
🪛 GitHub Actions: pre-commit / 0_pre-commit.txt
[error] 23-23: ruff check (F401):
typing.Optionalimported but unused🪛 GitHub Actions: pre-commit / pre-commit
[error] 23-23: ruff check (F401):
typing.Optionalimported but unused🤖 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/backends/mega/kernel/sm90_pull_fp8/weights.py` at line 23, Remove the unused Optional symbol from the typing import in weights.py, while preserving the remaining TYPE_CHECKING, Literal, and Tuple imports.Source: Pipeline failures
319-334: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Annotate the shape tuples to fix the failing mypy check (L326/L330).
fc1_sf_shape/fc2_sf_shapeare inferred astuple[int, int, int]from the blockwise branch, so the 2-tuple assignments in theelsebranch fail withIncompatible types in assignment. Declare both as variable-length shape tuples before the branch.Proposed fix
+ fc1_sf_shape: Tuple[int, ...] + fc2_sf_shape: Tuple[int, ...] if blockwise: fc1_sf_shape = (local_experts, fc1_out // 128, hidden_size // 128) fc2_sf_shape = (local_experts, hidden_size // 128, intermediate_size // 128) sf_dtype = torch.float32 else: from .....kernel_src.sm90.pull_style_cutedsl_megakernel import ceil_div fc1_sf_shape = ( local_experts, _swizzled_flat_e8m0_size(fc1_out, ceil_div(hidden_size, 32)), ) fc2_sf_shape = ( local_experts, _swizzled_flat_e8m0_size(hidden_size, ceil_div(intermediate_size, 32)), ) sf_dtype = torch.uint8📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.fc1_sf_shape: Tuple[int, ...] fc2_sf_shape: Tuple[int, ...] if blockwise: fc1_sf_shape = (local_experts, fc1_out // 128, hidden_size // 128) fc2_sf_shape = (local_experts, hidden_size // 128, intermediate_size // 128) sf_dtype = torch.float32 else: from .....kernel_src.sm90.pull_style_cutedsl_megakernel import ceil_div fc1_sf_shape = ( local_experts, _swizzled_flat_e8m0_size(fc1_out, ceil_div(hidden_size, 32)), ) fc2_sf_shape = ( local_experts, _swizzled_flat_e8m0_size(hidden_size, ceil_div(intermediate_size, 32)), ) sf_dtype = torch.uint8🧰 Tools
🪛 GitHub Actions: pre-commit / 0_pre-commit.txt
[error] 326-326: mypy: Incompatible types in assignment (expression has type "tuple[int, int]", variable has type "tuple[int, int, int]") [assignment]
[error] 330-330: mypy: Incompatible types in assignment (expression has type "tuple[int, int]", variable has type "tuple[int, int, int]") [assignment]
🪛 GitHub Actions: pre-commit / pre-commit
[error] 326-326: mypy: Incompatible types in assignment (expression has type "tuple[int, int]", variable has type "tuple[int, int, int]") [assignment]
[error] 330-330: mypy: Incompatible types in assignment (expression has type "tuple[int, int]", variable has type "tuple[int, int, int]") [assignment]
🤖 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/backends/mega/kernel/sm90_pull_fp8/weights.py` around lines 319 - 334, Annotate fc1_sf_shape and fc2_sf_shape before the blockwise conditional as variable-length integer shape tuples, such as tuple[int, ...], so both the 3-dimensional blockwise assignments and 2-dimensional non-blockwise assignments type-check without changing their values.Source: Pipeline failures
flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/common/megamoe_constants.py (1)
19-25: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Duplicate reciprocal-limit definitions.
Lines 23-25 re-declare
Nvfp4E2M1RcpLimit,Fp8E4M3RcpLimit, andFp8E5M2RcpLimitwith the exact same RHS as lines 19-21. The second block is dead and should be removed (or was one of them meant to differ?).🧹 Proposed cleanup
Nvfp4E2M1RcpLimit = 1.0 / Nvfp4E2M1Max Fp8E4M3RcpLimit = 1.0 / Fp8E4M3FNMax Fp8E5M2RcpLimit = 1.0 / Fp8E5M2Max - -Nvfp4E2M1RcpLimit = 1.0 / Nvfp4E2M1Max -Fp8E4M3RcpLimit = 1.0 / Fp8E4M3FNMax -Fp8E5M2RcpLimit = 1.0 / Fp8E5M2Max📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.Nvfp4E2M1RcpLimit = 1.0 / Nvfp4E2M1Max Fp8E4M3RcpLimit = 1.0 / Fp8E4M3FNMax Fp8E5M2RcpLimit = 1.0 / Fp8E5M2Max🤖 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/sm100/cutedsl_megamoe/src/common/megamoe_constants.py` around lines 19 - 25, Remove the duplicate second definitions of Nvfp4E2M1RcpLimit, Fp8E4M3RcpLimit, and Fp8E5M2RcpLimit in megamoe_constants.py, preserving the first reciprocal-limit block unchanged.flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/mega_reference.py (1)
97-97: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Return annotation understates arity (4 values returned, annotated as 3-tuple).
reference_expert_fc12returnsfc2_fp32, fc1_q, fc1_sf_out, fc1_fp32(Line 139) and the caller unpacks four (Line 336), but the annotation declares a 3-tuple.🩹 Proposed fix
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:🤖 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/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/mega_reference.py` at line 97, Update the return annotation of reference_expert_fc12 to declare a four-element tuple matching its four returned tensors: fc2_fp32, fc1_q, fc1_sf_out, and fc1_fp32. Keep the return values and caller unpacking unchanged.flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/mega_runner.py (1)
2758-2761: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Unprofessional / broken user-facing error message.
The message contains profanity and an incomplete clause ("unsupported by the current."). Please make it a clear, professional diagnostic.
✏️ Suggested wording
- raise ValueError( - f"num_topk ({problem.num_topk}) > 32 is unsupported by the current. Shit inherited from DeepGEMM." - ) + raise ValueError( + f"num_topk ({problem.num_topk}) > 32 is unsupported " + f"(inherited 32-topk limit from the DeepGEMM path)." + )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.if problem.num_topk > 32: raise ValueError( f"num_topk ({problem.num_topk}) > 32 is unsupported " f"(inherited 32-topk limit from the DeepGEMM path)." )🤖 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/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/mega_runner.py` around lines 2758 - 2761, Replace the ValueError message in the num_topk validation with a clear, professional diagnostic stating that values above 32 are unsupported by the current implementation; remove the profanity and incomplete wording while preserving the existing validation behavior.flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_common.py (1)
567-582: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash set -euo pipefail echo "== locate file ==" git ls-files | rg '^flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_common\.py$' || true echo "== relevant lines ==" if [ -f flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_common.py ]; then sed -n '460,640p' flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_common.py | nl -ba -v460 fi echo "== call sites for _rcp_approx_kernel and _swiglu_pair_kernel ==" rg -n "_rcp_approx_kernel|_swiglu_pair_kernel|runner_common::|from .*runner_common|import runner_common|get|rcp|swiglu" flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamOE flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe || true echo "== repository search for the file ==" rg -n "triton\.runtime|compile_if_definition|kernel|n_elements|BLOCK.*=|program_id|arange|int64|int32" flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab -g '*.py' || trueRepository: flashinfer-ai/flashinfer
Length of output: 330
🏁 Script executed:
#!/bin/bash set -euo pipefail echo "== relevant lines ==" sed -n '460,640p' flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_common.py echo "== call sites for _rcp_approx_kernel and _swiglu_pair_kernel ==" rg -n "_rcp_approx_kernel|_swiglu_pair_kernel|runner_common::|from .*runner_common|import runner_common|get|rcp|swiglu" flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe || true echo "== repository-wide references and tensor size terms ==" rg -n "runner_common|_rcp_approx_kernel|_swiglu_pair_kernel|world_size|n_elements|numel|num_tokens|intermediate|topk|hidden_size|rank|swap" flashinfer -g '*.py' | head -n 250 || trueRepository: flashinfer-ai/flashinfer
Length of output: 50381
🏁 Script executed:
#!/bin/bash set -euo pipefail echo "== runner_common size and remaining relevant lines ==" wc -l flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_common.py sed -n '180,526p' flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_common.py sed -n '640,670p' flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_common.py echo "== exact imports/usages of rcp/swiglu Triton helpers ==" rg -n "_rcp_approx_ftz_f32_cuda|_swiglu_pair_hw_match_cuda|_get_rcp_approx_triton_kernel|_get_swiglu_pair_hw_match_triton_kernel|swiglu_fold_interleave\(" flashinfer -g '*.py' || true echo "== deterministic type/arithmetic probe from source ==" python3 - <<'PY' import ast from pathlib import Path src = Path("flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_common.py").read_text() tree = ast.parse(src) for node in tree.body: if isinstance(node, ast.FunctionDef) and node.name in {"_get_rcp_approx_triton_kernel", "_get_swiglu_pair_hw_match_triton_kernel"}: funcs = [n for n in node.body if isinstance(n, ast.FunctionDef)] print(f"{node.name}:") for fn in funcs: assign = [n for n in fn.body if isinstance(n, ast.Assign) and len(n.targets)==1 and isinstance(n.targets[0], ast.Name) and n.targets[0].id=="offsets"] if assign: print(f" {fn.name}: assignment = {ast.unparse(assign[0].value)}") PY echo "== config / problem size terms for nvfp4/swiglu paths ==" rg -n "token_max|max_tokens|num_tokens|hidden_size|intermediate|topk|num_experts|ep_size|world_size|gate_up_interleave|swiglu_fold" flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe -g '*.py' | head -n 200 || trueRepository: flashinfer-ai/flashinfer
Length of output: 43841
Document the 32-bit offset assumption in
_rcp_approx_kernel/_swiglu_pair_kernel.These kernels compute
offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)without.to(tl.int64), unlike_pack_fp4_kernelwhich explicitly handles the billions-of-elements combine path. Add a small guard/comment here for the supportedn_elementsrange so callers cannot pass~2^31element tensors to these Triton helpers.🤖 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/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_common.py` around lines 567 - 582, Document and enforce the supported 32-bit offset range in _rcp_approx_kernel and _swiglu_pair_kernel. Add a small guard or assertion that rejects n_elements at or beyond the signed 32-bit limit, and clearly comment that these kernels intentionally use 32-bit offsets and do not support ~2^31-element tensors.flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_fc12_common.py (1)
111-130: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash set -euo pipefail echo "== locate file ==" fd -a 'runner_fc12_common.py' . || true echo "== relevant lines ==" sed -n '1,180p' flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_fc12_common.py echo "== searches for sf_vec_size/Block sizes and hidden checks ==" rg -n "Nvfp4BlockSize|Mxfp8BlockSize|sf_vec_size|hidden|ceil_div|ProblemDesc|intermediate" flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab -S echo "== files containing hidden block-size validation or SF coverage ==" rg -n "hidden.*Block|Block.*hidden|ceil_div\\(hidden|sf_vec_size|activation.*SF|SF block" flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab -SRepository: flashinfer-ai/flashinfer
Length of output: 50381
🏁 Script executed:
#!/bin/bash set -euo pipefail echo "== constants ==" fd -a 'megamoe_constants.py' . | xargs -r -I{} sh -c 'echo "-- {}"; sed -n "1,140p" "{}"' echo "== hidden validation all occurrences ==" rg -n "hidden.*must be a positive multiple|ProblemDesc|kind_sf_vec_size|Nvfp4BlockSize|Mxfp8BlockSize" flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src -S echo "== kind_host_utils excerpt ==" fd -a 'host_utils.py' flashinfer/moe_ep/kernel_src | xargs -r -I{} sh -c 'echo "-- {}"; sed -n "1,120p" "{}"' echo "== behavioral simulation of guard for representative values ==" python3 - <<'PY' from pathlib import Path import re paths = list(Path("flashinfer/moe_ep/kernel_src").rglob("megamoe_constants.py")) + \ list(Path("flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src").rglob("common.py")) for p in paths: text = p.read_text() for m in re.finditer(r"(?:Nvfp4BlockSize|Mxfp8BlockSize)\s*=\s*(\d+)", text): print(f"{p}:{text[:m.start()].count(chr(10))+1}: {m.group(0)}") nv=16 mx=32 for kind in ["nvfp4", "mxfp8_e4m3", "mxfp8_e5m2"]: guard = nv print(f"{kind}: current hidden guard requires multiple of {guard}; mxfp8 sf_vec=16? no 32") for hidden in [16, 32, 40, 64]: print(f"hidden={hidden}: current hidden guard passes={hidden % nv == 0}") PYRepository: flashinfer-ai/flashinfer
Length of output: 38510
Make
hiddenvalidation match the activation SF block size.
ProblemDesc.__post_init__currently enforceshidden % Nvfp4BlockSizefor everykind, sohidden=16passes formxfp8_e4m3/mxfp8_e5m2even though the MXFP8 activation and quantization paths usesf_vec_size = Mxfp8BlockSize = 32. Use the sameNvfp4BlockSizevsMxfp8BlockSizebranch used by the intermediate checks before blocking unsupportedhiddenvalues.🤖 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/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_fc12_common.py` around lines 111 - 130, The hidden-size validation in ProblemDesc.__post_init__ uses Nvfp4BlockSize for every kind. Select the activation block size based on self.kind, using Nvfp4BlockSize for nvfp4 and Mxfp8BlockSize for MXFP8 kinds, then validate hidden against that selected size while preserving the existing positive-multiple error behavior.flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/iket_compat.py (1)
19-19: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Replace the bare
warnings.warn(or drop it).Per the surrounding comment, public-release / CTK-12.9 wheels routinely land in this fallback, so this line prints to stdout on every import in the common case — noisy inside serving engines and test harnesses. Route it through
warnings/logging(or remove it), matching the no-op-shim intent.🔧 Proposed change
try: from cutlass.cute import iket # type: ignore except (ImportError, NotImplementedError): - print("!!!! Iket is not enabled !!!!") + import warnings + + warnings.warn( + "cute.iket dialect unavailable; using no-op IKET markers.", + RuntimeWarning, + stacklevel=2, + )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.try: from cutlass.cute import iket # type: ignore except (ImportError, NotImplementedError): import warnings warnings.warn( "cute.iket dialect unavailable; using no-op IKET markers.", RuntimeWarning, stacklevel=2, )🤖 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/sm100/cutedsl_megamoe/src/src/iket_compat.py` at line 19, Replace the import-time print in the Iket compatibility fallback with warnings.warn or remove it entirely, preserving the no-op shim behavior and avoiding stdout output on common imports. Update only the fallback notification around the “Iket is not enabled” message.flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/reference.py (1)
15-18: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Stale metadata-packing docstring. The module header describes a 12-byte
[rank_idx:u32][token_idx:u32][topk_idx:u32]layout, but_pack_metadata(andTokenSrcMetadataintoken_comm.py,nbytes=8) actually uses an 8-byte i64: low 32b =src_token, high 32b =(src_rank << 16) | src_topk. Update the header to the 8-byte convention so the oracle's ground-truth doc matches the implementation.🤖 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/sm100/cutedsl_megamoe/src/src/reference.py` around lines 15 - 18, Update the module header metadata-packing documentation to describe the 8-byte i64 layout implemented by _pack_metadata and TokenSrcMetadata: low 32 bits are src_token, while high 32 bits contain (src_rank << 16) | src_topk. Replace the stale 12-byte field description and ensure the padding sentinel behavior remains documented.flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/shim/hopper_fp8.py (2)
924-933: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
mypy failure: reuse of
shapeloop variable conflicts with the 3-tuple binding above.
shapeis first bound toTuple[int, int, int]in theweight_checksloop (Line 845); reusing it here for the 1-tuple/(e,)scale shapes triggers thepre-commitmypy errorIncompatible types in assignment ... tuple[int] vs tuple[int, int, int]. Rename the scale-loop target.Proposed fix
- for name, tensor, shape in scale_checks: + for name, tensor, expected_shape in scale_checks: _require_cuda(name, tensor) - if tuple(tensor.shape) != shape: + if tuple(tensor.shape) != expected_shape: raise ValueError( - f"{name} must have shape {shape}, got {tuple(tensor.shape)}." + f"{name} must have shape {expected_shape}, got {tuple(tensor.shape)}." )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.for name, tensor, expected_shape in scale_checks: _require_cuda(name, tensor) if tuple(tensor.shape) != expected_shape: raise ValueError( f"{name} must have shape {expected_shape}, got {tuple(tensor.shape)}." ) if tensor.dtype != torch.float32: raise ValueError( f"{name} must be float32, got {tensor.dtype}." )🧰 Tools
🪛 GitHub Actions: pre-commit / 0_pre-commit.txt
[error] 924-924: mypy: Incompatible types in assignment (expression has type "tuple[int]", variable has type "tuple[int, int, int]") [assignment]
🪛 GitHub Actions: pre-commit / pre-commit
[error] 924-924: mypy: Incompatible types in assignment (expression has type "tuple[int]", variable has type "tuple[int, int, int]") [assignment]
🤖 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/pull_style_cutedsl_megakernel/shim/hopper_fp8.py` around lines 924 - 933, Rename the `shape` loop variable in the `scale_checks` iteration to a distinct name, and update the associated shape comparison and error message in that block. Leave the earlier `weight_checks` loop binding unchanged.Source: Pipeline failures
1720-1720: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
mypy failure: narrow
FP8_SCALE_MODEto theLiteralbefore passing it tocreate_dummy_inputs.
os.environ.get(...)returnsstr, which failscreate_dummy_inputs'sfp8_scale_mode: Literal["per_tensor", "blockwise"]param (pre-commit mypy error at Line 1736).Proposed fix
- FP8_SCALE_MODE = os.environ.get("MEGA_FP8_SCALE_MODE", "per_tensor") + FP8_SCALE_MODE: Literal["per_tensor", "blockwise"] = ( + "blockwise" + if os.environ.get("MEGA_FP8_SCALE_MODE", "per_tensor") == "blockwise" + else "per_tensor" + )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.FP8_SCALE_MODE: Literal["per_tensor", "blockwise"] = ( "blockwise" if os.environ.get("MEGA_FP8_SCALE_MODE", "per_tensor") == "blockwise" else "per_tensor" )🤖 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/pull_style_cutedsl_megakernel/shim/hopper_fp8.py` at line 1720, Update the FP8_SCALE_MODE initialization and the call to create_dummy_inputs so the environment value is validated or narrowed to the accepted Literal["per_tensor", "blockwise"] values before being passed, while preserving "per_tensor" as the default and rejecting or handling unsupported values explicitly.Source: Pipeline failures
flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/epilogue_refactor.py (1)
464-475: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash # Locate warp_redux_sync definition/semantics and other half-warp usages to compare the intended pattern. fd -t f 'epilogue_refactor.py' flashinfer/moe_ep/kernel_src | xargs -I{} sed -n '460,500p' {} rg -nP '\bwarp_redux_sync\s*\(' flashinfer/moe_ep/kernel_src -C2Repository: flashinfer-ai/flashinfer
Length of output: 11764
🏁 Script executed:
#!/bin/bash set -euo pipefail echo "== CUTE repo locations ==" git ls-files | rg 'cute|flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/flashinfer_cutlass_frontend|README|requirements|pyproject|setup' echo echo "== Search for warp_redux_sync implementation ==" rg -n "def warp_redux_sync|warp_redux_sync|redux\.sync|membermask" . -g '*.py' -g '*.cpp' -g '*.cuh' -g '*.cu' | head -n 200 echo echo "== Inspect files likely containing cute backend ==" for f in $(git ls-files | rg '(^|/)(cute|flashinfer_cutlass_frontend|nvrtc|compiler).*\.py$|cute/__init__.py|cute/config.*\.py$'); do echo "--- $f" sed -n '1,220p' "$f" doneRepository: flashinfer-ai/flashinfer
Length of output: 50380
🌐 Web query:
Cute architecture warp_redux_sync redux.sync membermask divergent if CUDA💡 Result:
In NVIDIA CUTLASS, warp_redux_sync is a function within the CuTe DSL (cutlass.cute.arch) used to perform warp-level reduction operations across threads [1]. Key technical aspects include: Operation and Usage The function reduces values from participating threads in a warp using a specified operation (e.g., 'add', 'max', 'min', 'fmax', 'fmin', 'xor', etc.) [1]. The result is broadcast to all threads in the specified mask [1]. As of recent updates (specifically with CUDA Toolkit 13.1), the kind argument for the reduction operation must be provided as a string literal rather than an enum [2][3][4]. Membermask and Divergence The mask_and_clamp argument represents the warp participation mask (defaulting to 0xFFFFFFFF, which includes all threads in a warp) [1]. Regarding divergence: 1. Thread participation: Warp-level synchronization primitives generally require all active threads in the specified mask to participate [5][6]. 2. Divergence Behavior: In CUDA, using synchronization primitives within divergent control flow can lead to undefined behavior or deadlocks if not handled correctly, as threads in different execution paths may not reach the synchronization point simultaneously [7]. 3. PTX Context: At the PTX instruction level, redux.sync is a communication instruction that functions across the specified mask [5][6]. The hardware and PTX compiler expect the membermask to accurately reflect the threads arriving at the synchronization point [8]. Failure to match the mask to the active, arriving threads often leads to illegal instructions or execution hangs [9]. In summary, warp_redux_sync requires careful management of the thread mask to ensure all participating threads are converged at the point of the reduction to avoid divergence-related errors [1][5][9].
Citations:
- 1: https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/cute_dsl_api/cute_arch.html
- 2: https://docs.nvidia.com/cutlass/4.4.2/overview.html
- 3: https://docs.nvidia.com/cutlass/latest/CHANGELOG.html
- 4: https://github.com/NVIDIA/cutlass/blob/main/CHANGELOG.md
- 5: https://docs.nvidia.com/cuda/pdf/ptx_isa_8.7.pdf
- 6: https://docs.nvidia.com/cuda/archive/12.4.0/pdf/ptx_isa_8.4.pdf
- 7: https://docs.nvidia.com/cuda/parallel-thread-execution/index.html
- 8: https://gh.evko.io/crucible-notes/ptxas/sass-isa/control-flow.html
- 9: https://groups.google.com/g/llvm-dev/c/ZYmVxDbHbnI
Fix the
sf_vec_size < 32half-warp amax reduction.For the 16-lane SF blocks, the second
cute.arch.warp_redux_sync(v, "fmax", abs=True)must only reduce lanes 16..31; reducing the whole warp makes both per-16 blocks receive the same full-warp amax. In addition, keep the full-warp reductions converged: issuing them inside theif not first_halfbranch from only second-half lanes is divergent for the active mask and defeats the expectedredux.syncsemantics.🤖 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/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/epilogue_refactor.py` around lines 464 - 475, Update _amax_lane for sf_vec_size < 32 so each 16-lane half-warp computes its own amax: use a converged warp reduction with lane-masked values for the second half rather than calling warp_redux_sync only inside if not first_half. Preserve the existing full-warp reduction when sf_vec_size == 32 and ensure both half-warp reductions execute uniformly across participating lanes.flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/mega_reference.py (1)
97-97: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Return annotation is a 3-tuple but the function returns 4 tensors.
reference_expert_fc12returnsfc2_fp32, fc1_q, fc1_sf_out, fc1_fp32(Line 131), and both call sites unpack four values, so only the annotation is out of sync.🩹 Proposed annotation fix
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:🤖 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/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/mega_reference.py` at line 97, Update the return annotation of reference_expert_fc12 to describe a 4-tensor tuple, matching its four returned values and both call-site unpacking patterns.tests/moe_ep/test_moe_ep_sm90_pull_fp8_mega_multirank.py (1)
578-579: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
These two tests never run in CI (marker filter deselects them).
This file is excluded from
run_unitand is only executed by themega_sm90target with-m "gpu_4 and arch_hopper". That filter requires both markers, so:
test_sm90_pull_fp8_preprocess_mega_weights_from_bf16(Line 578, marked only@pytest.mark.arch_hopper) is deselected — losing its unique preprocess shape/stride/dtype coverage.test_sm90_pull_fp8_mega_kernel_is_registered(Line 615, no marker) is deselected.Add the
gpu_4marker (the multirank harness assignsnum_local_experts = num_experts // world_size, which holds for 4 ranks) so the preprocess test is actually collected; the registration check is already covered bytest_sm90_pull_fp8_config.py::test_registry_resolves_backend, so either mark it or drop it here.🧪 Suggested marker addition for the preprocess test
+@pytest.mark.gpu_4 `@pytest.mark.arch_hopper` def test_sm90_pull_fp8_preprocess_mega_weights_from_bf16():Also applies to: 615-616
🤖 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/test_moe_ep_sm90_pull_fp8_mega_multirank.py` around lines 578 - 579, Add the gpu_4 pytest marker to test_sm90_pull_fp8_preprocess_mega_weights_from_bf16 so it is selected by the mega_sm90 multirank filter alongside arch_hopper. For test_sm90_pull_fp8_mega_kernel_is_registered, either add gpu_4 and arch_hopper markers or remove the redundant test, since registry coverage already exists in test_registry_resolves_backend.
9 tasks
… (4.6.1 is a perf floor) Same sweep recipe/geometry as the headline table, DSL runtime pinned to 4.5.2: compiles and runs, but generated code is 34-54% slower than 4.6.1 across every nvfp4 variant and token count (dg baseline reproduces within 1%, isolating the delta to the DSL runtime). Treat 4.6.1 as a performance floor, not just a compile-compatibility floor. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… notes) Four actionable follow-ups from the 2026-07-15 vLLM 0.25.1 e2e integration (DeepSeek-V4-Flash, 4x GB200): fi_dg run-to-run nondeterminism, CUDA-graph capture (analysis updated with e2e context), MoEEpMegaLayer source-weight retention OOM at model load, and hot-path-free knob selection (offline tuning + geometry/token-bucket heuristic; in-engine knobs=auto is unusable and its tuner-harness winner lost 13% e2e). Plus four earlier analysis notes: MoEWeightPack discriminated-union refactor, multinode support, nixl_ep suppression, trtllm import. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…l load) MoEEpMegaLayer kept self._weights alive for the layer lifetime even though nothing reads it after preprocess_weights() produces the transformed tensors. In the fp4-checkpoint flow the pack is a ~3.2 GB per-layer bf16 dequant copy, so 43 MoE layers pinned ~140 GB of dead weight and OOMed a 186 GB GB200 at model load (vLLM e2e run 7, 2026-07-15). - MoEEpMegaLayer: type _weights Optional, skip storing the pack when backend.transformed_weights is supplied, release it at the end of _preprocess_weights(); memory invariant documented in the class docstring. - MoEEpSplitLayer: same pattern — the pack is only needed for init-time validation and kernel preprocessing (the kernel retains what it needs), so release it at the end of __init__. - Weakref-based release tests for all three paths; factory docstring and todo_weight_pack_retention_oom.md updated (downstream vLLM patch workaround is now a redundant no-op). Resolves flashinfer/moe_ep/todo_weight_pack_retention_oom.md. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ntract + capture guards) Implements the todo_cuda_graph.md restoration plan (single-rank scope). The per-forward sync was already gone (sync=False entry-point defaults); this adds the pieces that make capture safe and testable: - shim/comm.py: ensure_not_capturing() guard; wired after the no-op early returns in _ensure_mega_compiled, set_gate_up_clamp, apply_knobs, _release_workspace (both nvfp4 and mxfp8 frontends) and at the top of the autotune_knobs collective sweep — host-side compile/alloc/free now fails loudly instead of corrupting a capture. - shim run()/mega_moe entry points: sync=True is gated on is_current_stream_capturing (zero-break variant from the TODO). - MoEEpMegaLayer.warmup(): one full eager forward (workspace alloc, cute.compile, knobs="auto" sweep, real launch) + device sync; documents the call-on-all-ranks-before-capture contract. _ensure_workspace raises under capture with a warmup() hint. - tests: tests/moe_ep/test_mega_cuda_graph.py (GB200, MEGA_NO_DIST=1) — nvfp4 + mxfp8 capture post-warmup, 3x replay bit-exact vs eager, replay over in-place-mutated inputs, and capture-without-warmup raises; wired into run_tests.sh oracle section (excluded from unit). Plus mocked capture-guard unit tests in test_mega_layer_validation.py. Verified on GB200: unit 164 passed, graph tests 3/3 passed. Remaining (tracked in todo_cuda_graph.md): 2-rank lockstep replay + skewed staging stress tests, vLLM shared-workspace warmup adoption. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The kernel package ships DataPreprocess (src/src/inputs_process.py) — a fused CuTe-DSL kernel that quantizes bf16 activations (NVFP4 per-16 E4M3 scales / MXFP8 per-32 E8M0) and repacks routing to int64/fp32 in ONE launch — but nothing wired it: both cutedsl mega backends staged through a torch-composed path (fp32 upcast + ~a dozen elementwise/reduce launches + five buffer copies) on every forward of every layer. The vLLM e2e nsys showed these paths host-gap-bound (946k cudaLaunchKernel on fi_nvfp4 vs 100k native), so per-forward launch count is the first-order cost. - shim/quant_stage.py: fused_quant_stage() wrapping DataPreprocess with the frontends' caching pattern — one cute.compile per (topk, hidden, quant_type) with dynamic token extent, launch-args cache keyed on pointers + token count + stream, capture-guarded compile, and the capacity-tail topk_idx=-1 re-mask for torch-path parity. Offline norm_const mode only (online amax mode exists in the kernel, not wired). - nvfp4/mxfp8 backend staging: fused path is the default; FLASHINFER_MEGA_FUSED_STAGE=0 restores the torch path (bisection aid, documented in CLAUDE.md). quantize_input=False (pre-staged) unchanged. - tests/moe_ep/test_fused_quant_stage.py: fused vs torch staging must agree BIT-EXACTLY (the repo's torch quantizers are bit-matched to the kernel) across all three quant kinds, partial/full capacity, and launch-cache hit/rebuild; verified on GB200 (7 passed). CUDA-graph capture/replay (3 passed, staging kernel captured), torch-oracle section (PASS), and unit suite (171 passed) all green with the fused path active. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…hot path)
Implements the lookup half of todo_vllm_knob_heuristic.md: tuning becomes an
offline step whose winners persist, and engine-time knob resolution is a
dict lookup — no compiles, no collectives, no timing in the serving path.
- shim/knob_cache.py: JSON cache (FLASHINFER_MOE_EP_KNOB_CACHE, default
~/.cache/flashinfer/moe_ep_knob_cache.json) keyed on (device, dtype,
world_size, hidden, intermediate, num_experts, topk, combine wire) with
nearest-bucket max_tokens selection; atomic upsert writes; corrupted or
disabled cache degrades to the heuristic with a warning.
- get_symm_buffer_for_{mega,mxfp8_mega}_moe: knobs=None resolves cache-first,
then tuner.default_knobs; explicit dicts unchanged.
- autotune_knobs: on_winner callback; the nvfp4/mxfp8 wrappers record each
winner (rank 0) so any knobs="auto" run persists its result.
- knobs="auto" warns loudly at backend construction pointing at the offline
flow (it remains a collective multi-minute sweep — never in-engine).
- python -m flashinfer.moe_ep.tune: offline CLI (torchrun multi-rank or
MEGA_NO_DIST=1); deterministic candidates by default,
--allow-nondeterministic for ikr; --max-candidates for smoke runs.
- tests/moe_ep/test_knob_cache.py: round-trip/bucket/miss/upsert/disable/
corruption CPU tests + GPU test that buffer creation picks up a cached
entry. Verified on GB200: unit 181 passed, GPU test passed, CLI smoke
recorded a winner and a fresh process resolved it via lookup.
Remaining (tracked in todo_vllm_knob_heuristic.md): populate the cache for
production geometries, skewed-rank offline timing mode, vLLM matrix re-run.
AI-assisted.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
AI-assisted follow-up to the CodeRabbit review, applied to both the sm100 and sm90 kernel trees where the code is mirrored: - epilogue_refactor: _amax_lane issued full-mask redux.sync inside divergent half-warp branches (UB per PTX). Now branchless: each lane issues one redux with its own half's membermask, so each 16-lane half gets its own amax. - runner_fc12_common: hidden-size validation branches on quantization kind (Nvfp4BlockSize vs Mxfp8BlockSize), matching the existing _sf_vec convention. - runner_common: guard the three flat-index Triton helpers against >=2**31 elements (int32 offset arithmetic would wrap). - megamoe_constants: drop duplicated RcpLimit definitions. - mega_reference / kernel_mxfp8_glu_fc12: fix 3-tuple return annotations on functions returning 4 values; sync docstrings. - kernel_mxfp8_glu_fc12: remove leftover [fc12 stages] print and the dead GMEM-probing debug block in the fc2 spin-exit path. - mega_runner / sym_buffer: reword unprofessional error message and TODO comments. - iket_compat: import-time print -> warnings.warn. - reference.py: metadata docstring updated to the actual 8-byte i64 layout (low 32b src_token, high 32b (src_rank << 16) | src_topk). - ptx_helpers: _fence_rel_gpu docstring said system scope; it is device (gpu) scope. - contract.py: Contract.table cached via functools.cached_property (validation stays eager in __post_init__). - core/kernel/base.py: sm100-specific quant-stage memo eviction moved out of the shared destroy() into a _forget_workspace_state() hook, overridden by the two sm100 cutedsl backends. Validated on H100 (SLURM): run_tests.sh unit, oracle_sm90, mega_sm90 all pass; pre-commit clean. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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
`@flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/contract.py`:
- Around line 313-318: Replace the standalone self.table eager-access expression
with an assignment to _ in both
flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/contract.py
lines 313-318 and
flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/contract.py
lines 312-318, preserving eager validation and cached-property population.
🪄 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 Plus
Run ID: bfbe1e2a-1e36-4c5c-a2ed-802b70831393
📒 Files selected for processing (25)
flashinfer/moe_ep/backends/mega/kernel/mxfp8_cutedsl/backend.pyflashinfer/moe_ep/backends/mega/kernel/nvfp4_cutedsl/backend.pyflashinfer/moe_ep/core/kernel/base.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/common/megamoe_constants.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_mxfp8_glu/kernel_mxfp8_glu_fc12.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/contract.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/epilogue_refactor.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/mega_reference.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/mega_runner.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_common.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_fc12_common.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/iket_compat.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/ptx_helpers.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/reference.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/sym_buffer.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/common/megamoe_constants.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/contract.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/epilogue_refactor.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/mega_reference.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/mega_runner.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/runner_fc12_common.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/iket_compat.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/ptx_helpers.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/reference.pyflashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/sym_buffer.py
💤 Files with no reviewable changes (2)
- flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/common/megamoe_constants.py
- flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/common/megamoe_constants.py
🚧 Files skipped from review as they are similar to previous changes (17)
- flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/epilogue_refactor.py
- flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/iket_compat.py
- flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/reference.py
- flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/mega_runner.py
- flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_common.py
- flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_mxfp8_glu/kernel_mxfp8_glu_fc12.py
- flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/mega_reference.py
- flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/ptx_helpers.py
- flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_fc12_common.py
- flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/ptx_helpers.py
- flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/sym_buffer.py
- flashinfer/moe_ep/backends/mega/kernel/nvfp4_cutedsl/backend.py
- flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/mega_reference.py
- flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/src/sym_buffer.py
- flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/runner_fc12_common.py
- flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/mega_runner.py
- flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/epilogue_refactor.py
Collaborator
|
[FAILED] Pipeline #60371539 — 12/18 executed test jobs passed Compared with nightly #60276939. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 6/6 passed
Failure detailsTimeouts, infrastructure, or incomplete jobs
|
Collaborator
Author
|
/bot run moe_ep |
Collaborator
|
@mhoqueanik is not authorized to trigger this CI job. cc: @yzh119, @sricketts, @yongwww |
Collaborator
|
/bot run moe_ep |
Collaborator
Collaborator
|
[FAILED] Pipeline #60951733 — 0/0 executed test jobs passed No usable JUnit artifact was available; individual tests and nightly comparison could not be recovered. Failure detailsTimeouts, infrastructure, or incomplete jobs
|
Collaborator
|
/bot run moe_ep |
Collaborator
Collaborator
|
[FAILED] Pipeline #60952111 — 0/0 executed test jobs passed No usable JUnit artifact was available; individual tests and nightly comparison could not be recovered. Failure detailsTimeouts, infrastructure, or incomplete jobs
|
5 tasks
Anerudhan
enabled auto-merge (squash)
August 5, 2026 18:45
Anerudhan
approved these changes
Aug 5, 2026
bkryu
approved these changes
Aug 5, 2026
This was referenced Aug 7, 2026
5 tasks
aleozlx
pushed a commit
that referenced
this pull request
Aug 14, 2026
…ush-style FP8 backend; sync CuTe-DSL 4.7 quant-staging fix (#4449) ## Summary Three things: a layout/naming refactor of `flashinfer.moe_ep`'s mega-kernel layer, the incorporation of the SM90 push-style FP8 backend (#4069, since merged upstream) as the first new backend added in the restructured shape, and one vendored-kernel sync that fixes the fused activation-quant staging crash on CuTe-DSL 4.7 — un-blocking 4.7.x and lifting the temporary `==4.6.1` pin. The branch is merged up to upstream/main (2febce5, past the v0.6.17 line and the #4069 squash). The refactor organizes the layer around two orthogonal views: 1. **Taxonomy (user view)** — backends move to `backends/mega/kernel/sm<arch>/<act_dtype>_<weight_dtype>_<out_dtype>_<kernel_style>/`, and registry `kernel_name` strings plus config classes carry the same fully-qualified names. One glance at a name now tells you the architecture, the activation/weight/output dtypes, and the kernel style: | old kernel_name | new kernel_name | new config class | |---|---|---| | `deep_gemm_mega` | `sm100_fp8_fp4_bf16_deepgemm` | `Sm100_Fp8_Fp4_Bf16_Deepgemm_MegaMoeConfig` | | `nvfp4_cutedsl` | `sm100_nvfp4_nvfp4_bf16_cutedsl` | `Sm100_Nvfp4_Nvfp4_Bf16_Cutedsl_MegaMoeConfig` | | `mxfp8_cutedsl` | `sm100_mxfp8_mxfp8_bf16_cutedsl` | `Sm100_Mxfp8_Mxfp8_Bf16_Cutedsl_MegaMoeConfig` | | `sm90_pull_fp8` | `sm90_fp8_fp8_bf16_pull_cutedsl` | `Sm90_Fp8_Fp8_Bf16_PullCutedsl_MegaMoeConfig` | | `sm90_push_fp8` (new, from #4069) | `sm90_fp8_fp8_bf16_push_cuda` | `Sm90_Fp8_Fp8_Bf16_PushCuda_MegaMoeConfig` | Naming conventions: deep_gemm dtypes are plain `fp8`/`fp4`, matching upstream `deep_gemm.fp8_fp4_mega_moe`; the mx/nv prefixes are reserved for the cutedsl kernels' block-scaled formats. Output dtype is always bf16 — nvfp4's `combine_dtype` is comm-wire compression, not an output format. 2. **Provenance (kernel-dev view)** — vendored kernel sources in `kernel_src/` are keyed by upstream repo snapshot, not by architecture: `kernel_src/sm100/cutedsl_megamoe` moves to `kernel_src/cutedsl_megamoe` (the mother repo ships kernels for multiple arches, so an smXX level misrepresents it). Each drop mirrors the vendor repo layout — `src/` byte-for-byte upstream, all adaptation in `shim/` — and gains a `VENDOR.md` recording upstream repo/commit/sync state and pending local diffs. A new `kernel_src/README.md` states the contract explicitly: **no edits to `src/` of any kind — including docstrings, comments, and lint fixes**; tool warnings against vendored files (docstring-coverage gates, review bots) are handled by excluding the path, never by editing the file. The sm90 fork trees (`kernel_src/sm90/pull_style_cutedsl_megakernel` from #4113, `kernel_src/sm90/push_style_megamoe` from #4069) intentionally stay separate snapshots — one kernel_src dir = one upstream commit — and fold into the mother drop if/when upstream merges them. **Why:** verbatim snapshots must stay diffable against one upstream commit, and splitting vendored trees per-dtype or per-arch breaks re-sync; meanwhile users navigate by architecture and dtype, not by which vendor repo a kernel came from. Putting each concern where its audience looks resolves the tension. The layout rule is documented in `docs/design_docs/moe_ep_architecture.md`, and it is what makes new backend families routine — demonstrated in this very PR by the SM90 push-style incorporation below, and next by the follow-up backend-family PRs (SM100 BF16 #4386, SM120 MXFP8). ## Directories affected All changes live under `flashinfer/moe_ep/` plus its tests and docs: - `backends/mega/kernel/sm100/{fp8_fp4_bf16_deepgemm,nvfp4_nvfp4_bf16_cutedsl,mxfp8_mxfp8_bf16_cutedsl}/` and `backends/mega/kernel/sm90/{fp8_fp8_bf16_pull_cutedsl,fp8_fp8_bf16_push_cuda}/` — taxonomy backend wrappers (moved/renamed; push_cuda is new). - `kernel_src/cutedsl_megamoe/` (moved from `kernel_src/sm100/cutedsl_megamoe/`), `kernel_src/sm90/pull_style_cutedsl_megakernel/`, `kernel_src/sm90/push_style_megamoe/` (new) — provenance-keyed vendored drops, each with `VENDOR.md`; new `kernel_src/README.md` states the no-edits contract. - `backends/mega/kernel/tuning.py` + per-backend `tuner.py` files — tuning machinery moved out of `tune.py` (now a CLI shim). - `core/kernel/registry.py`, `moe_ep/__init__.py` — deprecated-alias resolution and re-exports. - `tests/moe_ep/`, `docs/design_docs/moe_ep_{architecture,runbook}.md`, `pyproject.toml`/`.pre-commit-config.yaml` excludes, `run_tests.sh` (new 2-GPU `sm90_push` target). ## Test results - **Full `run_tests.sh` matrix — all 12 targets green** on 4xH100 (job 2389821, 2026-08-13), including the new `sm90_push` Hopper target and the fault-tolerance suites after the deadlock fixes. - **B200** (jobs 2388315/2388326): registry/alias smoke, deprecated aliases, unit x3 green — 396 passed / 72 skipped (push cpu/packaging/contract tests run; Hopper-marked kernel tests skip). - **Unit target re-validated green** after the second upstream merge (job 2389880) and again after the round-2 CodeRabbit fixes (job 2389916), same 396/72 counts, B200. - **8x B200** (jobs 2384640/2384641/2384650): quant-staging sync matrix fully green on both dsl 4.6.1 and 4.7.0 (details in the vendored-sync section below). - **GB200 + B200**: mxfp8/nvfp4 multirank oracle suites with the per-cell tolerance band. - Microbenchmark re-run: no regressions vs pre-restructure reference numbers (deep_gemm parity; cutedsl kernels at or above their previous points). - `pre-commit run -a` fully green at the branch head (e9f791a). ## SM90 push-style FP8 backend (incorporates #4069) Ports #4069 (head 301f8ce; since merged to main as f9b13ef — re-diffed, byte-identical, no post-review deltas) onto the taxonomy/provenance layout, serving as the first proof of the "one taxonomy backend dir + one provenance-keyed kernel drop" recipe: - **`kernel_src/sm90/push_style_megamoe/`** — verbatim byte-for-byte drop from the PR head (`src/{a2a,fp8_gemm}` CUDA sources, `shim/`, ACKNOWLEDGEMENT.md) plus a `VENDOR.md` provenance record. - **`backends/mega/kernel/sm90/fp8_fp8_bf16_push_cuda/`** — the five wrapper files relocated from upstream's flat `kernel/sm90_push_fp8/`, config renamed to `Sm90_Fp8_Fp8_Bf16_PushCuda_MegaMoeConfig`, registered with `deprecated_aliases=("sm90_push_fp8",)`. - **Core deltas carried from the PR:** `mega_layer.py` allocates the output before `stage_inputs`; pyproject package-data ships the drop's `.cu`/`.cuh` for non-editable installs; the `isolated_deep_gemm_cache` conftest fixture; the mega-layer allocation-order regression test. - **Tests:** the nine sm90_push_fp8 test files (names kept to minimize re-sync friction) rewritten to the taxonomy. Deviation from upstream: `run_tests.sh` exposes `sm90_push` as its own 2-GPU Hopper target instead of folding it into multirank — on non-Hopper nodes the arch-marked files collect 0 tests and torchrun turns pytest exit 5 into a failure. ## CuTe-DSL 4.7 quant-staging fix (vendored sync) The `CUDA_ERROR_MISALIGNED_ADDRESS` crash on cutlass-dsl 4.7.0 — which presented as a deep_gemm mega multirank failure — was root-caused to the **fused bf16→quantized activation staging** (`DataPreprocess` in the vendored cutedsl_megamoe tree), which every mega staging path shares, deep_gemm's included. The kernel team's fix is synced in as a single-file partial re-sync per the vendoring policy: - `kernel_src/cutedsl_megamoe/src/src/inputs_process.py` + `src/common/host_utils.py` taken **verbatim** from upstream `bangyus/cutedsl_megamoe @ 50117315d`, recorded in `VENDOR.md` under pending-diffs (resolves at the next full re-sync). The mxfp8 quant kernel is reworked so each lane owns one contiguous 16-byte fp8 store (adjacent lanes reduce the 32-element block amax via `shuffle_sync_bfly`, even lane writes the E8M0 scale), and `__init__` gains a hidden-size row-alignment guard. - Also fixes a stale pre-commit exclude left by the directory move (`kernel_src/sm100/cutedsl_megamoe` → `kernel_src/cutedsl_megamoe`) so hooks stop reformatting the verbatim `src/` tree. Validated on 8x B200 (jobs 2384640/2384641/2384650), full matrix green on **both** DSL versions: | section | dsl 4.6.1 | dsl 4.7.0 | |---|---|---| | drop's own harness (`python -m src.inputs_process`: bit-exact scales + SNR vs reference, nvfp4 offline/online + mxfp8) | 3/3 | 3/3 | | `test_fused_quant_stage.py` | 11/11 | 11/11 | | mega multirank x4 ranks (deep_gemm + nvfp4 + mxfp8) | 20/rank | 20/rank | | single-rank kernel-vs-reference oracles | 6/6 | 6/6 | The deep_gemm multirank suite previously crashed deterministically on 4.7.0; it now passes there. On the strength of this, the runbook's temporary `==4.6.1` pin is lifted (see the DSL guidance bullet below). ## Also in this PR - **Per-backend tuners.** `flashinfer/moe_ep/tune.py` becomes a pure CLI shim (surface unchanged: `python -m flashinfer.moe_ep.tune`); dtype-specific tuning moves into the backends (`sm100/{nvfp4,mxfp8}.../tuner.py`), shared sweep machinery (dist lifecycle, skewed restage, schedule grid, timed sweep tail) into `backends/mega/kernel/tuning.py`. - **CUTLASS DSL guidance updated (pin lifted).** The test-container recipe briefly carried a hard `nvidia-cutlass-dsl==4.6.1` pin because 4.7.0 crashed the mega multirank path; with the crash root-caused and fixed above, the runbook now allows `-U` installs again. 4.6.1 remains the perf-validated reference (pin it when producing numbers meant to compare against the TUNING.md tables); 4.7.0 is correctness-validated. The library's supported floor remains 4.5.2 (the MR!27 WAR already in main). - **Per-cell bf16 term-magnitude tolerance band** for the mxfp8 multirank oracle compares — a principled per-cell bound derived from the bf16 accumulation term magnitudes, replacing the global rtol that produced rare single-cell false failures. Validated on GB200 and B200. - **One-direction import layering rules** codified in the architecture doc, with all `cutedsl_megamoe` access routed through the drop's `__init__` rather than deep-path imports. ## Merge with upstream/main and follow-up fixes The branch is merged up to upstream/main in two steps. First to aaf97df (95 commits, incl. the v0.6.17 release line): conflict resolution keeps the restructure spellings everywhere; upstream's one real kernel advance in the moved tree — the 4fbac49 singleton-expert TMA-modes fix (#4296) — is ported onto the renamed paths and recorded in `VENDOR.md`. Notable upstream picks now in-tree: `BootstrapConfig.device` (#4348) and the E_local=1 nvfp4 oracle regression test. Second merge to 2febce5 (13 commits), resolving the conflicts created when #4069 itself squash-merged upstream (f9b13ef) with the same moe_ep files in the pre-restructure flat layout. Every conflict resolves to the taxonomy spellings (upstream's side is the flat spelling of content this branch already carries); upstream's flat `backends/mega/kernel/sm90_push_fp8/` wrapper and its re-folding of `sm90_push` into the multirank target are dropped in favor of this branch's layout. The vendored push drop was re-diffed against the merged SHA: byte-for-byte identical, no post-review deltas (recorded in `VENDOR.md`). Post-merge hardening found and fixed by full-suite runs: - **Merge fallout:** auto-merged regions had re-introduced pre-restructure `kernel_src.sm100.cutedsl_megamoe` spellings in 12 files, silently skipping entire GPU test files via `importorskip`; restored, and upstream's re-added flat `sm90_pull_fp8/` wrapper removed. - **FT test deadlocks (4xH100):** the fault-tolerance multirank test's evicted victim ran a collective `destroy()` against the survivors' barrier sequence, deadlocking until the NCCL watchdog — the victim tail now mirrors the survivors' barrier→destroy→barrier shape. The FT smoke's survivors now keep forwarding past the kill window so they actually observe the fault, and `run_tests.sh` judges the smoke by counting `SMOKE_RESULT` markers (torchrun interleaves lines). - **Unit-suite crasher isolation:** the long-known in-suite-only interpreter abort (heap corruption accumulating over the ~200-test single-process run, firing during a plain module import or in CPython teardown) is worked around by running the trigger test in its own pytest process and exiting the unit invocations via `os._exit(pytest_rc)`; rationale in the runbook, root cause tracked (needs ASAN). All tests pass — this is process-teardown hygiene, not a kernel bug. **CodeRabbit review responses.** Two rounds of actionable findings are fixed in-branch (640b75f, 57926a9) — highlights from round 2: the push packaging test's import-boundary gate was building the pre-taxonomy flat backend path and passing vacuously (fixed, now validates all 5 wrapper files); the test baseline's weight cache gains weakref eviction; `cutedsl_megamoe/shim/__main__.py` added so the documented `python -m ...shim` commands resolve; the cutedsl_megamoe `VENDOR.md` provenance TODOs are filled. Findings inside verbatim-vendored `kernel_src/**/src/` trees are deliberately not patched locally — they route upstream per the vendoring policy in `kernel_src/README.md`. **Lint.** `pre-commit run -a` is fully green (clang-format, mypy, ruff check/format, whitespace hooks). The final e9f791a is a pure ruff-format pass over 13 moe_ep files — line wraps where the longer taxonomy class names pushed calls past the limit. The vendored `src/` trees are untouched by hooks (the exclude set holds). ## Backward compatibility External callers keep working unchanged — both the old config-class names and the old kernel_name strings remain as deprecated aliases: - **Config classes**: `DeepGemmMegaMoeConfig`, `Nvfp4CutedslMegaMoeConfig`, `Mxfp8CutedslMegaMoeConfig`, `Sm90PullFp8MegaMoeConfig`, and `Sm90PushFp8MegaMoeConfig` are plain aliases of the new `Sm<arch>_..._MegaMoeConfig` classes, defined (with a removal note) in `flashinfer/moe_ep/__init__.py` right below the taxonomy imports, and still exported via `__all__`. - **Registry kernel_name strings**: `deep_gemm_mega`, `nvfp4_cutedsl`, `mxfp8_cutedsl`, `sm90_pull_fp8`, and `sm90_push_fp8` resolve to the taxonomy backends through the `deprecated_aliases=` parameter of each backend's `@register_mega_kernel(...)` decoration; the resolution machinery lives in `flashinfer/moe_ep/core/kernel/registry.py`. Using one emits a `DeprecationWarning`, and aliases are excluded from the available-kernels listing. - Both alias families WILL BE REMOVED in a future release (noted at both locations above). ## Testing - Directory moves and renames are behavior-preserving by construction; registry tests exercise both the taxonomy names and the deprecated aliases (alias use warns; the kernel listing shows taxonomy names only). - Full `run_tests.sh` matrix (all 12 targets) green on 4xH100 (job 2389821); B200 unit/registry/alias validation (jobs 2388315/2388326) — see Test results above. - The quant-staging sync validated on both dsl 4.6.1 and 4.7.0 (matrix above); mxfp8/nvfp4 multirank oracle suites validated on GB200 and B200. - The standalone MoE-EP microbenchmark was re-run against this branch with no regressions vs the pre-restructure reference numbers (deep_gemm parity; cutedsl kernels at or above their previous points). ## Relation to other PRs Re-layering on top of #4113 (SM90 pull-style FP8 backend, merged) and incorporating #4069 (SM90 push-style FP8 backend, merged upstream 2026-08-12; the vendored drop was re-diffed against the merged SHA f9b13ef and is byte-identical). This is the base branch for the upcoming backend-family PRs — SM100 BF16 (#4386) and SM120 MXFP8 — each of which adds one taxonomy backend directory plus one provenance-keyed kernel drop in the shape this restructure establishes. Both follow-up branches are already rebased onto this branch's head (unit target green on each), so they apply as exactly their backend-specific commits once this merges. --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Co-authored-by: Md Anik <mhoqueanik@cw-dfw-cs-001-login-01.cm.cluster> Co-authored-by: Md Saidul Hoque Anik <mhoqueanik@login-preos01.a51.clusters.nvidia.com>
This was referenced Aug 18, 2026
feat(moe_ep): SM90 NVFP4 push mega-MoE backend with W4A8 compute and load-time weight policies
#4589
Open
jefby
pushed a commit
to jefby/flashinfer
that referenced
this pull request
Aug 19, 2026
…ush-style FP8 backend; sync CuTe-DSL 4.7 quant-staging fix (flashinfer-ai#4449) ## Summary Three things: a layout/naming refactor of `flashinfer.moe_ep`'s mega-kernel layer, the incorporation of the SM90 push-style FP8 backend (flashinfer-ai#4069, since merged upstream) as the first new backend added in the restructured shape, and one vendored-kernel sync that fixes the fused activation-quant staging crash on CuTe-DSL 4.7 — un-blocking 4.7.x and lifting the temporary `==4.6.1` pin. The branch is merged up to upstream/main (2febce5, past the v0.6.17 line and the flashinfer-ai#4069 squash). The refactor organizes the layer around two orthogonal views: 1. **Taxonomy (user view)** — backends move to `backends/mega/kernel/sm<arch>/<act_dtype>_<weight_dtype>_<out_dtype>_<kernel_style>/`, and registry `kernel_name` strings plus config classes carry the same fully-qualified names. One glance at a name now tells you the architecture, the activation/weight/output dtypes, and the kernel style: | old kernel_name | new kernel_name | new config class | |---|---|---| | `deep_gemm_mega` | `sm100_fp8_fp4_bf16_deepgemm` | `Sm100_Fp8_Fp4_Bf16_Deepgemm_MegaMoeConfig` | | `nvfp4_cutedsl` | `sm100_nvfp4_nvfp4_bf16_cutedsl` | `Sm100_Nvfp4_Nvfp4_Bf16_Cutedsl_MegaMoeConfig` | | `mxfp8_cutedsl` | `sm100_mxfp8_mxfp8_bf16_cutedsl` | `Sm100_Mxfp8_Mxfp8_Bf16_Cutedsl_MegaMoeConfig` | | `sm90_pull_fp8` | `sm90_fp8_fp8_bf16_pull_cutedsl` | `Sm90_Fp8_Fp8_Bf16_PullCutedsl_MegaMoeConfig` | | `sm90_push_fp8` (new, from flashinfer-ai#4069) | `sm90_fp8_fp8_bf16_push_cuda` | `Sm90_Fp8_Fp8_Bf16_PushCuda_MegaMoeConfig` | Naming conventions: deep_gemm dtypes are plain `fp8`/`fp4`, matching upstream `deep_gemm.fp8_fp4_mega_moe`; the mx/nv prefixes are reserved for the cutedsl kernels' block-scaled formats. Output dtype is always bf16 — nvfp4's `combine_dtype` is comm-wire compression, not an output format. 2. **Provenance (kernel-dev view)** — vendored kernel sources in `kernel_src/` are keyed by upstream repo snapshot, not by architecture: `kernel_src/sm100/cutedsl_megamoe` moves to `kernel_src/cutedsl_megamoe` (the mother repo ships kernels for multiple arches, so an smXX level misrepresents it). Each drop mirrors the vendor repo layout — `src/` byte-for-byte upstream, all adaptation in `shim/` — and gains a `VENDOR.md` recording upstream repo/commit/sync state and pending local diffs. A new `kernel_src/README.md` states the contract explicitly: **no edits to `src/` of any kind — including docstrings, comments, and lint fixes**; tool warnings against vendored files (docstring-coverage gates, review bots) are handled by excluding the path, never by editing the file. The sm90 fork trees (`kernel_src/sm90/pull_style_cutedsl_megakernel` from flashinfer-ai#4113, `kernel_src/sm90/push_style_megamoe` from flashinfer-ai#4069) intentionally stay separate snapshots — one kernel_src dir = one upstream commit — and fold into the mother drop if/when upstream merges them. **Why:** verbatim snapshots must stay diffable against one upstream commit, and splitting vendored trees per-dtype or per-arch breaks re-sync; meanwhile users navigate by architecture and dtype, not by which vendor repo a kernel came from. Putting each concern where its audience looks resolves the tension. The layout rule is documented in `docs/design_docs/moe_ep_architecture.md`, and it is what makes new backend families routine — demonstrated in this very PR by the SM90 push-style incorporation below, and next by the follow-up backend-family PRs (SM100 BF16 flashinfer-ai#4386, SM120 MXFP8). ## Directories affected All changes live under `flashinfer/moe_ep/` plus its tests and docs: - `backends/mega/kernel/sm100/{fp8_fp4_bf16_deepgemm,nvfp4_nvfp4_bf16_cutedsl,mxfp8_mxfp8_bf16_cutedsl}/` and `backends/mega/kernel/sm90/{fp8_fp8_bf16_pull_cutedsl,fp8_fp8_bf16_push_cuda}/` — taxonomy backend wrappers (moved/renamed; push_cuda is new). - `kernel_src/cutedsl_megamoe/` (moved from `kernel_src/sm100/cutedsl_megamoe/`), `kernel_src/sm90/pull_style_cutedsl_megakernel/`, `kernel_src/sm90/push_style_megamoe/` (new) — provenance-keyed vendored drops, each with `VENDOR.md`; new `kernel_src/README.md` states the no-edits contract. - `backends/mega/kernel/tuning.py` + per-backend `tuner.py` files — tuning machinery moved out of `tune.py` (now a CLI shim). - `core/kernel/registry.py`, `moe_ep/__init__.py` — deprecated-alias resolution and re-exports. - `tests/moe_ep/`, `docs/design_docs/moe_ep_{architecture,runbook}.md`, `pyproject.toml`/`.pre-commit-config.yaml` excludes, `run_tests.sh` (new 2-GPU `sm90_push` target). ## Test results - **Full `run_tests.sh` matrix — all 12 targets green** on 4xH100 (job 2389821, 2026-08-13), including the new `sm90_push` Hopper target and the fault-tolerance suites after the deadlock fixes. - **B200** (jobs 2388315/2388326): registry/alias smoke, deprecated aliases, unit x3 green — 396 passed / 72 skipped (push cpu/packaging/contract tests run; Hopper-marked kernel tests skip). - **Unit target re-validated green** after the second upstream merge (job 2389880) and again after the round-2 CodeRabbit fixes (job 2389916), same 396/72 counts, B200. - **8x B200** (jobs 2384640/2384641/2384650): quant-staging sync matrix fully green on both dsl 4.6.1 and 4.7.0 (details in the vendored-sync section below). - **GB200 + B200**: mxfp8/nvfp4 multirank oracle suites with the per-cell tolerance band. - Microbenchmark re-run: no regressions vs pre-restructure reference numbers (deep_gemm parity; cutedsl kernels at or above their previous points). - `pre-commit run -a` fully green at the branch head (e9f791a). ## SM90 push-style FP8 backend (incorporates flashinfer-ai#4069) Ports flashinfer-ai#4069 (head 301f8ce; since merged to main as f9b13ef — re-diffed, byte-identical, no post-review deltas) onto the taxonomy/provenance layout, serving as the first proof of the "one taxonomy backend dir + one provenance-keyed kernel drop" recipe: - **`kernel_src/sm90/push_style_megamoe/`** — verbatim byte-for-byte drop from the PR head (`src/{a2a,fp8_gemm}` CUDA sources, `shim/`, ACKNOWLEDGEMENT.md) plus a `VENDOR.md` provenance record. - **`backends/mega/kernel/sm90/fp8_fp8_bf16_push_cuda/`** — the five wrapper files relocated from upstream's flat `kernel/sm90_push_fp8/`, config renamed to `Sm90_Fp8_Fp8_Bf16_PushCuda_MegaMoeConfig`, registered with `deprecated_aliases=("sm90_push_fp8",)`. - **Core deltas carried from the PR:** `mega_layer.py` allocates the output before `stage_inputs`; pyproject package-data ships the drop's `.cu`/`.cuh` for non-editable installs; the `isolated_deep_gemm_cache` conftest fixture; the mega-layer allocation-order regression test. - **Tests:** the nine sm90_push_fp8 test files (names kept to minimize re-sync friction) rewritten to the taxonomy. Deviation from upstream: `run_tests.sh` exposes `sm90_push` as its own 2-GPU Hopper target instead of folding it into multirank — on non-Hopper nodes the arch-marked files collect 0 tests and torchrun turns pytest exit 5 into a failure. ## CuTe-DSL 4.7 quant-staging fix (vendored sync) The `CUDA_ERROR_MISALIGNED_ADDRESS` crash on cutlass-dsl 4.7.0 — which presented as a deep_gemm mega multirank failure — was root-caused to the **fused bf16→quantized activation staging** (`DataPreprocess` in the vendored cutedsl_megamoe tree), which every mega staging path shares, deep_gemm's included. The kernel team's fix is synced in as a single-file partial re-sync per the vendoring policy: - `kernel_src/cutedsl_megamoe/src/src/inputs_process.py` + `src/common/host_utils.py` taken **verbatim** from upstream `bangyus/cutedsl_megamoe @ 50117315d`, recorded in `VENDOR.md` under pending-diffs (resolves at the next full re-sync). The mxfp8 quant kernel is reworked so each lane owns one contiguous 16-byte fp8 store (adjacent lanes reduce the 32-element block amax via `shuffle_sync_bfly`, even lane writes the E8M0 scale), and `__init__` gains a hidden-size row-alignment guard. - Also fixes a stale pre-commit exclude left by the directory move (`kernel_src/sm100/cutedsl_megamoe` → `kernel_src/cutedsl_megamoe`) so hooks stop reformatting the verbatim `src/` tree. Validated on 8x B200 (jobs 2384640/2384641/2384650), full matrix green on **both** DSL versions: | section | dsl 4.6.1 | dsl 4.7.0 | |---|---|---| | drop's own harness (`python -m src.inputs_process`: bit-exact scales + SNR vs reference, nvfp4 offline/online + mxfp8) | 3/3 | 3/3 | | `test_fused_quant_stage.py` | 11/11 | 11/11 | | mega multirank x4 ranks (deep_gemm + nvfp4 + mxfp8) | 20/rank | 20/rank | | single-rank kernel-vs-reference oracles | 6/6 | 6/6 | The deep_gemm multirank suite previously crashed deterministically on 4.7.0; it now passes there. On the strength of this, the runbook's temporary `==4.6.1` pin is lifted (see the DSL guidance bullet below). ## Also in this PR - **Per-backend tuners.** `flashinfer/moe_ep/tune.py` becomes a pure CLI shim (surface unchanged: `python -m flashinfer.moe_ep.tune`); dtype-specific tuning moves into the backends (`sm100/{nvfp4,mxfp8}.../tuner.py`), shared sweep machinery (dist lifecycle, skewed restage, schedule grid, timed sweep tail) into `backends/mega/kernel/tuning.py`. - **CUTLASS DSL guidance updated (pin lifted).** The test-container recipe briefly carried a hard `nvidia-cutlass-dsl==4.6.1` pin because 4.7.0 crashed the mega multirank path; with the crash root-caused and fixed above, the runbook now allows `-U` installs again. 4.6.1 remains the perf-validated reference (pin it when producing numbers meant to compare against the TUNING.md tables); 4.7.0 is correctness-validated. The library's supported floor remains 4.5.2 (the MR!27 WAR already in main). - **Per-cell bf16 term-magnitude tolerance band** for the mxfp8 multirank oracle compares — a principled per-cell bound derived from the bf16 accumulation term magnitudes, replacing the global rtol that produced rare single-cell false failures. Validated on GB200 and B200. - **One-direction import layering rules** codified in the architecture doc, with all `cutedsl_megamoe` access routed through the drop's `__init__` rather than deep-path imports. ## Merge with upstream/main and follow-up fixes The branch is merged up to upstream/main in two steps. First to aaf97df (95 commits, incl. the v0.6.17 release line): conflict resolution keeps the restructure spellings everywhere; upstream's one real kernel advance in the moved tree — the 4fbac49 singleton-expert TMA-modes fix (flashinfer-ai#4296) — is ported onto the renamed paths and recorded in `VENDOR.md`. Notable upstream picks now in-tree: `BootstrapConfig.device` (flashinfer-ai#4348) and the E_local=1 nvfp4 oracle regression test. Second merge to 2febce5 (13 commits), resolving the conflicts created when flashinfer-ai#4069 itself squash-merged upstream (f9b13ef) with the same moe_ep files in the pre-restructure flat layout. Every conflict resolves to the taxonomy spellings (upstream's side is the flat spelling of content this branch already carries); upstream's flat `backends/mega/kernel/sm90_push_fp8/` wrapper and its re-folding of `sm90_push` into the multirank target are dropped in favor of this branch's layout. The vendored push drop was re-diffed against the merged SHA: byte-for-byte identical, no post-review deltas (recorded in `VENDOR.md`). Post-merge hardening found and fixed by full-suite runs: - **Merge fallout:** auto-merged regions had re-introduced pre-restructure `kernel_src.sm100.cutedsl_megamoe` spellings in 12 files, silently skipping entire GPU test files via `importorskip`; restored, and upstream's re-added flat `sm90_pull_fp8/` wrapper removed. - **FT test deadlocks (4xH100):** the fault-tolerance multirank test's evicted victim ran a collective `destroy()` against the survivors' barrier sequence, deadlocking until the NCCL watchdog — the victim tail now mirrors the survivors' barrier→destroy→barrier shape. The FT smoke's survivors now keep forwarding past the kill window so they actually observe the fault, and `run_tests.sh` judges the smoke by counting `SMOKE_RESULT` markers (torchrun interleaves lines). - **Unit-suite crasher isolation:** the long-known in-suite-only interpreter abort (heap corruption accumulating over the ~200-test single-process run, firing during a plain module import or in CPython teardown) is worked around by running the trigger test in its own pytest process and exiting the unit invocations via `os._exit(pytest_rc)`; rationale in the runbook, root cause tracked (needs ASAN). All tests pass — this is process-teardown hygiene, not a kernel bug. **CodeRabbit review responses.** Two rounds of actionable findings are fixed in-branch (640b75f, 57926a9) — highlights from round 2: the push packaging test's import-boundary gate was building the pre-taxonomy flat backend path and passing vacuously (fixed, now validates all 5 wrapper files); the test baseline's weight cache gains weakref eviction; `cutedsl_megamoe/shim/__main__.py` added so the documented `python -m ...shim` commands resolve; the cutedsl_megamoe `VENDOR.md` provenance TODOs are filled. Findings inside verbatim-vendored `kernel_src/**/src/` trees are deliberately not patched locally — they route upstream per the vendoring policy in `kernel_src/README.md`. **Lint.** `pre-commit run -a` is fully green (clang-format, mypy, ruff check/format, whitespace hooks). The final e9f791a is a pure ruff-format pass over 13 moe_ep files — line wraps where the longer taxonomy class names pushed calls past the limit. The vendored `src/` trees are untouched by hooks (the exclude set holds). ## Backward compatibility External callers keep working unchanged — both the old config-class names and the old kernel_name strings remain as deprecated aliases: - **Config classes**: `DeepGemmMegaMoeConfig`, `Nvfp4CutedslMegaMoeConfig`, `Mxfp8CutedslMegaMoeConfig`, `Sm90PullFp8MegaMoeConfig`, and `Sm90PushFp8MegaMoeConfig` are plain aliases of the new `Sm<arch>_..._MegaMoeConfig` classes, defined (with a removal note) in `flashinfer/moe_ep/__init__.py` right below the taxonomy imports, and still exported via `__all__`. - **Registry kernel_name strings**: `deep_gemm_mega`, `nvfp4_cutedsl`, `mxfp8_cutedsl`, `sm90_pull_fp8`, and `sm90_push_fp8` resolve to the taxonomy backends through the `deprecated_aliases=` parameter of each backend's `@register_mega_kernel(...)` decoration; the resolution machinery lives in `flashinfer/moe_ep/core/kernel/registry.py`. Using one emits a `DeprecationWarning`, and aliases are excluded from the available-kernels listing. - Both alias families WILL BE REMOVED in a future release (noted at both locations above). ## Testing - Directory moves and renames are behavior-preserving by construction; registry tests exercise both the taxonomy names and the deprecated aliases (alias use warns; the kernel listing shows taxonomy names only). - Full `run_tests.sh` matrix (all 12 targets) green on 4xH100 (job 2389821); B200 unit/registry/alias validation (jobs 2388315/2388326) — see Test results above. - The quant-staging sync validated on both dsl 4.6.1 and 4.7.0 (matrix above); mxfp8/nvfp4 multirank oracle suites validated on GB200 and B200. - The standalone MoE-EP microbenchmark was re-run against this branch with no regressions vs the pre-restructure reference numbers (deep_gemm parity; cutedsl kernels at or above their previous points). ## Relation to other PRs Re-layering on top of flashinfer-ai#4113 (SM90 pull-style FP8 backend, merged) and incorporating flashinfer-ai#4069 (SM90 push-style FP8 backend, merged upstream 2026-08-12; the vendored drop was re-diffed against the merged SHA f9b13ef and is byte-identical). This is the base branch for the upcoming backend-family PRs — SM100 BF16 (flashinfer-ai#4386) and SM120 MXFP8 — each of which adds one taxonomy backend directory plus one provenance-keyed kernel drop in the shape this restructure establishes. Both follow-up branches are already rebased onto this branch's head (unit target green on each), so they apply as exactly their backend-specific commits once this merges. --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Co-authored-by: Md Anik <mhoqueanik@cw-dfw-cs-001-login-01.cm.cluster> Co-authored-by: Md Saidul Hoque Anik <mhoqueanik@login-preos01.a51.clusters.nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds the first Hopper mega kernel to the moe_ep stack: the kernel team's SM90 pull-style FP8 CuTeDSL MegaMoE (Vincent's
hopper_megamoe, kernel repo commit1275b8b), integrated as a newsm90_pull_fp8backend behind the existingMegaKernelBackendcontract — fused NVSHMEM dispatch + FC1 + SwiGLU + FC2 + combine in one launch, FP8 E4M3/E5M2 with per-tensor or DeepGEMM-style blockwise scaling, native and swap-A/B layouts.New oracle test paths: every mega kernel (SM90 and the three existing SM100 paths) now has a multi-rank torch oracle — real cross-rank EP launches judged against pure-torch global math instead of same-kernel parity.
What's in here
Two-tree
kernel_src/restructure — the existing Blackwell tree moves tokernel_src/sm100/cutedsl_megamoe/; the SM90 drop lands askernel_src/sm90/pull_style_cutedsl_megakernel/(a fork of the same kernel repo, shared runtime intentionally duplicated at its own drop revision). The trees expose colliding top-level module names, so eachshim/_paths.pynow guards process exclusivity (a process runs on one arch anyway). All paths/docs/lint excludes updated; sm100 tree is byte-identical after the move.Vendored drop + shim — verbatim
src/plus our adapter layer (shim/hopper_fp8.pymirrors the sm100mxfp8.pydesign: frozen validated config, lazycute.compile, launch cache, symm-buffer allocator + compute entry per the runbook mega-kernel contract).Backend + tests —
Sm90PullFp8MegaMoeConfig(registered"sm90_pull_fp8"), weight preprocessing for both scale modes (gate/up interleave-8, K-major without repack), sm_90 arch gate, runtime requirements, and a core fix:_init_nvshmem_after_distno longer imports a kernel tree (it would have broken every sm90 multirank session via the exclusivity guard).Benchmark + TUNING.md —
benchmarks/bench_moe_ep_sm90_mega.pyreproduces the kernel team's 7-point token sweep through the FI layer; measured results and methodology documented in the sm90 tree'sTUNING.md.Multi-rank torch oracles — parity tests can't catch a mega kernel that is wrong but self-consistent at
world_size > 1(comm + compute are fused, so both sides of a parity test run the same CUDA kernel). Newtest_moe_ep_*_mega_multirank_torch_oracletests close that gap: each rank launches the fused kernel with real cross-rank NVSHMEM traffic, all-gathers the actual operands the kernel consumed (plain pre-swizzle weight legs; for mxfp8/sm90 also staged payloads + routing), and checks its own output slice against torch math over the global expert set, with forced cross-rank routing. Covers sm90 ({per_tensor, blockwise} × {native, swap_ab}) and all three SM100 paths, including the variant knobs: nvfp4in_kernel_fc2_reduce+ quantized combine wires (16e2m1xbf16,32e4m3xe8m0, wire modeled exactly viacombine_roundtrip_to_fp32— newly exported through the sm100 shim boundary), mxfp8in_kernel_fc2_reduce. All picked up by the existingrun_tests.sh mega/mega_sm90targets.Docs —
docs/design_docs/moe_ep_architecture.mdgains a "Torch oracles" section (methodology, independence contract, per-kernel last-passed table) and a refreshedrun_tests.shtarget list.Testing
SM90 (4×H100 80GB, EP4):
run_tests.sh unit— 189+9 tests pass (new host-only config/registry tests included)run_tests.sh oracle_sm90— kernel vs the drop's fp32 torch reference,{per_tensor, blockwise} × {native, swap_ab}: 5/5 passrun_tests.sh mega_sm90— 4-GPUMoEEpLayerparity vs a direct-shim session with forced cross-rank routing: bit-exact on separate-reduce paths (incl. pre-staged fp8 inputs and repeat-forward launch-cache guards), roundoff-envelope forin_kernel_fc2_reduce: 5/5 × 4 ranks pass; plus the multi-rank torch oracle, 4/4 params × 4 ranks passSM100 regression (4×GB200, 2026-07-31):
run_tests.sh all— all 8 sections pass (unit 303/303, single-GPU oracles, split multirank, split-path correctness bf16/nvfp4/ht, mega multirank, smoke): the SM90 integration and restructure leave every Blackwell path unregressedrel_l2 < 0.02per rankPerformance
28-point sweep at the drop's DSv4-Pro geometry (384 experts, top-6, hidden 7168, inter 3072, 512–32K tokens/rank): the FI integration adds no measurable kernel-path overhead — within ±5% of the kernel team's own harness at most points (their numbers exclude the TopkReduce tail and report min-rank; ours don't), peak 562 TFLOPS/rank. Full tables, comparison caveats, and open items (16K-token iteration variance, fused staging port, DSL 4.6.1 A/B) in
kernel_src/sm90/pull_style_cutedsl_megakernel/TUNING.md.Notes for reviewers
kernel_src/*/src/directories are verbatim kernel-team drops — review the shim/backend layers, not those.PrequantizedMoEWeightsnot wired (bf16 canonical orpreprocess_weights=False);reuse_dispatch_warpstoken-back is perf-exercised but not yet in the correctness matrix.Summary by CodeRabbit
New Features
Performance
Documentation
Tests