feat(moe_ep): SM120 MXFP8 swap-AB CuTeDSL MegaMoE kernel - #4387
Conversation
|
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:
ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (235)
💤 Files with no reviewable changes (1)
🚧 Files skipped from review as they are similar to previous changes (86)
📝 WalkthroughWalkthroughThe PR restructures the MegaMoE backend layout with architecture- and datatype-specific naming (for example ChangesMegaMoE multi-architecture backend restructuring
Estimated code review effort: 5 (Critical) | ~180 minutes Mergeability Score: 🟠 High · up to The new SM120 backend can silently produce incorrect results or fail with illegal memory access and distributed deadlocks in supported or configurable paths. Descriptor construction, routing-buffer sizing, shared state, compiled-workspace lifetime, and collective failure handling still have concrete unresolved risks, so the PR is not ready to merge until these are fixed or explicitly accepted by owners. Sequence Diagram(s)Not applicable: this PR is a large-scale rename, restructuring, and vendored-source addition across many files without a single coherent new user-facing control flow suitable for a concise sequence diagram. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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}
There was a problem hiding this comment.
Actionable comments posted: 4
Note
Due to the large number of review comments, Critical severity comments were prioritized as inline comments.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/moe_ep/test_deep_gemm_mega_kernel_vs_reference.py (1)
226-242: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winSkip instead of erroring when the test runs outside torchrun.
The self-bootstrap path was removed, so the test now depends on the torchrun-provided rendezvous environment. Under plain
pytest,WORLD_SIZEis unset, line 226 defaults it to1, and the guard at line 227 passes.dist.init_process_group(backend="nccl")then runs withoutMASTER_ADDR/MASTER_PORTand raises, so the test errors rather than skips.run_tests.shexcludes this file fromrun_unit, but a directpytest tests/moe_ep/invocation still collects it.Gate on the presence of
WORLD_SIZE, matching the launcher requirement stated in the module docstring.🛠️ Proposed fix
- world_size = int(os.environ.get("WORLD_SIZE", "1")) + if "WORLD_SIZE" not in os.environ: + pytest.skip("requires torchrun launch (WORLD_SIZE unset)") + world_size = int(os.environ["WORLD_SIZE"]) if world_size != 1: pytest.skip("single-rank oracle test; run with --nproc_per_node=1")🤖 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_deep_gemm_mega_kernel_vs_reference.py` around lines 226 - 242, Update the test’s launcher guard before dist.init_process_group to skip when WORLD_SIZE is absent, rather than defaulting the missing environment variable to 1. Preserve the existing single-rank skip for configured world sizes other than 1, and keep torchrun-launched execution unchanged.
🟠 Major comments (34)
.pre-commit-config.yaml-21-21 (1)
21-21: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winCorrect and complete the vendored-source exclusions.
The configured SM100 path is not the documented SM100 vendor tree. The exclusions also omit the SM120 vendored
src/tree. Hooks, Ruff, and mypy will process some verbatim upstream files and can block kernel-drop re-syncs.
.pre-commit-config.yaml#L21-L21: excludeflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/,flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/, andflashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/src/.pyproject.toml#L110-L111: use those same three paths in the mypy exclusion list.pyproject.toml#L125-L126: use those same three paths in the Ruff exclusion list.Confidence: high. As per coding guidelines, kernel source changes must remain re-syncable without formatting or lint churn.
🤖 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 @.pre-commit-config.yaml at line 21, Update .pre-commit-config.yaml:21, pyproject.toml:110-111, and pyproject.toml:125-126 so the pre-commit, mypy, and Ruff exclusions consistently include flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/, flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/, and flashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/src/.Source: Coding guidelines
docs/design_docs/moe_ep_runbook.md-210-423 (1)
210-423: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winRemove the duplicate Benchmarking section.
This added section is followed by another
## Benchmarkingsection at Line 426. Keep one canonical section. Merge the SM90 additions into it so benchmark commands, version pins, and measured-result claims have one source of truth.Confidence: high. As per coding guidelines, keep documentation synchronized with code changes.
🤖 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 `@docs/design_docs/moe_ep_runbook.md` around lines 210 - 423, Remove the duplicate `## Benchmarking` section and retain one canonical benchmarking section. Merge any SM90-specific additions, commands, version pins, and measured results from the later section into the retained section, resolving conflicts so each instruction and claim has a single authoritative version.Source: Coding guidelines
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/VENDOR.md-9-12 (1)
9-12: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick winRecord the upstream snapshot provenance before merge.
VENDOR.mdstill leaves the repository URL, vendored commit, and sync date as TODOs. Without these values, reviewers cannot reproduce thesrc/drop or distinguish upstream content from local edits.Confidence: high.
🤖 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/cutedsl_megamoe/VENDOR.md` around lines 9 - 12, Update the provenance entries in VENDOR.md: replace the repository URL TODO with the canonical upstream URL, record the exact upstream commit SHA for the current src/ drop, and set Last synced to the date that snapshot was imported. Preserve the existing authors/contacts reference and document only the actual vendored snapshot metadata.flashinfer/moe_ep/backends/mega/kernel/sm100/nvfp4_nvfp4_bf16_cutedsl/backend.py-349-361 (1)
349-361: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winClear the thunk cache during workspace teardown.
Lines 305-306 store
thunkandworkspace.output_activationin_thunk_state. The thunk also retains the workspace launch inputs. This hook removes the staged-token memo but retains those tensors after the workspace pool releases the workspace.Clear
_thunk_statewhen its workspace key matchesid(workspace).Proposed fix
def _forget_workspace_state(self, workspace) -> None: + if ( + self._thunk_state is not None + and self._thunk_state[0][0] == id(workspace) + ): + self._thunk_state = None + # The fused-stage memos key on topk_idx.data_ptr(); the symmetricConfidence: High. Based on learnings, every tensor
data_ptr()used as part of a cache key must participate in eviction or invalidation.🤖 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/sm100/nvfp4_nvfp4_bf16_cutedsl/backend.py` around lines 349 - 361, Update _forget_workspace_state to remove the _thunk_state entry associated with the workspace being torn down, using id(workspace) as the key. Preserve the existing staged-token eviction and only clear the thunk state when its stored workspace key matches the current workspace.Source: Learnings
tests/moe_ep/test_moe_ep_sm120_mxfp8_cutedsl_mega_multirank.py-652-676 (1)
652-676: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winQuantitize the SM120 FC1 weight with hidden rows first.
_fc1_weight_from_w13returns the interleaved(2I, hidden)tensor, but_quantize_mxfp8_weight_k_majortreats its inputs asN, KwithKtrailing. Passfc1_interleaved[expert].transpose(0, 1)before quantizing, like the SM100 oracle.🤖 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_sm120_mxfp8_cutedsl_mega_multirank.py` around lines 652 - 676, Update the FC1 quantization in the expert loop to pass fc1_interleaved[expert].transpose(0, 1) into _quantize_mxfp8_weight_k_major, ensuring the hidden dimension is treated as K-trailing. Keep the subsequent transpose and scale-factor handling unchanged.flashinfer/moe_ep/backends/mega/kernel/sm120/mxfp8_mxfp8_bf16_cutedsl/staging.py-80-155 (1)
80-155: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winValidate pre-staged routing metadata before staging.
When
quantize_input=False, this path validates only routing shapes. It accepts invalid expert IDs and incompatible routing dtypes. The backend then copiestopk_idsinto the shim'storch.int64workspace andtopk_weightsinto its typed workspace.Validate
topk_ids.dtype,topk_weights.dtype, device placement, and each live expert ID before the copy. Reuse the common routing checks where possible. Invalid IDs can select the wrong expert or cause an invalid kernel access.Confidence: High. The SM120 shim allocates
topk_idxastorch.int64and consumes it as kernel routing metadata.🤖 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/sm120/mxfp8_mxfp8_bf16_cutedsl/staging.py` around lines 80 - 155, Extend validate_sm120_mxfp8_forward_inputs for the quantize_input=False path to validate routing metadata before staging: enforce the expected device and dtypes for topk_ids and topk_weights, and reject every live expert ID outside the configured expert range. Reuse validate_mega_forward_inputs or its routing-validation helpers where compatible, while preserving the existing shape and pre-staged MXFP8 checks.flashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/src/moe_nvfp4_swapab/run_mega_tests.sh-127-127 (1)
127-127: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winDo not run M01 with the known-broken N=128 default.
M01_single_balanced_topk2setsMEGA_NO_DIST=1but does not set--mma_tiler_mnk. The SM120 default uses N=128.VENDOR.mddocuments incorrect world-size-one output with N=128.Select a validated N=64 tiler for M01, or disable this case until the upstream kernel fixes the single-rank path. Otherwise, this functional test can report a kernel result that is known to be numerically wrong.
Confidence: High.
🤖 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/sm120/swapab_cutedsl_megakernel/src/moe_nvfp4_swapab/run_mega_tests.sh` at line 127, Update the M01_single_balanced_topk2 test entry to pass a validated N=64 --mma_tiler_mnk configuration, or remove/disable the case until the upstream single-rank kernel is fixed; do not leave it using the SM120 N=128 default.flashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/shim/mxfp8.py-290-311 (1)
290-311: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winKnob application does not respect
gate_up_clampprecedence in either entry point. Both sites accept every dataclass field as a knob, but neither keepsgate_up_clampconsistent with the value the rest of the frontend actually reads.MegaMoESm120Mxfp8Frontendstores the clamp twice, inself._config.gate_up_clampand inself._gate_up_clamp, and_mega_compile_keyand the kernel constructor read only the latter.
flashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/shim/mxfp8.py#L290-L311: after line 310 assignsnew_config, also assignself._gate_up_clamp = new_config.gate_up_clamp, and change the early-return comparison on line 306 to compare againstself.configrather thanself._config.flashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/shim/mxfp8.py#L1047-L1056: addoverrides.pop("gate_up_clamp", None)andoverrides.pop("token_back_mode", None)next to the existingin_kernel_fc2_reducepop, so all three explicit parameters ofget_symm_buffer_for_sm120_mxfp8_mega_moewin over the knob dict as the comment on lines 1053-1054 states.🤖 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/sm120/swapab_cutedsl_megakernel/shim/mxfp8.py` around lines 290 - 311, Update flashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/shim/mxfp8.py lines 290-311 in MegaMoESm120Mxfp8Frontend.apply_knobs to compare the new configuration with self.config for the no-op check, then synchronize self._gate_up_clamp with new_config.gate_up_clamp after assignment. Also update lines 1047-1056 in get_symm_buffer_for_sm120_mxfp8_mega_moe to remove gate_up_clamp and token_back_mode from overrides alongside in_kernel_fc2_reduce, ensuring all explicit parameters take precedence over knobs.flashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/shim/mxfp8.py-496-505 (1)
496-505: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winA failed compile leaves a freed workspace reachable through the cache.
Line 504 calls
_release_workspace(), which freesself._mega.shared_workspaceandself._mega.combine_root. It does not clearself._megaorself._mega_key. Lines 639-640 set the new state only aftercute.compilesucceeds.If
cute.compileat line 611,sym_zerosat line 572, orcompile_topk_reduceat line 629 raises,self._megastill references the old_CompiledMegawhose symmetric buffers are already freed, andself._mega_keystill holds the old key. A retry with the same config then matches line 500 and returns that entry. The launch runs against freed symmetric-heap storage.
set_gate_up_clampandapply_knobsavoid this because they call_invalidate_compile_cache()right after_release_workspace(). Apply the same ordering here.🛡️ Proposed fix to invalidate before releasing
ensure_not_capturing("cute.compile + symmetric-heap allocation") self._release_workspace() + # Drop the cache entry before any step below can raise: the released + # workspace must never stay reachable through _mega / _mega_key. + self._invalidate_compile_cache() self._assert_mirrored_constants()🤖 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/sm120/swapab_cutedsl_megakernel/shim/mxfp8.py` around lines 496 - 505, In _ensure_mega_compiled, invalidate the compiled cache immediately before calling _release_workspace(), matching the ordering used by set_gate_up_clamp and apply_knobs. Ensure both self._mega and self._mega_key are cleared before any potentially failing allocation or compilation step, while preserving the existing cache-return path for valid matching entries.flashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/src/moe_mxfp8_glu/epilogue_mxfp8.py-762-773 (1)
762-773: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winThe
rmem_sfif-chain silently drops scales beyond subtile 3.
rmem_sfis allocated with 4 entries and the chain handlessubtile_idx0..3.self._subtile_cntiscta_tile_n // 2 // EpilogueTileN, so a largercta_tile_nproduces more subtiles. Those extraqpvscalevalues are then discarded without any error, and_stg_sf_fc1writes stale scale factors.Derive the buffer length and the loop bound from
self._subtile_cnt, as_write_sf_fc2_bufferalready does for the FC2 side.♻️ Proposed fix
- if subtile_idx == 0: - rmem_sf[0] = qpvscale - elif subtile_idx == 1: - rmem_sf[1] = qpvscale - elif subtile_idx == 2: - rmem_sf[2] = qpvscale - elif subtile_idx == 3: - rmem_sf[3] = qpvscale + for j in cutlass.range_constexpr(self._subtile_cnt): + if subtile_idx == cutlass.Int32(j): + rmem_sf[j] = qpvscaleThe matching allocation at line 578 must use
cute.make_layout(self._subtile_cnt), and_stg_sf_fc1must storeself._subtile_cntentries instead of a fixed 4.🤖 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/sm120/swapab_cutedsl_megakernel/src/moe_mxfp8_glu/epilogue_mxfp8.py` around lines 762 - 773, Update the FC1 scale-factor handling around rmem_sf allocation and _stg_sf_fc1: derive the layout length and storage loop bound from self._subtile_cnt, using cute.make_layout(self._subtile_cnt) and writing every qpvscale entry through that count instead of a fixed four-entry if-chain. Preserve the existing per-subtile scale generation and FC2 sizing pattern in _write_sf_fc2_buffer.flashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/src/moe_sm120_mxfp8_swapab/mega_runner.py-1400-1417 (1)
1400-1417: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winThe CUDA-event path calls
all_gatherwithout checking thattorch.distributedis initialized.Every other collective in this file is guarded (see lines 887-891, 1266-1269, 2101-2104, and the
world_size == 1branch at 1759). This one is not. WithMEGA_NO_DIST=1and--use_cuda_events,torch.distributed.all_gatherraises because no process group exists.🐛 Proposed fix
- gathered = [torch.empty_like(local_us) for _ in range(self.world_size)] - torch.distributed.all_gather(gathered, local_us) + if ( + torch.distributed.is_available() + and torch.distributed.is_initialized() + ): + gathered = [ + torch.empty_like(local_us) for _ in range(self.world_size) + ] + torch.distributed.all_gather(gathered, local_us) + else: + gathered = [local_us] if self.rank == 0:🤖 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/sm120/swapab_cutedsl_megakernel/src/moe_sm120_mxfp8_swapab/mega_runner.py` around lines 1400 - 1417, Guard the CUDA-event aggregation in the shown timing path before calling torch.distributed.all_gather, using the same distributed-initialization and single-rank handling established by the other collectives in this file. When distributed execution is unavailable, compute critical_us from local_us directly; otherwise preserve the existing gathered amax calculation.flashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/src/moe_mxfp8_glu/epilogue_mxfp8.py-206-209 (1)
206-209: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winLine 209 discards the
_overlapping_accumguard.Lines 206-208 compute
_overlapping_accumfromallow_overlap_accand from the shape checkself._cta_tile_n == EpiWarpCount * EpilogueTileN * 2. Line 209 then forces it toTrueunconditionally. Every downstream decision that depends on the flag (_num_acc_pipeline_stages,_num_accumulator_tmem_cols, the odd/even subtile walk, the TMEM stage stride inkernel_mxfp8_glu_fc12.py) assumes the guarded shape relation holds. A configuration that fails the check now silently takes the overlap path with mismatched TMEM column arithmetic.If overlap is always required for this kernel, replace the computation with an explicit validation error. If it is a debug override, gate it behind a named flag.
🐛 Proposed fix
- self._overlapping_accum = allow_overlap_acc and ( - self._cta_tile_n == EpiWarpCount * EpilogueTileN * 2 - ) - self._overlapping_accum = True + if self._cta_tile_n != EpiWarpCount * EpilogueTileN * 2: + raise ValueError( + "GluMxfp8Epilogue requires cta_tile_n == " + f"{EpiWarpCount * EpilogueTileN * 2}; got {self._cta_tile_n}." + ) + self._overlapping_accum = True🤖 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/sm120/swapab_cutedsl_megakernel/src/moe_mxfp8_glu/epilogue_mxfp8.py` around lines 206 - 209, Remove the unconditional `self._overlapping_accum = True` assignment so the guard computed from `allow_overlap_acc` and the `_cta_tile_n` shape check remains authoritative. If overlap must always be enabled, replace the override with explicit validation that rejects configurations failing the shape relation; otherwise expose any debug override through a named flag.flashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/src/moe_sm120_mxfp8_swapab/moe_persistent_scheduler.py-1205-1212 (1)
1205-1212: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winDerive the static scheduler params shape from the params object.
When
static_expert_shapeis bound,MoEStaticSchedulerParams.__extract_mlir_values__emits no MLIR values.MoEStaticPersistentTileScheduler.__new_from_mlir_values__still slicesvalues[idx : idx + 3], which deserializes the wrong fields and advancesidxintooffs.🐛 Proposed fix
- new_params = new_from_mlir_values(self.params, values[idx : idx + 3]) - idx += 3 + params_len = len(extract_mlir_values(self.params)) + new_params = new_from_mlir_values( + self.params, values[idx : idx + params_len] + ) + idx += params_len🤖 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/sm120/swapab_cutedsl_megakernel/src/moe_sm120_mxfp8_swapab/moe_persistent_scheduler.py` around lines 1205 - 1212, Update MoEStaticPersistentTileScheduler.__new_from_mlir_values__ to derive the serialized parameter count from self.params instead of hardcoding three values. When self.params has static_expert_shape bound and emits no MLIR values, consume zero entries and leave idx positioned at the start of self.offs; preserve normal deserialization for dynamic parameters.flashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/src/moe_sm120_mxfp8_swapab/runner_fc12_common.py-1268-1281 (1)
1268-1281: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick winThe correctness path always runs the torch profiler.
Line 1269 inverts the condition: the profiler wraps the launch whenever
run_target_kernel_onlyis False, which is the default correctness path. Line 1278 then prints the fullkey_averagestable on every run. Profiling adds significant launch overhead and floods the harness output thatrun_mega_tests.shcaptures.Gate the profiler on an explicit switch.
MiscDescalready carriesverbose.🐛 Proposed fix
- if not self.misc.run_target_kernel_only: + if self.misc.verbose: with torch.profiler.profile(🤖 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/sm120/swapab_cutedsl_megakernel/src/moe_sm120_mxfp8_swapab/runner_fc12_common.py` around lines 1268 - 1281, Update the launch logic around compiled_kernel so torch.profiler.profile and its key_averages table run only when MiscDesc.verbose is enabled. Keep the existing unprofiled launch and synchronization behavior for the default correctness path and run_target_kernel_only mode.flashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/src/moe_sm120_mxfp8_swapab/token_comm.py-62-75 (1)
62-75: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winThe constructor silently overwrites a caller-supplied
combine_format.Lines 64-72 assign
kwargs["combine_format"]unconditionally. If a caller passescombine_formatthrough**kwargs, this discards it without any error. The baseTokenInPullTokenBackPushaccepts that keyword, so the call site looks valid and the substitution is invisible.Line 73 has the same problem for
token_back_by_dispatch: it derives the value fromfc2_output_dtypeand overwrites any explicit setting.Reject the conflicting keywords instead of dropping them.
🐛 Proposed fix
local_rank = kwargs.pop("local_rank") fc2_output_dtype = kwargs.pop("fc2_output_dtype", None) + for _reserved in ("combine_format", "token_back_by_dispatch"): + if _reserved in kwargs: + raise ValueError( + f"{_reserved} is derived from fc2_output_dtype on the SM120 " + f"SYSMEM path; do not pass it explicitly." + ) kwargs["combine_format"] = CombineFormat(🤖 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/sm120/swapab_cutedsl_megakernel/src/moe_sm120_mxfp8_swapab/token_comm.py` around lines 62 - 75, Update the constructor around the `combine_format` and `token_back_by_dispatch` assignments to detect whether either keyword was already supplied in `kwargs`; raise an error for caller-provided values instead of overwriting them, then derive and assign the defaults only when absent before calling `TokenInPullTokenBackPush.__init__`.flashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/src/moe_sm120_mxfp8_swapab/kernel_fc12.py-279-282 (1)
279-282: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winValidate the K tile against the hardcoded four inner MMA steps.
_validate_mma_tiler_and_cluster_shapeaccepts anyk % 32 == 0. Both compute loops then hardcode the inner K walk to four steps:for k_inner_mma in cutlass.range_constexpr(0, 4)at Line 1883 and Line 3880, against themake_swapab_m64n8k128_tiled_mmaatom. Four steps of the K32 atom consume exactly 128 K elements per staged tile.If a caller passes
mma_tiler_mnk[2] = 256(theImplDescdefault inrunner_fc12_common.pyis(128, 128, 256)), the TMA stages a K256 tile but the MMA loop only accumulates the first 128 K elements. The result is silently wrong, with no error._setup_attributesonly asserts divisibility by the instruction K, so it does not catch this either.fc1_tiles_per_fc2_k_tileat Line 3039 also assumes this relation.Reject any K other than 128, or drive the inner loop from
mma_tiler_mnk[2] // 32.🐛 Proposed guard
if k % 32 != 0: raise ValueError( f"SM120 MXFP8 K ({k}) must be a multiple of the m16n8k32 K atom." ) + if k != 128: + raise NotImplementedError( + f"SM120 MXFP8 swap-AB hardcodes 4 inner K32 MMA steps per staged " + f"tile, so mma_tiler K must be 128; got {k}." + )🤖 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/sm120/swapab_cutedsl_megakernel/src/moe_sm120_mxfp8_swapab/kernel_fc12.py` around lines 279 - 282, Update _validate_mma_tiler_and_cluster_shape so the K tile matches the four hardcoded m16n8k32 MMA steps used by both compute loops and fc1_tiles_per_fc2_k_tile: reject any mma_tiler_mnk[2] value other than 128, or replace those fixed inner-loop bounds with mma_tiler_mnk[2] // 32 and update dependent staging logic consistently.flashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/src/moe_nvfp4_swapab/megamoe_kernel.py-909-938 (1)
909-938: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
name()duplicates the base cache key by hand.This method restates every field of
Sm100SwapABSwigluFp4Fc12Kernel.name()(kernel_fc12.pylines 193-211) and appends the MegaMoE fields. The docstring already records that the shared part must be kept in sync manually.
name()is the compiled-kernel cache key. If a future change adds a codegen-affecting field to the basename()and this copy is not updated, two functionally different kernels collide on one cache entry and the wrong binary is reused.Have the base expose the shared segment once and have this method append to it.
♻️ Proposed structure
In
kernel_fc12.py:+ def _name_common_suffix(self) -> str: + m, n, k = self.mma_tiler_mnk + cm, cn = self.cluster_shape_mn + exp = "x".join(map(str, self.static_expert_shape)) if self.static_expert_shape else "dyn" + epiflag = "x".join(map(str, self.epi_flag_batch)) if self.epi_flag_batch else "none" + cta = "2_cta" if self.use_2cta_instrs else "1_cta" + fc2store = "fc2store_stg" if self.non_ubulk_fc2_store else "fc2store_ublk" + inkred = "inkernel_redg" if self.in_kernel_fc2_reduce else "no_inkernel_redg" + apply_topk = "apply_topk_fc1_pre_quant" if self.apply_topk_in_fc1 else "apply_topk_after_fc2" + return ( + f"_mmatiler_{m}x{n}x{k}_cluster_{cm}x{cn}_{cta}_sched_{self.load_balance_mode}" + f"_expert_shape_{exp}_grouphint_{self.group_hint}" + f"_padding_{self.token_padding_block}x{self.sf_padding_block}" + f"_{fc2store}_{inkred}_{apply_topk}" + f"_fc2out{self.fc2_output_dtype.__name__}_sfvec{self.sf_vec_size}" + f"_acc{self.acc_dtype.__name__}_clamp{self.gate_up_clamp}_epiflag{epiflag}" + ) + def name(self) -> str: - ... + return "moe_fc12_fuse_nvfp4" + self._name_common_suffix()Then in this file build
"megamoe_nvfp4" + self._name_common_suffix() + <megamoe fields>.🤖 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/sm120/swapab_cutedsl_megakernel/src/moe_nvfp4_swapab/megamoe_kernel.py` around lines 909 - 938, Refactor the cache-key construction so the base kernel exposes the shared name segment through a reusable method such as _name_common_suffix(), and update the MegaMoE name() method to build "megamoe_nvfp4" from that shared suffix before appending its MegaMoE-specific fields. Remove the duplicated base-field formatting from name() while preserving the existing key format and excluding local_rank as before.flashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/src/moe_nvfp4_swapab/epilogue_refactor.py-1414-1451 (1)
1414-1451: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winUse
self.warp_idxfor the per-warp TMEM slice index.
_subtile_local_tmem_tensorpasses a TMEM subtile view that is already offset by each warp’s_warp_lane_offset, but this local 32-row split removes that warp-only stride. Selecting slice0makes epi warps 1-3 load warp 0's accumulator rows for the first raw load, so the non-preload fc1 path produces wrong outputs for non-zero epilogue warps. Useself.warp_idxhere to keep the sliced view per-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/sm120/swapab_cutedsl_megakernel/src/moe_nvfp4_swapab/epilogue_refactor.py` around lines 1414 - 1451, Update the per-warp TMEM slice selection in the epilogue load path to use self.warp_idx instead of the hardcoded 0 when indexing tmem_subtile_tensor_per_warp. Preserve the existing local 32-row split and copy operations so each epilogue warp loads its own accumulator rows.flashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/src/moe_nvfp4_swapab/kernel_fc12.py-1338-1343 (1)
1338-1343: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winUse the scheduler’s 2-CTA position for
is_leader_cta.Under
is_swap_ab, launch grid shape becomes(cluster_m, 1)and the 2-CTA pairs run alongcta_id_in_cluster[0] == bidy % cluster_shape_mn[0];bidxis fixed at 0, so the current derivation makes both CTAs in each pair leaders. Source the V pair index fromblock_in_cluster_coord_vmnk[0]instead.🤖 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/sm120/swapab_cutedsl_megakernel/src/moe_nvfp4_swapab/kernel_fc12.py` around lines 1338 - 1343, Update is_leader_cta in the kernel’s CTA-index setup to derive the V pair index from block_in_cluster_coord_vmnk[0], using the scheduler’s 2-CTA position rather than bidx or mma_tile_coord_v. Preserve the existing leader comparison against the tiled-MMA thread-ID extent and leave cta_rank_in_cluster unchanged.flashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/src/moe_nvfp4_swapab/moe_utils.py-49-50 (1)
49-50: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winUse 128 bytes for
TensormapDescBytesand remove the duplicate assignment.This file defines the TMA descriptor slot stride and allocation multiplier, including
TensormapWorkspace.get_ptrandTensormapWorkspace.size_bytes. ACUtensorMapdescriptor is 128 bytes; using 64 bytes makes descriptor slots underallocate and can write into the next executor’s slot. KeepTensormapDescBytes = 128once.🤖 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/sm120/swapab_cutedsl_megakernel/src/moe_nvfp4_swapab/moe_utils.py` around lines 49 - 50, Update the module-level TensormapDescBytes definition to a single assignment of 128, removing the duplicate 64-byte assignment. Ensure TensormapWorkspace.get_ptr and TensormapWorkspace.size_bytes continue using this constant for descriptor stride and allocation.Source: Linters/SAST tools
flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/shim/hopper_fp8.py-457-479 (1)
457-479: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winPin the weight tensors used in the launch-key check.
MegaMoEHopperFp8Frontend.run()hits the cached path when the data_ptr tuple matches, and_CompiledMegastoreslaunch_kwargsandlaunch_outputbut not the FC weights/scales they were built from. If a caller drops a tensor and CUDA reuses the same address,_validate_inputsis skipped and_build_mega_runtime_kwargs()will run on the replaced source. Include the weight tensor identities/objects in the cache entry or enforce a caller contract that keeps these tensors alive for the session lifetime.🤖 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 457 - 479, The cache key built by _launch_cache_key must retain the FC weight and scale tensor objects, not only their data_ptr values, so cached launches cannot reuse an address for replacement tensors. Include the relevant weight tensors in the cached entry or otherwise pin them for the compiled session, and ensure _validate_inputs remains effective when the original tensors are no longer alive.Source: Learnings
flashinfer/moe_ep/backends/mega/kernel/sm90/fp8_fp8_bf16_pull_cutedsl/staging.py-67-73 (1)
67-73: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winRecord the staged token count before the zero-token early return.
The early return at Lines 68-69 skips
_note_staged_tokensat Line 122. Two consequences follow:
- On a fresh workspace,
compute(output=None)callsstaged_tokens(), getsNone, and raisesValueError.- After an earlier stage of N > 0 tokens, the attribute keeps the stale N. A later zero-token stage leaves that stale value, so
compute(output=None)runs over N rows of stale routing data.The pre-staged branch in
backend.pyalways calls_note_staged_tokens, so the two staging paths currently disagree.🔧 Proposed fix
num_tokens, hidden = hidden_states.shape if num_tokens == 0: + _note_staged_tokens(topk_idx_out, 0) returnAlso applies to: 116-122
🤖 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/fp8_fp8_bf16_pull_cutedsl/staging.py` around lines 67 - 73, Update the staging function around the num_tokens zero-token guard to call _note_staged_tokens with the current num_tokens before returning, ensuring zero-token stages record 0 instead of leaving staged_tokens unset or stale. Preserve the existing validation and non-empty staging flow, and keep the behavior consistent with the pre-staged backend path.flashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/src/src/bootstrap.py-444-461 (1)
444-461: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winThe CPU-staged
all_gathershim dropsasync_opsemantics.
_cpu_staged_all_gatheracceptsasync_opbut ignores it on the CUDA path and always returnsNone. A caller that passesasync_op=TruereceivesNoneinstead of a work handle, so a following.wait()raisesAttributeError. The shim replacesdist.all_gatherprocess-wide, so any other code in the same process is affected onceMEGA_SINGLE_GPU_GLOO=1is set.Reject the unsupported mode explicitly instead of returning a value that violates the API contract.
🔧 Proposed fix
def _cpu_staged_all_gather(tensor_list, tensor, group=None, async_op=False): if not tensor.is_cuda: return _orig_all_gather( tensor_list, tensor, group=group, async_op=async_op ) + if async_op: + raise NotImplementedError( + "MEGA_SINGLE_GPU_GLOO CPU-staged all_gather does not " + "support async_op=True for CUDA tensors." + ) cpu_out = [ torch.empty_like(t, device="cpu") for t in tensor_list ] _orig_all_gather(cpu_out, tensor.cpu(), group=group) - for dst, src in zip(tensor_list, cpu_out): + for dst, src in zip(tensor_list, cpu_out, strict=True): dst.copy_(src) return None🤖 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/sm120/swapab_cutedsl_megakernel/src/src/bootstrap.py` around lines 444 - 461, Update _cpu_staged_all_gather to explicitly reject async_op=True, raising an appropriate unsupported-operation error before entering the CUDA staging path; preserve the existing synchronous behavior and return value for async_op=False, including the non-CUDA delegation.Source: Linters/SAST tools
flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/runner_fc12_common.py-205-234 (1)
205-234: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winUse the kind-aware data dtype in the TMA alignment checks.
ProblemDesc.__post_init__supportsnvfp4,mxfp8_e4m3,mxfp8_e5m2,fp8_e4m3, andfp8_e5m2, but the four_check_tma_leading_dim_aligncalls useNvfp4DataDtype. For FP8 kinds,leading_dim_bytesmeasures FP8 data as 1 byte/element, whileNvfp4DataDtypemeasures it as 0.5 byte/element, so valid FP8 layouts can fail and the error names the wrong dtype. Select the kind data dtype once withkind_data_dtype(self.kind)and pass it to those checks.🤖 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/runner_fc12_common.py` around lines 205 - 234, Use kind_data_dtype(self.kind) in the surrounding initialization logic to select the data dtype once, then pass that value to the four TMA alignment checks for activation, fc1_weight, fc2_weight, and fc1_output in the runner flow. Replace their hardcoded Nvfp4DataDtype arguments while preserving the existing self.fc2_output_dtype argument for fc2_output.flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/moe_persistent_scheduler.py-1205-1212 (1)
1205-1212: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winKeep
paramsdeserialization aligned with the serialized value count.
MoEStaticSchedulerParams.__extract_mlir_values__emits zero values when only Python int fields are used inexpert_shape.MoEStaticPersistentTileScheduler.__new_from_mlir_values__still slicesvalues[idx : idx + 3]and advancesidx += 3, so every subsequent field is deserialized from the wrong offset. Use the emittedparamsvalue length, asMoEDynamicPersistentTileSchedulerdoes.🔧 Proposed fix
idx = 0 - new_params = new_from_mlir_values(self.params, values[idx : idx + 3]) - idx += 3 + params_len = len(extract_mlir_values(self.params)) + new_params = new_from_mlir_values(self.params, values[idx : idx + params_len]) + idx += params_len🤖 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/moe_persistent_scheduler.py` around lines 1205 - 1212, Update MoEStaticPersistentTileScheduler.__new_from_mlir_values__ to determine the params slice length from MoEStaticSchedulerParams.__extract_mlir_values__ (or the corresponding emitted params values) instead of assuming three values. Advance idx by that same computed length so new_offs and all subsequent fields remain aligned, including when expert_shape contains only Python ints.flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/token_comm.py-653-664 (1)
653-664: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winUse the fixed per-rank capacity when deriving the advertise buffer stride.
src_token_topk_idxis allocated as(num_experts_per_rank, world_size, max_tokens_per_rank * num_topk), butdispatch_warp_bodypassesinput_token_buffer.shape[0]intodispatch_prep; this createsMAX_SLOT_C = num_tokens * num_topk. Under a per-rank capacity contract, this must be the runtime per-rank token count, not the local buffer length. Add an assertion or pass an explicit constant for the stride used at line 678.🤖 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/token_comm.py` around lines 653 - 664, Update dispatch_warp_body and dispatch_prep so the advertise-buffer stride uses the fixed runtime per-rank token capacity required by src_token_topk_idx, rather than input_token_buffer.shape[0]. Ensure MAX_SLOT_C in the token_comm.py path is derived from that capacity and add an assertion validating the buffer shape when appropriate.flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/common/host_utils.py-170-173 (1)
170-173: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winDetect NaN values as mismatches.
At Line 173,
diff > thresholdis false whendiffis NaN. A NaN kernel result can therefore print a successful validation result against a finite reference value. Usetorch.isclose(..., equal_nan=False)and invert its result, or explicitly include non-finite values inmismatch_mask.Confidence: High.
Proposed fix
- diff = _torch.abs(gpu_data.float() - ref_data.float()) - threshold = atol + rtol * _torch.abs(ref_data.float()) - mismatch_mask = diff > threshold + mismatch_mask = ~_torch.isclose( + gpu_data.float(), ref_data.float(), atol=atol, rtol=rtol, equal_nan=False + )🤖 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/common/host_utils.py` around lines 170 - 173, Update the mismatch calculation in the validation logic around `mismatch_mask` to classify NaN results as mismatches, rather than relying solely on `diff > threshold`. Use `_torch.isclose` with `equal_nan=False` and invert it, or explicitly combine the existing comparison with a non-finite-value check while preserving the atol/rtol thresholds.flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/mega_runner.py-2644-2664 (1)
2644-2664: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winThe
NotImplementedErrorhandler reports success and hides real failures.
return_codestays0on every path, andos._exit(return_code)runs unconditionally. AnyNotImplementedErrorraised anywhere insidetester.run()is caught, printed on rank 0 only, and the process exits with status 0.The comment says the handler exists "until the MegaMoE kernel side is wired", but this PR wires the kernel. The handler now also swallows genuine failures.
validateraisesNotImplementedErrorat lines 1504-1508 for thetransformers+ in-kernel-reduce combination, and that failure would be reported as a pass. Non-rank-0 processes print nothing at all.🐛 Proposed fix
return_code = 0 try: tester.run() except NotImplementedError as exc: - # Expected until the MegaMoE kernel side is wired; the host - # orchestration above is the part being smoke-tested for now. - if rank == 0: - print(f"[mega_runner] kernel launch skipped: {exc}") + print(f"[rank {rank}] [mega_runner] unsupported configuration: {exc}") + return_code = 1🤖 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_runner.py` around lines 2644 - 2664, Update the tester.run() error handling and return_code flow so NotImplementedError is no longer silently treated as success now that the kernel is wired. Preserve intentional handling for the unsupported transformers plus in-kernel-reduce validation case by reporting the failure appropriately, propagate a nonzero exit status for genuine failures, and ensure all ranks participate consistently before os._exit(return_code).flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/mega_reference.py-2559-2597 (1)
2559-2597: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winInclude per-operand layout in the compiled-kernel cache key.
_to_cute_tensor()marks CuTe tensors with the input’sleading_dimbefore compiling, butself._compiledonly uses dtype. Identical dtypes from different calls, such as fc1 and fc2 inreference_expert_fc12(), share one compiled kernel; if their 3D operand layouts differ, the second call reuses a cache entry with the wrongmark_layout_dynamic(leading_dim=...)shape. Add per-operand shape/stride/leading-dimension data to the cache key.🤖 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` around lines 2559 - 2597, Update the _compiled cache key in the surrounding execution method to include each operand’s layout metadata, not just dtype: capture per-operand shape, stride, and leading-dimension information for a_cute, b_cute, sfa_cute, sfb_cute, and c_cute. Ensure calls with different 3D layouts compile and cache separate kernels while identical layouts continue reusing the correct compiled entry.flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/epilogue_fp8_swapab.py-885-931 (1)
885-931: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy liftDisable or rebuild the in-kernel FC2 REDG path.
fc2_in_kernel_topk_reducereaches this epilogue without a swap-AB guard, and the swap-AB hidden mapping strides bylane_group(hidden0/hidden1plus+8). The packing shuffles lanes+0/+4/+8/+12, so those lanes ownh,h+4,h+8,h+12rather than four adjacent 2-bytebf16cells for one 8-byte REDG. This write path can use misaligned REDG targets or write wrong columns.🤖 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_hopper_fp8/epilogue_fp8_swapab.py` around lines 885 - 931, Disable the in-kernel FC2 top-k reduction path in the swap-AB configuration, or rebuild its lane mapping so the shuffle lanes and hidden indices target four adjacent bf16 cells for each REDG segment. Update the branch guarded by _fc2_in_kernel_topk_reduce and preserve aligned destinations and correct columns before calling red_add_relaxed_sys_v2_bf16x2.flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/kernel_fp8_glu_fc12.py-2265-2284 (1)
2265-2284: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy liftUse the full stage transaction size for
tx_count.
tx_countis the expected transaction bytes for one stage barrier. The shared AB pipeline has two TMA producers, andnum_tma_load_bytesalready covers the A bytes, B bytes, and activation-scale bytes for one stage. Passingnum_tma_load_bytes // 2makes the full mbarrier complete before the TMA-B load finishes, so WGMMA consumers can read uninitialized SMEM. Settx_count=self.num_tma_load_bytes.🤖 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_hopper_fp8/kernel_fp8_glu_fc12.py` around lines 2265 - 2284, Update the AB pipeline construction in the kernel initialization to pass the full stage transaction size as tx_count. Specifically, change the PipelineTmaAsync.create call to use self.num_tma_load_bytes rather than halving it, while leaving the producer and consumer group configuration unchanged.flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/bootstrap.py-27-31 (1)
27-31: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winDefer
nvshmem.coreimport until device bootstrap.
flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/__init__.pyimportsget_symm_buffer_for_hopper_fp8_mega_moe, which transitively importsshim/comm.pyand thensrc/bootstrap.py. That makesflashinfer.moe_ep.kernel_src.sm90.pull_style_cutedsl_megakernelfail at import time in environments withoutnvshmem4py. Movenvshmem.coreinto the function path where sym-buffer allocation/bootstrap runs, or makesrc/bootstrap.pyraise/defer a usable conditional error only whendist.is_initialized().🤖 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/bootstrap.py` around lines 27 - 31, Remove the module-level nvshmem.core import from src/bootstrap.py and import it only inside the device bootstrap or symmetric-buffer allocation function that requires it. Keep package imports usable without nvshmem4py, while preserving a clear conditional error when that runtime path is invoked without the dependency.flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/custom_ext.py-565-566 (1)
565-566: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
_cluster_mis global class state and can be overwritten between kernels.Line 566 mutates the class attribute
GluMxFp8WorkTileInfo._cluster_m. EveryGluMxFp8WorkTileInfo.from_rmemcall reads that class attribute at Line 253 to derivefc1_counter_index. If twoGluMxFp8Fc12SchedExtensioninstances with differentcluster_mexist in one process (for example during a tuning sweep or a mixed SM90/SM100 run), the last constructed extension wins, and the earlier kernel derives a wrong FC1 counter slot.__new_from_mlir_values__at Lines 568-579 restoresresult.cluster_mbut never restores the class attribute, so the hazard also survives MLIR round-trips.Pass
cluster_mexplicitly intofrom_rmeminstead, or make the work-tile class per-extension.🐛 Sketch of an explicit-parameter fix
`@classmethod` - def from_rmem(cls, rmem: cute.Tensor) -> "GluMxFp8WorkTileInfo": + def from_rmem(cls, rmem: cute.Tensor, cluster_m: int = 1) -> "GluMxFp8WorkTileInfo": return cls( ... - fc1_counter_index=rmem[1] // cutlass.Int32(cls._cluster_m), + fc1_counter_index=rmem[1] // cutlass.Int32(cluster_m), )Each
from_rmemcall site then supplies the owning extension'sself.cluster_m.🤖 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/custom_ext.py` around lines 565 - 566, Remove the shared `GluMxFp8WorkTileInfo._cluster_m` mutation and make `cluster_m` an explicit input to `GluMxFp8WorkTileInfo.from_rmem`. Update every `from_rmem` call in `GluMxFp8Fc12SchedExtension` to pass the owning extension’s `self.cluster_m`, while preserving `result.cluster_m` handling in `__new_from_mlir_values__`.flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/mega_runner.py-52-56 (1)
52-56: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winAdd or move the missing
tester.host_utilsimports.
flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/mega_runner.pyimportsreduce_add_deterministic_check_dim_size_limitandreduce_add_ordering_matchfromtester.host_utils, but there is notesterpackage in the vendored SM90 file set. Add those helpers alongside the existing shared test utilities or point the import at the package that defines them. Without this,mega_runnerfails during module load andvalidate()cannot run.🤖 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_hopper_fp8/mega_runner.py` around lines 52 - 56, Update the imports used by mega_runner so reduce_add_deterministic_check_dim_size_limit and reduce_add_ordering_match resolve from an available shared utility package rather than the nonexistent tester package. Preserve validate() behavior and ensure the module loads successfully in the vendored SM90 file set.
| if sf_vec_size in self.VALID_AB_DTYPE_SF_SIZE: | ||
| valid_ab = self.VALID_AB_DTYPE_SF_SIZE[sf_vec_size] | ||
| if ab_dtype not in valid_ab: | ||
| raise ValueError( | ||
| f"ab_dtype={ab_dtype.__name__} is not valid for " | ||
| f"sf_vec_size={sf_vec_size}. " | ||
| f"Expected one of: {[t.__name__ for t in valid_ab_tuple]}." | ||
| ) | ||
| else: | ||
| raise NotImplementedError( | ||
| f"sf_vec_size must be {Mxfp8BlockSize} (MXFP8)" | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
valid_ab_tuple is undefined; the validation path raises NameError.
The local variable is named valid_ab at line 102. Line 107 references valid_ab_tuple. When a caller passes an unsupported ab_dtype, Python raises NameError while building the message instead of the intended ValueError. The diagnostic is lost.
🐛 Proposed fix
if ab_dtype not in valid_ab:
raise ValueError(
f"ab_dtype={ab_dtype.__name__} is not valid for "
f"sf_vec_size={sf_vec_size}. "
- f"Expected one of: {[t.__name__ for t in valid_ab_tuple]}."
+ f"Expected one of: {[t.__name__ for t in valid_ab]}."
)📝 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 sf_vec_size in self.VALID_AB_DTYPE_SF_SIZE: | |
| valid_ab = self.VALID_AB_DTYPE_SF_SIZE[sf_vec_size] | |
| if ab_dtype not in valid_ab: | |
| raise ValueError( | |
| f"ab_dtype={ab_dtype.__name__} is not valid for " | |
| f"sf_vec_size={sf_vec_size}. " | |
| f"Expected one of: {[t.__name__ for t in valid_ab_tuple]}." | |
| ) | |
| else: | |
| raise NotImplementedError( | |
| f"sf_vec_size must be {Mxfp8BlockSize} (MXFP8)" | |
| ) | |
| if sf_vec_size in self.VALID_AB_DTYPE_SF_SIZE: | |
| valid_ab = self.VALID_AB_DTYPE_SF_SIZE[sf_vec_size] | |
| if ab_dtype not in valid_ab: | |
| raise ValueError( | |
| f"ab_dtype={ab_dtype.__name__} is not valid for " | |
| f"sf_vec_size={sf_vec_size}. " | |
| f"Expected one of: {[t.__name__ for t in valid_ab]}." | |
| ) | |
| else: | |
| raise NotImplementedError( | |
| f"sf_vec_size must be {Mxfp8BlockSize} (MXFP8)" | |
| ) |
🧰 Tools
🪛 Ruff (0.16.1)
[error] 107-107: Undefined name valid_ab_tuple
(F821)
🤖 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/sm120/swapab_cutedsl_megakernel/src/moe_mxfp8_glu/kernel_mxfp8_glu_fc12.py`
around lines 101 - 112, Update the validation error message in the sf_vec_size
branch to reference the existing valid_ab collection instead of the undefined
valid_ab_tuple, while preserving the intended ValueError for unsupported
ab_dtype values.
Source: Linters/SAST tools
| @cute.jit | ||
| def _amax_lane(self, v: cutlass.Float32) -> cutlass.Float32: | ||
| if cutlass.const_expr(self.sf_vec_size == 32): | ||
| return cute.arch.warp_redux_sync(v, "fmax", abs=True) | ||
| first_half = (self.lane_idx % cutlass.Int32(32)) < cutlass.Int32(16) | ||
| vsel = cutlass.Float32(0.0) | ||
| if first_half: | ||
| vsel = v | ||
| amax = cute.arch.warp_redux_sync(vsel, "fmax", abs=True) | ||
| if not first_half: | ||
| amax = cute.arch.warp_redux_sync(v, "fmax", abs=True) | ||
| return amax |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Inspect cute.arch.warp_redux_sync signature and its default member mask.
set -euo pipefail
CUTLASS_DIR="$(python -c 'import cutlass, pathlib; print(pathlib.Path(cutlass.__file__).parent)' 2>/dev/null || true)"
if [ -z "${CUTLASS_DIR:-}" ]; then echo "cutlass not importable in sandbox"; exit 0; fi
rg -nP -C10 'def\s+warp_redux_sync' "$CUTLASS_DIR" --glob '*.py'
rg -nP -C3 'redux\.sync' "$CUTLASS_DIR" --glob '*.py' | head -40Repository: flashinfer-ai/flashinfer
Length of output: 196
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate target file =="
fd -a 'epilogue_refactor.py' . || true
echo "== relevant source context =="
if [ -f flashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/src/moe_nvfp4_swapab/epilogue_refactor.py ]; then
sed -n '440,490p' flashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/src/moe_nvfp4_swapab/epilogue_refactor.py | cat -n
fi
echo "== cute.arch.warp_redux_sync call sites =="
rg -n 'warp_redux_sync|redux\.sync' . --glob '*.py' | head -200
echo "== imports and sf_vec_size context =="
sed -n '1,80p' flashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/src/moe_nvfp4_swapab/epilogue_refactor.py | cat -n
rg -n 'sf_vec_size|lane_idx' flashinfer/moe_ep/kernel_src/sm120/swapab_cutedsl_megakernel/src/moe_nvfp4_swapab/epilogue_refactor.py | head -80Repository: flashinfer-ai/flashinfer
Length of output: 16871
🌐 Web query:
CUTE warp_redux_sync default member mask redux.sync full member mask
💡 Result:
In the context of the NVIDIA CuTe (CUTLASS Template) library, the warp_redux_sync function and the underlying PTX redux.sync instruction utilize a membership mask to define which threads in a warp participate in a collective reduction operation [1][2]. Warp Membership Mask The membership mask is a 32-bit integer where each bit corresponds to a specific thread (lane) in a warp [3]. When executing warp-level primitives, this mask determines which threads are synchronized and participate in the operation [3]. Default and Full Masks - Full Member Mask: In most NVIDIA CUDA and CuTe contexts, the full member mask is defined as 0xFFFFFFFF, which represents all 32 threads in a warp [1][3][4]. - Default Mask: The CuTe warp_redux_sync API explicitly uses this full mask (0xFFFFFFFF) as the default value for its mask_and_clamp parameter [1]. Usage Notes - PTX Specification: The underlying PTX instruction redux.sync.<op>.<type> dst, src, membermask expects the src value followed by the membermask [2][5]. Developers should be cautious of alignment issues; historical discrepancies between intrinsic definitions and PTX documentation have occasionally led to confusion regarding argument order [6][5]. - Best Practices: While 0xFFFFFFFF is the standard default, you should analyze program logic to ensure only threads intended to participate are included in the mask [3]. Using an incorrect mask—especially failing to include all threads that reach the primitive—can lead to program hangs or undefined behavior [3]. It is important to note that redux.sync operations are architecture-specific (typically requiring SM 8.0 or higher) and function by synchronizing all participating threads before performing the reduction [2][7][8].
Citations:
- 1: https://docs.nvidia.com/cutlass/4.5.3/media/docs/pythonDSL/cute_dsl_api/cute_arch.html
- 2: https://reviews.llvm.org/D100124
- 3: https://developer.nvidia.com/blog/using-cuda-warp-level-primitives/
- 4: https://github.com/tile-ai/tilelang/blob/a417b38a/tilelang/contrib/cutedsl/warp.py
- 5: https://lists.llvm.org/pipermail/llvm-commits/Week-of-Mon-20250324/1595275.html
- 6: [cuclang 20.1.0] Using __reduce_max_sync in a cuda kernel fails with an illegal instruction was encountered llvm/llvm-project#131415
- 7: https://lists.llvm.org/pipermail/cfe-commits/Week-of-Mon-20210517/370284.html
- 8: https://github.com/NVIDIA/TensorRT-LLM/blob/7a8bd87f/tensorrt_llm/_torch/cute_dsl_kernels/argmax.py
Keep _amax_lane reductions on uniform warp control flow.
cute.arch.warp_redux_sync has a full-warp member mask by default, but the second reduction is reached only by the 16 lanes in first_half == False. That violates redux.sync’s full member-mask contract and can leave the warp hanging.
Compute both half-warp amax candidates with masked values, then select the lane’s result afterward.
🤖 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/sm120/swapab_cutedsl_megakernel/src/moe_nvfp4_swapab/epilogue_refactor.py`
around lines 464 - 475, Update _amax_lane to keep every warp_redux_sync call on
uniform warp control flow: compute both half-warp candidates using masked
values, perform the reductions unconditionally for all lanes, then select the
appropriate result based on first_half. Preserve the existing full-warp path for
sf_vec_size == 32.
| in_aligned_iter = cute.make_ptr( | ||
| in_tile.element_type, | ||
| in_tile.iterator.toint(), | ||
| AddressSpace.gmem, | ||
| assumed_align=16, | ||
| ) | ||
| in_tile = cute.make_tensor(in_aligned_iter, in_tile.layout) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
The kernel asserts 16-byte alignment that the validator never checks on the base pointer.
Lines 66-71 and 115-120 rebuild the tile iterators with assumed_align=16 and then issue a 128-bit CopyUniversalOp. _validate_tensors checks only the strides at lines 167-172. It never checks combine_output.data_ptr() or reduced_output.data_ptr().
_infer_assumed_align at lines 187-192 confirms that the base pointer can be less than 16-byte aligned. A non-owning slice of a larger buffer reaches this path with an 8-byte aligned base. Every stride check still passes, and the 128-bit copy then reads or writes a misaligned address, which faults at runtime.
Add a base-pointer check to _validate_tensors.
🛡️ Proposed fix
if reduced_output.stride(0) % BF16_HIDDEN_PER_THREAD != 0:
raise ValueError("reduced_output rows must preserve 16-byte alignment.")
+ for _name, _t in (("combine_output", combine_output), ("reduced_output", reduced_output)):
+ if int(_t.data_ptr()) % 16 != 0:
+ raise ValueError(
+ f"{_name} base pointer must be 16-byte aligned for the "
+ f"128-bit vector copy; got {int(_t.data_ptr()) % 16} byte offset."
+ )Also applies to: 115-121, 167-172
🤖 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/sm120/swapab_cutedsl_megakernel/src/moe_sm120_mxfp8_swapab/topk_reduce.py`
around lines 66 - 72, Update _validate_tensors to validate that
combine_output.data_ptr() and reduced_output.data_ptr() are 16-byte aligned, in
addition to the existing stride checks. Reject misaligned base pointers before
the tile iterators use assumed_align=16 for the 128-bit copies.
mhoqueanik
changed the title
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/VENDOR.md (1)
9-12: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick winRecord the immutable upstream baseline.
Line 10 through Line 12 leave the repository URL, base commit, and sync date unresolved. The partial update at Line 29 cannot be audited or reproduced without this baseline. Record the exact upstream URL, full base SHA, and sync date before merge.
🤖 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/cutedsl_megamoe/VENDOR.md` around lines 9 - 12, Update VENDOR.md to replace the TODO placeholders with the exact immutable upstream repository URL, full source baseline commit SHA, and date when the current drop was synced, preserving the existing repository, Vendored commit, and Last synced fields so the partial update can be audited and reproduced.
🤖 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.
Outside diff comments:
In `@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/VENDOR.md`:
- Around line 9-12: Update VENDOR.md to replace the TODO placeholders with the
exact immutable upstream repository URL, full source baseline commit SHA, and
date when the current drop was synced, preserving the existing repository,
Vendored commit, and Last synced fields so the partial update can be audited and
reproduced.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 8fe69aaf-04a6-4839-b1b7-87748290e6b6
📒 Files selected for processing (6)
.pre-commit-config.yamldocs/design_docs/moe_ep_runbook.mdflashinfer/moe_ep/kernel_src/README.mdflashinfer/moe_ep/kernel_src/cutedsl_megamoe/VENDOR.mdflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/common/host_utils.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/src/inputs_process.py
🚧 Files skipped from review as they are similar to previous changes (2)
- .pre-commit-config.yaml
- docs/design_docs/moe_ep_runbook.md
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/VENDOR.md (1)
9-12: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winPin the current upstream snapshot.
Repo,Vendored commit, andLast syncedremain TODO. Lines 19-20 require a verbatim drop and a clean upstream diff, but these placeholders prevent reproducible auditing ofsrc/. Record the canonical repository URL, exact upstream SHA, and sync date before merge.🤖 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/cutedsl_megamoe/VENDOR.md` around lines 9 - 12, Replace the TODO placeholders in VENDOR.md with the canonical upstream repository URL, the exact commit SHA corresponding to the vendored src/ snapshot, and the date of that synchronization. Keep the existing attribution and verbatim-drop requirements unchanged.
🧹 Nitpick comments (1)
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/VENDOR.md (1)
29-47: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winMake the standalone
inputs_processharness compatible with the mixed snapshot.Production fused staging imports
DataPreprocessand does not load the mismatchedhost_utilssymbol. The test suite references the harness but does not execute it. Vendor the matchinghost_utilsAPI, or add a CI check that detects the missingmxfp8_quantize_per_block_32_rowsymbol when the harness runs.🤖 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/cutedsl_megamoe/VENDOR.md` around lines 29 - 47, Update the standalone inputs_process harness to work with the mixed snapshot by providing the expected mxfp8_quantize_per_block_32_row API in the recorded-drop host_utils, or by adding a CI check that explicitly detects and reports its absence when the harness executes. Keep production DataPreprocess imports and unrelated nvfp4, shim, and kernel paths unchanged.
🤖 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.
Outside diff comments:
In `@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/VENDOR.md`:
- Around line 9-12: Replace the TODO placeholders in VENDOR.md with the
canonical upstream repository URL, the exact commit SHA corresponding to the
vendored src/ snapshot, and the date of that synchronization. Keep the existing
attribution and verbatim-drop requirements unchanged.
---
Nitpick comments:
In `@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/VENDOR.md`:
- Around line 29-47: Update the standalone inputs_process harness to work with
the mixed snapshot by providing the expected mxfp8_quantize_per_block_32_row API
in the recorded-drop host_utils, or by adding a CI check that explicitly detects
and reports its absence when the harness executes. Keep production
DataPreprocess imports and unrelated nvfp4, shim, and kernel paths unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 4d7a0758-5827-4d28-b0a8-5adddb3bd8c3
📒 Files selected for processing (2)
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/VENDOR.mdflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/common/host_utils.py
🚧 Files skipped from review as they are similar to previous changes (1)
- flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/common/host_utils.py
56209bc to
34444e8
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
[SUCCESS] Pipeline #65772681: 16/16 executed test jobs passed |
…on main) (#4903) ## 📌 Description `tests/moe/test_unified_moe_activation_matrix.py::test_documented_activation_matrix_matches_runner_registry` is failing on `main`, which blocks CI for **every open PR**: ``` AssertionError: docs/design_docs/flashinfer_moe_api.md is stale; run: python scripts/generate_moe_activation_matrix.py --write ``` This is a **semantic merge conflict**, not a defect in any single PR. #4805 added the generator, its check test, and a matrix block rendered from `_BACKEND_RUNNERS` as it stood on that PR's base. Two changes landed on `main` in between, and neither could have known to re-render the block: | Change | Effect on the matrix | |---|---| | #4793 (`f7d4b167`) | renamed `CuteDslRunner.backend_key` `cute_dsl_nvfp4` → `cute_dsl` and added `QuantVariant.MXFP4` to its supported variants | | #4646 (`0cbace05`) | registered `CuTileBf16Runner` / `CuTileNvfp4Runner` in `_BACKEND_RUNNERS` | Each PR was green on its own base; the merged tree is what is stale. Because the check compares the committed block against the live registry, it has been red for everyone since #4805 merged. ## 🔍 Change Only the generated block changes — this commit is the mechanical output of the documented regeneration command: ``` python scripts/generate_moe_activation_matrix.py --write ``` - adds `cutile_bf16` (`BF16`) and `cutile_nvfp4` (`NVFP4`), both `SwiGLU`, `ReLU2` - replaces the two `cute_dsl_nvfp4` rows with three `cute_dsl` rows (`MXFP4`, `NVFP4`, `W4A16`) No source, test, or prose changes. ## 🧪 Testing The authoritative check is `test_documented_activation_matrix_matches_runner_registry`, which runs in this PR's own CI. ## 🔗 Related Surfaced while triaging CI on #4387, whose H100 job ran the full suite with 221,273 passing and this as the sole failure. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Documentation** * Updated the MoE activation matrix table to document MXFP4, NVFP4, and W4A16 support across additional activation functions. * Added documented BF16 and NVFP4 configuration entries for CuTile implementations. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Verbatim snapshot of the four kernel packages (common/, src/, moe_sm120_mxfp8_swapab/, moe_mxfp8_glu/) from the sm120_swapab_wt fork worktree at d19d30a (branch run/sm120-mxfp8-perf), incl. two uncommitted worktree edits recorded in VENDOR.md. No shim yet — src/ only. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Adapter layer over the verbatim drop, mirroring the sm90 fork-tree conventions: _paths sibling-tree guard, per-tree comm.py (plus zero_local_counter_regions -- this drop's kernel does not tail-clean its local counters), all-lazy kernel_helpers (constants pull cutlass at import), and the Sm120 MXFP8 frontend with combine_output + topk-reduce second stage, per-expert all-ones epilogue args, K-major weight views, and mirrored-ABI-constant guard. Also vendors moe_nvfp4_swapab (import dependency of the mxfp8 torch reference). AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Registers the SM120 swap-AB MXFP8 kernel behind the standard MegaKernelBackend contract: Sm120_Mxfp8_Mxfp8_Bf16_Cutedsl_MegaMoeConfig (native token_back_mode enum, knob dict only -- no autotune yet), K-major + interleave-8 weight preprocessing, torch-composed staging (no fused DataPreprocess shim in this tree), validate_mega_arch_sm120 (sm_120/sm_121 exact family), and the sm120 runtime-requirements alias. Exported from flashinfer.moe_ep; runbook tree list updated. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Mirrors the sm100 mxfp8 multirank file for the sm120 swap-AB backend: layer-vs-direct-shim parity (staged / prestaged / ikr / large-token dispatch token-back), the all-gather torch-oracle anchor (the sm120 reference pins gate_up_interleave=8 + apply_topk_in_fc1=True itself), sm120-local _plain_mxfp8_from_bf16, and registry/preprocess checks. The term-magnitude band helper moves to an arch-free shared module (tests/moe_ep/mega_oracle_compare.py) so both suites import one definition without dragging in each other's importorskip. Adds the arch_sm120 marker (exact sm_12x family) and tightens arch_blackwell to the sm_10x family so sm100-tree tests stop collecting on sm_11x/12x hosts where their kernels cannot compile. run_tests.sh gains mega_sm120 (own torchrun process; trees are process-exclusive). AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Every sm_120/121 node on the cluster has one GPU, and the SM120 kernel drop's own bootstrap runs N ranks per GPU (GB10 / DGX Spark). Mirror that in the FI runtime: fold LOCAL_RANK onto the physical GPUs (identity when there are enough), and under MEGA_SINGLE_GPU_GLOO=1 (the drop's env) init the process group as gloo with a CPU-side NVSHMEM UID broadcast (NCCL cannot host two ranks on one device). tests/conftest.py: an explicit torchrun WORLD_SIZE >= N overrides the physical-GPU-count skip for gpu_N markers. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The drop's own mega_runner crashes with cudaErrorIllegalAddress under --in_kernel_fc2_reduce (verified on RTX PRO 6000, 4 ranks / 1 GPU, DSL 4.6.1) and the flag is absent from the kernel team's test scripts. Reject it in the FI backend, skip the ikr tests with the documented reason, and record the gap in VENDOR.md; the shim keeps the plumbing for a fixed drop. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
cluster_m > 1 fails cute.compile ('expects num_multicast to be 1 for
non multicast G2S copies') -- reproduced with the drop's own
mega_runner; its test scripts always use 1,1,1. Shim config rejects
anything else; large-token test profile drops the cluster knob.
AI-assisted.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Under the rank-sharing gloo process group, all_gather of CUDA tensors silently corrupts (gloo has no CUDA all_gather; the kernel drop's own bootstrap monkey-patches identical CPU staging) -- the oracle was comparing the kernel against a reference built from garbage operands. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Two more upstream gaps pinned by A/B against the drop's own runner: (1) gate_up_clamp is dead plumbing (kernel_fc12 stores it, never reads it; output bit-identical with/without clamp -- invisible to the drop's ±0.5-sparse test data). Backend rejects a set clamp; tests go clampless. (2) world_size=1 (MEGA_NO_DIST) numerics are silently wrong for mma_tiler N=128 (5-20% cells, reproduced upstream at their own standard geometry); the same tile is bit-exact at world_size=4. Both recorded in VENDOR.md. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
pre-commit run -a caught that the sm120 vendored drop was never added to the lint excludes — mypy/ruff/whitespace hooks were rewriting ~50 files under kernel_src/sm120/swapab_cutedsl_megakernel/src/ (verbatim tree, must stay byte-identical to upstream per kernel_src/README.md). Added the drop to the three exclude points (.pre-commit-config.yaml global exclude, pyproject [tool.mypy] exclude, [tool.ruff] extend-exclude) and reverted the hook damage. Legit findings in OUR files, fixed: - shim/mxfp8.py: the symm-buffer factory annotated token_back_mode as str while MegaMoESm120Mxfp8Config wants the Literal — annotated with the same Literal (mypy arg-type). - ruff-format pass over shim/, sm120+sm100+sm90 backend wrappers, and the moe_ep tests (line wraps from the longer taxonomy names; same set the base branch reformatted in e9f791a). AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…e check The error path referenced an undefined valid_ab_tuple, masking the intended ValueError with a NameError; re-aligns with the reference copy in kernel_src/cutedsl_megamoe (which uses valid_ab). Addresses CodeRabbit review on PR flashinfer-ai#4387. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…28 numerics) Dense-data microbenchmarks (moe_ep_benchmark, deepseek_v3 geometry) show the sm120 drop's ws1 N=128 numerics gap extends to world_size=2: rel-L2 vs the bf16 dense reference degrades from the ~6.35% MXFP8 band to 10-28% once tokens fill past an N=64 tile, with run-to-run magnitude variation (race-like). A/B on RTX 6000D: ws2 tokens/rank=64 N=64 -> 6.365%, N=128 -> 10.266%. ws4 N=128 stays in band (6.32-6.34%). Backend now defaults mma_tiler_mnk=(64, 64, 128) when ep_world_size <= 2; an explicit mma_tiler_mnk knob still overrides. VENDOR.md updated. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Dense-data validation shows the N=128 corruption also hits world_size=4 (8-25% rel-L2 across 8..4096 tokens/rank vs the ~6.35% MXFP8 band; ws4 tokens/rank=64 with N=64 returns 6.356%). The drop's ws4 'bit-exact' check used 1%-sparse test data, which cannot see it. Pin mma_tiler_mnk=(64,64,128) unconditionally (explicit knob still overrides); ~23% large-batch throughput cost vs the broken N=128. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Upstream flashinfer-ai#4348 asserts _resolve_local_device returns the raw LOCAL_RANK; the unconditional local_rank % device_count fold broke that test on 1-GPU CI (5 % 1 == 0). Fold only under the sm_12x rank-sharing flow (MEGA_SINGLE_GPU_GLOO=1), where it is required, and add a unit test for the gated fold. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The test-sharding runner collects the whole primary test scope in one pytest process. tests/moe_ep/test_moe_ep_sm120_mxfp8_cutedsl_mega_multirank.py imports the vendored SM120 swapab_cutedsl_megakernel shim at module scope, whose _paths guard raises RuntimeError (not ImportError, so pytest.importorskip does not catch it) when the SM100 cutedsl_megamoe tree already owns the top-level `common` module in that process. That aborts collection for the entire scope, which is why unit_test_b300 and JIT Unittest (H100) fail with "kernel module 'common' is already imported from .../cutedsl_megamoe/src/common" before a single test runs. SM90 already had a collection isolation partition for exactly this reason; the three vendored trees are mutually exclusive with one another, so give SM120 its own partition rather than folding it into the SM90 group. Execution is already per-source-file subprocesses, so only collection needed the split. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
97e4720 to
3a9e911
Compare
3 participants
feat(moe_ep): SM120 MXFP8 swap-AB CuTeDSL MegaMoE kernel
Summary
Adds an SM120 (sm_120/sm_121) MXFP8 MegaMoE backend to
flashinfer.moe_ep— the swap-AB CuTeDSL megakernel (fused dispatch + grouped GEMM + combine) for the RTX PRO 6000 / GB10 (DGX Spark) class of Blackwell parts. Registered assm120_mxfp8_mxfp8_bf16_cutedslbehind the standardMegaKernelBackendcontract, on the same taxonomy layout as the SM100/SM90 backends.What's included
kernel_src/sm120/— verbatim snapshot of the four kernel packages (common/,src/,moe_sm120_mxfp8_swapab/,moe_mxfp8_glu/) from thesm120_swapab_wtfork atd19d30a(branchrun/sm120-mxfp8-perf), including two uncommitted worktree edits recorded inVENDOR.md.src/stays byte-for-byte upstream; all adaptation lives in the shim._pathssibling-tree guard, per-treecomm.py(pluszero_local_counter_regions; this drop's kernel does not tail-clean its local counters), all-lazykernel_helpers, and the SM120 MXFP8 frontend withcombine_output+ topk-reduce second stage, per-expert epilogue args, K-major weight views, and a mirrored-ABI-constant guard.backends/mega/kernel/sm120/mxfp8_mxfp8_bf16_cutedsl/withSm120_Mxfp8_Mxfp8_Bf16_Cutedsl_MegaMoeConfig(nativetoken_back_modeenum), K-major + interleave-8 weight preprocessing, torch-composed staging (no fused DataPreprocess in this tree),validate_mega_arch_sm120(exact sm_120/sm_121 family), and the sm120 runtime-requirements alias. Exported fromflashinfer.moe_ep.LOCAL_RANKonto the physical GPUs (identity when there are enough), and underMEGA_SINGLE_GPU_GLOO=1inits the process group as gloo with a CPU-side NVSHMEM UID broadcast (NCCL cannot host two ranks on one device).arch_sm120pytest marker (witharch_blackwelltightened to the sm_10x family so sm100 suites stop collecting on sm_11x/12x hosts), and amega_sm120entry inrun_tests.sh.Upstream gaps found and guarded
Each was pinned by A/B against the drop's own runner and is recorded in
VENDOR.mdas a pending-upstream item; the FI backend rejects the broken configuration rather than silently running it:in_kernel_fc2_reducecrashes withcudaErrorIllegalAddress(verified RTX PRO 6000, 4 ranks / 1 GPU, DSL 4.6.1); absent from the kernel team's own test scriptscluster_m > 1failscute.compile("expects num_multicast to be 1 for non multicast G2S copies"); upstream scripts always use (1,1,1)gate_up_clampis dead plumbing —kernel_fc12stores it but never reads it; output bit-identical with/without (invisible to the drop's ±0.5-sparse test data)world_size=1(MEGA_NO_DIST) numerics silently wrong formma_tilerN=128 (5-20% of cells; reproduced upstream at their own standard geometry); same tile bit-exact at world_size=4Two harness-level fixes that the correctness claims depend on: the oracle's
all_gatheris CPU-staged under the rank-sharing gloo group (gloo has no CUDA all_gather — the reference was previously built from corrupted operands), and the shim's_mainsmoke device folds onto physical GPUs.Usage
On a single-GPU sm_12x box, run multirank via the drop's env:
MEGA_SINGLE_GPU_GLOO=1 torchrun --nproc_per_node=4 ....Constraints
validate_mega_arch_sm120).in_kernel_fc2_reduce,cluster != (1,1,1), andgate_up_clampare rejected until the upstream drop fixes them (see table above).Testing
gate_up_interleave=8+apply_topk_in_fc1=Trueitself).run_tests.sh mega_sm120(own torchrun process; kernel trees are process-exclusive).Performance
Ongoing. Kernel-level tuning (tiler/knob sweeps) and the standalone microbenchmark/e2e runs are in progress; numbers will be posted to this PR (or a follow-up) once the tuning sweep completes. Nothing in this PR should be read as a perf claim yet.
Summary by CodeRabbit
New Features
Improvements
Documentation