feat(moe_ep): SM100 BF16 CuTeDSL MegaMoE kernel - #4386
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:
📝 WalkthroughWalkthroughThis PR renames MegaMoE backends into architecture- and dtype-qualified identifiers, for example ChangesMegaMoE Backend Restructuring
Estimated code review effort: 5 (Critical) | ~180 minutes Possibly related issues
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1⚔️ Resolve merge conflicts 💡
🧪 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: 3
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 (2)
tests/moe_ep/test_deep_gemm_mega_kernel_vs_reference.py (1)
226-242: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winSkip instead of error when the torchrun rendezvous environment is absent.
The plain-pytest TCP rendezvous fallback was removed. Under plain
pytest,WORLD_SIZEis unset, so line 226 defaults it to"1"and the guard on lines 227-228 does not skip. Execution then reachesdist.init_process_group(backend="nccl")on line 242 with theenv://rendezvous and noMASTER_ADDRorMASTER_PORT. PyTorch raisesValueErrorfor the missing environment variables. The test errors instead of skipping.The docstring on lines 20-24 documents the torchrun requirement, but a collection under plain
pytest tests/moe_epstill fails.Add a skip when the rendezvous environment is missing.
🛡️ Proposed fix to skip without a rendezvous environment
if not dist.is_initialized(): + if "MASTER_ADDR" not in os.environ or "MASTER_PORT" not in os.environ: + pytest.skip( + "needs a torchrun rendezvous; run with " + "torchrun --standalone --nproc_per_node=1 -m pytest" + ) dist.init_process_group(backend="nccl", timeout=_PG_TIMEOUT)Run the following script to check how the runner launches this file:
#!/bin/bash # Check whether run_tests.sh always launches this file under torchrun. fd -t f 'run_tests.sh' tests/moe_ep --exec rg -n -C5 'deep_gemm_mega_kernel_vs_reference|torchrun' {} # Check whether other moe_ep tests still self-bootstrap a process group. rg -n -C3 'init_process_group' tests/moe_ep🤖 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, Add an early pytest skip in the single-rank setup before dist.init_process_group, requiring the torchrun rendezvous variables MASTER_ADDR and MASTER_PORT (along with the existing WORLD_SIZE check). Preserve execution when the rendezvous environment is present and prevent plain pytest collection from reaching env:// initialization without it.flashinfer/moe_ep/kernel_src/cutedsl_megamoe/SKILL.md (1)
90-113: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winPropagate the BF16 addition into every enumerated documentation surface. The PR adds a third dtype to this kernel tree, but several existing lists that enumerate the supported dtypes, shims, backends, and src dependencies were not extended. Each site below needs the BF16 entry added.
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/SKILL.md#L90-L113: addshim/bf16.pyto the step-3 signature-audit sentence and to the audit tables. Includemoe_bf16_glu.megamoe_kernel_bf16.Sm100MegaMoEBf16Kernelandsrc.sym_buffer.SymBufferHostfromshim/bf16.py, plus the three newkernel_helpers.pylazy imports:src.token_comm.CombineFormat,moe_nvfp4_swapab.mega_reference.combine_roundtrip_to_fp32, andmoe_bf16_glu.mega_reference_bf16.compute_megamoe_reference.flashinfer/moe_ep/kernel_src/cutedsl_megamoe/SKILL.md#L47-L49: addbf16_bf16_bf16_cutedslto the FI backend path list at line 47 and to the "What NOT to update here" list at lines 129-130.flashinfer/moe_ep/kernel_src/cutedsl_megamoe/TUNING.md#L304-L321: add a BF16 bullet to the knob-system section that states the single fixed geometry from_BF16_TOKEN_KNOBSand the one-entrybf16_candidates()autotune list.flashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/tuner.py#L284-L291: extend thewith_knobsdocstring to list BF16 as a supported config dataclass, and note thatMegaMoEBf16Configdeclarestoken_back_modedirectly instead oftoken_back_by_dispatch.As per coding guidelines: "Keep documentation synchronized with code changes, including
CLAUDE.md, skill files, examples, and documented infrastructure, patterns, error handling, and deprecated approaches."🤖 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/SKILL.md` around lines 90 - 113, Propagate BF16 documentation updates across flashinfer/moe_ep/kernel_src/cutedsl_megamoe/SKILL.md lines 90-113, adding shim/bf16.py and all specified kernel_helpers.py and src dependencies to the audit text and tables; also update lines 47-49 and 129-130 to include bf16_bf16_bf16_cutedsl. In flashinfer/moe_ep/kernel_src/cutedsl_megamoe/TUNING.md lines 304-321, document the fixed geometry from _BF16_TOKEN_KNOBS and the single-entry bf16_candidates() list. In flashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/tuner.py lines 284-291, extend the with_knobs docstring with BF16 support and explain MegaMoEBf16Config uses token_back_mode directly.Source: Coding guidelines
🟠 Major comments (21)
.pre-commit-config.yaml-59-59 (1)
59-59: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winPropagate the SM100 upstream exclusion to every tool-specific configuration.
The new upstream tree is
flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/. Add it to the hook-local mypy exclusion,[tool.mypy].exclude, and Ruffextend-exclude. Also update the nearby comments to document both upstream source roots.
.pre-commit-config.yaml#L59-L59: add the SM100 path to the hook-local mypy exclusion.pyproject.toml#L111-L111: add the SM100 path to[tool.mypy].exclude.pyproject.toml#L126-L126: add the SM100 path to Ruffextend-exclude.Proposed configuration update
- exclude: ^(flashinfer-cubin/|3rdparty/|build/|flashinfer/cute_dsl/attention/fmha/(fmha\.py|fmha_blockscaled\.py|helpers/)|flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/|flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/) + exclude: ^(flashinfer-cubin/|3rdparty/|build/|flashinfer/cute_dsl/attention/fmha/(fmha\.py|fmha_blockscaled\.py|helpers/)|flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/|flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/|flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/) "flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/", + "flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/", "flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/", "flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src", + "flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src", "flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src",As per coding guidelines, keep documentation synchronized with code changes, including documented infrastructure and patterns.
🤖 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 59, The exclusion for the new SM100 upstream source root must be propagated across all tool configurations. Update .pre-commit-config.yaml:59, pyproject.toml:111, and pyproject.toml:126 to include flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/ in the hook-local mypy exclusion, [tool.mypy].exclude, and Ruff extend-exclude respectively; also update the nearby comments in each configuration to document both upstream source roots.Source: Coding guidelines
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_common.py-498-500 (1)
498-500: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winDo not apply the int32 limit to the int64 pack kernel.
_pack_fp4_kernelcomputestl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)(line 522). Its own comment states the int64 cast exists because the combine round-trip packs billions of elements at ep4 / 32768. The new guard at line 569 rejects exactly those launches, so a path that worked now raisesValueError. The block comment at lines 498-499 also does not describe this kernel.Keep the guard on the two int32 kernels only, and correct the comment.
🐛 Proposed fix
-# The flat-index Triton helpers below compute ``program_id * BLOCK + arange`` -# in int32, so any launch with >= 2**31 elements would silently wrap. +# ``_rcp_approx_kernel`` and ``_swiglu_pair_kernel`` compute +# ``program_id * BLOCK + arange`` in int32, so any launch with >= 2**31 +# elements would silently wrap. ``_pack_fp4_kernel`` casts to int64 and is +# exempt. _TRITON_FLAT_INDEX_LIMIT = 2**31if n_pairs > 0: - _check_triton_flat_index(n_pairs, "_pack_f32_to_fp4") triton, kernel = _get_pack_fp4_triton_kernel()Also applies to: 569-569
🤖 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/src/moe_nvfp4_swapab/runner_common.py` around lines 498 - 500, Restrict _TRITON_FLAT_INDEX_LIMIT validation to the two Triton helpers that compute flat indices in int32; remove it from the _pack_fp4_kernel launch path so valid int64-indexed launches remain supported. Update the nearby comment to describe only those int32 kernels.tests/moe_ep/test_moe_ep_sm90_pull_fp8_mega_multirank.py-819-826 (1)
819-826: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winAdd the
arch_hoppermarker to the registration test.
test_sm90_pull_fp8_mega_kernel_is_registeredcarries no marker. Two consequences follow.First, the documented invocation on line 4 uses
-m "gpu_4 and arch_hopper". That expression deselects this test, so it never runs through the documented command.Second, an unmarked test is collected by a bare
pytest tests/moe_eprun. The module docstring on lines 37-41 states that the SM90 and SM100 kernel trees share top-level module names and are mutually exclusive per process. Importing the SM90 backend in a shared process can therefore break the SM100 tests in the same session.
test_sm90_pull_fp8_preprocess_mega_weights_from_bf16on line 782 hasarch_hopperbut nogpu_4, so the documented command also deselects it. Confirm that themega_sm90target inrun_tests.shselects both tests.🛡️ Proposed fix
+@pytest.mark.arch_hopper def test_sm90_pull_fp8_mega_kernel_is_registered():Run the following script to check the runner's selection expressions:
#!/bin/bash # Inspect the mega_sm90 target and every -m expression in the runner. fd -t f 'run_tests.sh' tests/moe_ep --exec rg -n -C6 'mega_sm90|sm90|-m ' {} # Confirm arch_hopper is registered as a marker. rg -n -C3 'arch_hopper' tests/conftest.py pyproject.toml🤖 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 819 - 826, Mark test_sm90_pull_fp8_mega_kernel_is_registered with arch_hopper and verify the mega_sm90 target in run_tests.sh selects both it and test_sm90_pull_fp8_preprocess_mega_weights_from_bf16. Ensure the runner’s marker expression includes both tests while preserving their existing GPU marker requirements and preventing unmarked SM90 tests from running in shared sessions.flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/grid_sync.py-105-129 (1)
105-129: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winAvoid duplicate inline PTX labels in
software_grid_sync.This template emits
SPINandDONElabels throughllvm.inline_asm; if either grid sync block is inlined or duplicated inside one PTX function, the labels collide and PTX assembly can fail. Use unique labels, for example UUID-style label names, and keep each jump target local to that inline asm instance.🤖 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/grid_sync.py` around lines 105 - 129, Update the inline PTX template in software_grid_sync to generate unique per-instance names for the SPIN and DONE labels, such as UUID-style identifiers, and use those names consistently in the corresponding branch instructions. Keep each jump target local to its llvm.inline_asm instance while preserving the existing synchronization logic.flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/dispatch_kernel.py-88-90 (1)
88-90: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winUse DSL bitwise operators for dynamic predicate combinations. Python
and/orfall back to Python bool short-circuiting instead of the CuTe DSL overloads.
- In
flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/dispatch_kernel.py#L88-L90, replaceandwith&in the_iket_activepredicate so both CTA and warp components contribute to the runtime branch.- In
flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/flag_batch.py#L75-L84, replaceorwith|in the flush condition so the phase-change flush is selected at runtime rather than by Python short-circuiting.🤖 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/dispatch_kernel.py` around lines 88 - 90, Replace Python boolean composition with CuTe DSL bitwise predicates: update _iket_active in flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/dispatch_kernel.py lines 88-90 to use & instead of and, and update the flush condition in flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/flag_batch.py lines 75-84 to use | instead of or, preserving runtime evaluation of both dynamic conditions.flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/shim/hopper_fp8.py-543-551 (1)
543-551: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winClear
_megaand_mega_keyright after_release_workspace().Line 549 frees
self._mega.shared_workspace, butself._megaandself._mega_keykeep pointing at the old_CompiledMega. If any step between line 549 and line 641 raises —cute.compile, the kernel constructor,get_workspace_sizes, or the symmetric allocation — the frontend is left holding a_CompiledMegawhoseshared_workspaceis already freed.A later call that resolves to the old compile key then hits the early return at lines 545-546 and launches with a dangling symmetric-heap pointer.
set_gate_up_clampandreleasealready pair the free with_invalidate_compile_cache(); this path should do the same.🐛 Proposed fix
ensure_not_capturing("cute.compile + symmetric-heap allocation") self._release_workspace() + # Drop the stale entry immediately: its shared_workspace is now freed, + # so an exception below must not leave it reachable via the early + # return above. + self._invalidate_compile_cache()🤖 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 543 - 551, Update _ensure_mega_compiled immediately after _release_workspace() to invalidate the cached compilation state by clearing both _mega and _mega_key, matching the existing _invalidate_compile_cache() behavior. Ensure failures during compilation, kernel construction, workspace sizing, or symmetric allocation cannot leave the old _CompiledMega eligible for the early return.flashinfer/moe_ep/backends/mega/kernel/sm90/fp8_fp8_bf16_pull_cutedsl/staging.py-67-69 (1)
67-69: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winMask the pad rows and record 0 tokens before the zero-token early return.
When
num_tokens == 0, this returns beforetopk_idx_out[num_tokens:capacity].fill_(-1)and before_note_staged_tokens. Two stale-state effects follow if a larger staging ran earlier on the same workspace:
topk_idx_outkeeps the previous live routes, so the next launch dispatches stale tokens.staged_tokens()still returns the previous count, socompute(output=None)computesnum_tokensfrom the old staging instead of 0.Reset the routing tail and the staged count before returning.
🐛 Proposed fix for the zero-token path
num_tokens, hidden = hidden_states.shape if num_tokens == 0: + topk_idx_out.fill_(-1) + _note_staged_tokens(topk_idx_out, 0) return🤖 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 - 69, Update the zero-token branch in the staging function to fill topk_idx_out[num_tokens:capacity] with -1 and call _note_staged_tokens with 0 before returning. Preserve the existing behavior for non-empty hidden_states while ensuring staged_tokens() reports zero and no stale routes remain.flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/token_comm.py-26-29 (1)
26-29: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winThe
iketguard catches onlyImportErrorand will not reach the fallback.The sibling
iket_compat.pyin this same directory documents the exact trap:cutlass.cute.experimentalraisesNotImplementedError, notImportError, on CUDA toolkits below 13.1, and it therefore catches(ImportError, NotImplementedError)on both of its import attempts. Its comment names "epilogue.py's ImportError-only guard" as the propagation path to fix.This guard has the same defect. If
from cutlass.cute import iketraisesNotImplementedErroron a CTK-12.9 wheel, the exception escapes andtoken_commfails at module import. That takes the whole SM90 backend down at load time, not at kernel launch.Import through
iket_compatunconditionally. That module already performs the full fallback chain and installs the no-op shim.🐛 Proposed fix
-try: - from cutlass.cute import iket as _iket # type: ignore -except ImportError: # pragma: no cover -- fallback for wheels without cute.iket - from .iket_compat import iket as _iket +# iket_compat performs the full experimental -> cute -> no-op shim fallback and +# catches NotImplementedError, which CTK < 13.1 raises instead of ImportError. +from .iket_compat import iket as _iket🤖 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 26 - 29, Update the _iket import in token_comm to import unconditionally from the local iket_compat module, removing the direct cutlass.cute import guard. Reuse iket_compat’s existing fallback chain and no-op shim so module loading succeeds when cutlass raises either ImportError or NotImplementedError.flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/src/reference.py-237-251 (1)
237-251: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winTrack a separate SF pool accumulator.
expert_pool_block_offsetandexpert_sf_pool_block_offsetcan diverge because the device advances them withceil(prev_valid_count / token_padding_block)andceil(prev_valid_count / sf_padding_block). This oracle advances onlypool_block_offsetand reuses it aspool_block_offset * SFBM, so wrongl1_sf_bufferaddresses are expected unlessBM == SFBM. Use a separatesf_pool_block_offsetadvanced by(T_e + SFBM - 1) // SFBM.🤖 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/reference.py` around lines 237 - 251, Update the reference pooling loop to maintain a separate sf_pool_block_offset for l1_sf_buffer addressing instead of deriving the offset from pool_block_offset. Use sf_pool_block_offset * SFBM when computing sf_pool_token_idx, and advance sf_pool_block_offset by (T_e + SFBM - 1) // SFBM alongside the existing pool_block_offset update.flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_hopper_fp8/kernel_fp8_glu_fc12.py-2276-2284 (1)
2276-2284: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winUse the full per-stage TMA transfer count for
tx_count.
PipelineTmaAsync.createreceives atx_countvalue from multiple TMA producers as one barrier threshold; it is not a per-warp split.self.num_tma_load_bytesalready includes the A-tile, B-tile, and activation-scale bytes for one stage, so divide only when calculating per-producer byte counts, not for the sharedtx_count.🤖 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 2276 - 2284, Update the `PipelineTmaAsync.create` call for `ab_pipeline` to pass the full per-stage transfer count as `tx_count` by removing the division by two from `self.num_tma_load_bytes`; retain any per-producer byte splitting elsewhere in the load calculations.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 winThe TMA alignment checks always assume the NVFP4 element size.
ProblemDescsupports five kinds (line 99-101), and the surrounding code branches onself.kindfor the SF block size (line 115) and the gate/up interleave (line 121). These five_check_tma_leading_dim_aligncalls hard-codeNvfp4DataDtypeforactivation,fc1_weight,fc2_weight, andfc1_output.For an MXFP8 or FP8 kind the element size is 1 byte, not 0.5. The check therefore computes half the real row size and demands
hidden % 32 == 0where the true TMA requirement ishidden % 16 == 0. Valid FP8 shapes are rejected.Derive the dtype from the kind, as
_generate_inputs_skeletonalready does withkind_data_dtype(problem.kind).🐛 Proposed fix
from moe_nvfp4_swapab.runner_common import ( check_tma_leading_dim_align as _check_tma_leading_dim_align, ) + _data_dtype = kind_data_dtype(self.kind) _check_tma_leading_dim_align( "activation", {"k_major": self.hidden}[self.fc1_activation_layout], - Nvfp4DataDtype, + _data_dtype, ) _check_tma_leading_dim_align( "fc1_weight", {"k_major": self.hidden}[self.fc1_weight_layout], - Nvfp4DataDtype, + _data_dtype, ) _check_tma_leading_dim_align( "fc2_weight", {"k_major": self.intermediate // 2, "n_major": self.hidden}[ self.fc2_weight_layout ], - Nvfp4DataDtype, + _data_dtype, ) _check_tma_leading_dim_align( "fc1_output (kernel-internal, fixed k_major)", self.intermediate // 2, - Nvfp4DataDtype, + _data_dtype, )🤖 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, Update the TMA alignment checks in the relevant ProblemDesc initialization flow to derive the element dtype from self.kind via kind_data_dtype, matching _generate_inputs_skeleton, instead of hard-coding Nvfp4DataDtype for activation, fc1_weight, fc2_weight, and fc1_output. Preserve the existing fc2_output_dtype handling so alignment requirements remain correct for all supported kinds.flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/runner_fc12_common.py-1531-1539 (1)
1531-1539: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
misc.enable_debug_checksis never consulted; the debug paths always run.
MiscDesc.enable_debug_checksis defined at line 421 and exposed as--enable_debug_checksat line 1930, whose help text states "Run determinism and fc1 workspace diagnostics during validate."validatecalls_check_kernel_determinismand_validate_fc1_phaseunconditionally.
_check_kernel_determinismperforms a full extra kernel launch plus three device-to-device byte clones, and_validate_fc1_phaseperforms a per-expert dequant readback. Both run on every non-skip_ref_checkinvocation. Gate them on the flag.🐛 Proposed gate
- self._check_kernel_determinism() - self._validate_fc1_phase() + if self.misc.enable_debug_checks: + self._check_kernel_determinism() + self._validate_fc1_phase()🤖 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 1531 - 1539, Update validate around _check_kernel_determinism and _validate_fc1_phase to run these diagnostics only when self.misc.enable_debug_checks is true. Preserve the existing skip_ref_check and reference/input validation behavior, while avoiding both diagnostic calls when the flag is disabled.flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/runner_fc12.py-220-232 (1)
220-232: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winAvoid dropping Blackwell-only kernels into the SM90 tree.
Sm100SwapABSwigluFp4Fc12Kerneltargets Blackwell and is called fromkernel_src/sm90/; keep it under the SM100 source tree or record the unsupported-architecture exception in a clear comment. The same issue applies to the SM100 device gate inbenchmark_p2p.py.🤖 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.py` around lines 220 - 232, The SM90 tree references Blackwell-only functionality without documenting or relocating the architecture-specific code. In flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/runner_fc12.py:220-232, move the Sm100SwapABSwigluFp4Fc12Kernel usage into the SM100 source tree, or add a clear comment documenting the intentional unsupported-architecture exception; apply the same treatment to the SM100 device gate in flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/benchmark_p2p.py:390-392.flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/mega_runner.py-2644-2664 (1)
2644-2664: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winA failing rank skips the barrier-aligned teardown and can hang the job.
Only
NotImplementedErroris caught.validate()raisesAssertionError, andrun_kernel()can raiseValueErrororRuntimeError. Any of those propagates out of thetryand skips the wholeif not _NO_DIST:block.The comment on Lines 2653-2657 states that block exists because unsynchronized teardown deadlocks: the failing rank exits through the interpreter's normal shutdown (running NVSHMEM finalizers under GC) while the passing ranks block in
torch.distributed.barrier(). The multi-rank job then hangs instead of failing.
return_codeis also never set to non-zero, so a skipped kernel launch reports success to the shell.Wrap the run in
try/except/finallyso every rank reaches the barrier andos._exitpath, and propagate a non-zero exit code on failure.🐛 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}") - - if not _NO_DIST: - # nvshmem_free/finalize are collective barriers; an unsynchronized or - # GC-driven teardown deadlocks once per-rank free order diverges. So - # just barrier-align, then os._exit and let the driver reclaim the heap - # on exit. os._exit skips finalizers/GC, hence the manual flush. - torch.cuda.synchronize() - if torch.distributed.is_initialized(): - torch.distributed.barrier() - sys.stdout.flush() - sys.stderr.flush() - os._exit(return_code) - return return_code + return_code = 2 + except BaseException: + # Every rank MUST reach the barrier-aligned teardown below; letting the + # exception escape leaves the passing ranks blocked in barrier() while + # this rank runs NVSHMEM finalizers under GC. + import traceback + traceback.print_exc() + return_code = 1 + finally: + if not _NO_DIST: + # nvshmem_free/finalize are collective barriers; an unsynchronized + # or GC-driven teardown deadlocks once per-rank free order + # diverges. So just barrier-align, then os._exit and let the + # driver reclaim the heap on exit. os._exit skips finalizers/GC, + # hence the manual flush. + torch.cuda.synchronize() + if torch.distributed.is_initialized(): + torch.distributed.barrier() + sys.stdout.flush() + sys.stderr.flush() + os._exit(return_code) + return return_code🤖 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() flow to catch validation and kernel-launch failures such as AssertionError, ValueError, and RuntimeError, set return_code to a non-zero value, and report the failure appropriately. Move the barrier-aligned teardown under a finally path so every rank reaches torch.cuda.synchronize(), the distributed barrier, output flushes, and os._exit, while preserving the existing NotImplementedError handling and ensuring skipped launches do not report success.flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/src/moe_nvfp4_swapab/moe_persistent_scheduler.py-1611-1648 (1)
1611-1648: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winGuard drain fields when
num_drain_warps=0.
make_storage_structacceptsnum_drain_warps=0by default, which creates zero-lengthdrain_mbaranddrain_responsefields. If a kernel using this storage callsdrain_empty_tiles, it writes an mbarrier and CLC response outside the allocated SMEM range, corrupting the neighboring scheduler fields.Require at least one drain slot when drain is enabled, or add a guard inside
drain_empty_tiles.🤖 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 1611 - 1648, Guard the zero-drain configuration in make_storage_struct and drain_empty_tiles: when drain functionality is enabled, require num_drain_warps to be at least one before constructing or using drain_mbar and drain_response. Ensure drain_empty_tiles does not access these fields when no drain slots were allocated, preventing writes into neighboring scheduler storage.flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_bf16_glu/kernel_bf16_glu_fc12.py-485-494 (1)
485-494: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winGuard against a non-positive AB stage count.
num_ab_stagehas no lower bound. Iffixed_overheadapproaches or exceedssmem_capacity // occupancy, the expression yields0or a negative value. Python floor division of a negative numerator rounds toward negative infinity, so the result can be-1or lower.That value then flows into
sm100_utils.make_smem_layout_a(..., self.num_a_stage)at line 397 and intopipeline.PipelineTmaUmma.create(num_stages=self.num_a_stage, ...)at line 1147. The failure surfaces as an opaque layout or pipeline error instead of an SMEM-budget error.The path is reachable:
_smem_misc_budget_bytesis documented at lines 443-447 as a hook the MegaMoE subclass extends, andc_bytes_totalgrows withgenerate_c=True.🛡️ Proposed guard
num_ab_stage = ( smem_capacity // occupancy - fixed_overhead ) // ab_bytes_per_stage + if num_ab_stage < 1: + raise ValueError( + f"SMEM budget exhausted: capacity/occupancy=" + f"{smem_capacity // occupancy} B, fixed_overhead={fixed_overhead} B " + f"(misc={self._smem_misc_budget_bytes()}, c={c_bytes_total}), " + f"ab_bytes_per_stage={ab_bytes_per_stage} B leaves " + f"num_ab_stage={num_ab_stage}; at least 1 stage is required." + ) num_a_stage = num_ab_stage num_b_stage = num_ab_stage🤖 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/src/moe_bf16_glu/kernel_bf16_glu_fc12.py` around lines 485 - 494, Guard the computed num_ab_stage in the stage-count setup so it cannot be zero or negative when fixed_overhead consumes the available shared memory. Validate the SMEM budget before assigning num_a_stage and num_b_stage, and raise the established SMEM-budget error with relevant capacity and overhead details rather than allowing invalid values to reach make_smem_layout_a or PipelineTmaUmma.create.flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_bf16_glu/kernel_bf16_glu_fc12.py-1081-1085 (1)
1081-1085: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winCompute the fc2 spin threshold once and reuse it.
Lines 1081-1085 and lines 1270-1273 compute the same expression:
ceil(fc1_weight_gemm.shape[0] / cta_tile_shape_mnk[1]) * epilogue._atom_thr_sizeThe first value goes to
GluBf16Fc12SchedExtensionat line 1089. The second is the spin bound used by the TMA-A warp at line 1412. The comment at lines 1266-1269 asserts the two must match, but nothing enforces it.If one expression is edited and the other is not, the TMA-A warp stops spinning before every fc1 N-tile has landed. The fc2 phase then reads a partially written
fc1_output, which produces wrong results without any error. The scheduler extension docstring inflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_bf16_glu/custom_ext_bf16.pydescribes the same coupling.Bind the value once above both uses.
♻️ Proposed refactor
- ext_fc2_spin_threshold = ( + fc2_spin_threshold = ( (fc1_weight_gemm.shape[0] + self.cta_tile_shape_mnk[1] - 1) // self.cta_tile_shape_mnk[1] * self.epilogue._atom_thr_size ) ext = GluBf16Fc12SchedExtension( fc1_done_counter_ptr=fc1_done_counter.iterator, - fc2_spin_threshold=ext_fc2_spin_threshold, + fc2_spin_threshold=fc2_spin_threshold,Then delete the duplicate at lines 1266-1273 and keep the explanatory comment above the single definition.
🤖 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/src/moe_bf16_glu/kernel_bf16_glu_fc12.py` around lines 1081 - 1085, Compute the fc2 spin threshold once before the scheduler extension setup, using the existing expression and explanatory coupling comment. Reuse that single value for both GluBf16Fc12SchedExtension and the TMA-A warp spin bound, and remove the duplicate calculation near the later fc2 phase.flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_bf16_glu/mega_reference_bf16.py-104-156 (1)
104-156: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
ref_compute_graphis validated and then ignored.Line 81 declares
ref_compute_graphas a required parameter. Lines 104-108 validate it. Line 156 then discards it and derivesexpert_graphfromapply_topk_in_fc1instead. Every call toreference_expert_fc12at line 195 passesexpert_graph, neverref_compute_graph.A caller that passes
ref_compute_graph="deepgemm"together withapply_topk_in_fc1=Falsereceives the"transformers"graph without any warning. The top-k weight is then applied in a different place than the caller requested. This is a correctness oracle, so a silent parameter override can mask a real kernel mismatch or manufacture a false failure.Either remove the parameter, or make the two inputs consistent and reject a conflicting combination.
🛠️ Proposed fix
expert_graph = "deepgemm" if apply_topk_in_fc1 else "transformers" + if ref_compute_graph != expert_graph: + raise ValueError( + f"ref_compute_graph={ref_compute_graph!r} conflicts with " + f"apply_topk_in_fc1={apply_topk_in_fc1}, which implies " + f"{expert_graph!r}. Pass matching values." + )🤖 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/src/moe_bf16_glu/mega_reference_bf16.py` around lines 104 - 156, Use the validated ref_compute_graph value when setting expert_graph instead of deriving it solely from apply_topk_in_fc1. Enforce consistency between ref_compute_graph and apply_topk_in_fc1, rejecting combinations where the requested graph disagrees with the top-k placement, before calling reference_expert_fc12.flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_bf16_glu/kernel_bf16_glu_fc12.py-748-763 (1)
748-763: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winCross-validate the fc1/fc2 shape relationships on the host.
Three derived dimensions are read independently and never compared:
- Line 748:
intermediate_downproj = fc1_output.shape[1].- Line 734:
experts, hidden_b, intermediate_gateup = fc1_weight.shape.- Line 763:
experts2, intermediate_downproj_b2, hidden_b2 = fc2_weight.shape.The kernel assumes
experts2 == experts,intermediate_downproj_b2 == intermediate_gateup // 2,hidden_b2 == hidden, andintermediate_downproj == intermediate_gateup // 2. Nothing enforces those relations. A caller that passes mismatched fc1 and fc2 weights gets out-of-bounds device reads through the TMA atoms built at lines 857 and 914, not a host-side error. The scheduler at line 964 derives its tile counts fromfc1_weightonly, so the fc2 phase indexesfc2_weightwith tile indices sized for a different tensor.All the operands are already unpacked, so the check costs nothing.
🛡️ Proposed guard
experts2, intermediate_downproj_b2, hidden_b2 = fc2_weight.shape + if cutlass.const_expr( + experts2 != experts + or hidden_b2 != hidden_b + or intermediate_downproj_b2 * 2 != intermediate_gateup + or intermediate_downproj != intermediate_downproj_b2 + ): + raise ValueError( + f"inconsistent fused fc12 shapes: " + f"fc1_weight={(experts, hidden_b, intermediate_gateup)}, " + f"fc2_weight={(experts2, intermediate_downproj_b2, hidden_b2)}, " + f"fc1_output N={intermediate_downproj}; expected " + f"experts and hidden to match and " + f"intermediate_downproj == intermediate_gateup // 2." + )Apply the guard only where the dimensions are codegen-time constants. If any of them stays runtime-dynamic,
cutlass.const_exprcannot evaluate the comparison, so move the check to the host runner 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/cutedsl_megamoe/src/moe_bf16_glu/kernel_bf16_glu_fc12.py` around lines 748 - 763, Validate the unpacked fc1/fc2 dimensions before building the fc2 GEMM transforms, ensuring experts2 equals experts, intermediate_downproj_b2 equals intermediate_gateup // 2, hidden_b2 equals hidden, and intermediate_downproj equals intermediate_gateup // 2. Apply these checks only when the dimensions are codegen-time constants using the existing const-expression mechanism; otherwise add equivalent validation in the host runner, and reject mismatches before constructing TMA atoms or scheduling work.flashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/bf16.py-175-181 (1)
175-181: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winClear the compiled state when the workspace is released.
_release_workspace()freesself._mega.shared_workspacebut leavesself._megaandself._mega_keyset._ensure_compiledcalls it at line 180 and only reassigns both fields at lines 239-240, aftercute.compilereturns.If
Sm100MegaMoEBf16Kernel(...),get_workspace_sizes(),sym_zeros(...), orcute.compile(...)raises, the frontend is left holding a_CompiledMegawhose symmetric workspace is already freed, together with the old_mega_key. Two failure modes follow:
- The caller retries with the same config. Lines 177-178 hit the cache and return the stale object.
mega.compiledis the old value andmega.shared_workspaceis freed memory.- The caller changes the config.
_release_workspace()runs again on the same already-freedshared_workspace, which double-frees the symmetric-heap allocation.Null the state inside
_release_workspace()so every caller gets the same invariant.🐛 Proposed fix
def _release_workspace(self) -> None: if self._mega is not None: free_sym_tensor(self._mega.shared_workspace) + self._mega = None + self._mega_key = NoneThe existing
self._mega = None/self._mega_key = Nonelines inset_gate_up_clamp,apply_knobs, andrelease()then become redundant and can be removed.Also applies to: 360-362
🤖 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/shim/bf16.py` around lines 175 - 181, Update _release_workspace() to clear both self._mega and self._mega_key after releasing the compiled workspace, ensuring failures during _ensure_compiled() cannot leave stale compiled state or cause a double free. Remove the now-redundant state-nulling assignments from set_gate_up_clamp, apply_knobs, and release().flashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/_paths.py-41-53 (1)
41-53: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winRequire
__file__to exist before bypassing the sentinel guard.
sys.modules.get(name)can return a namespace package, where__file__isNonewhile__path__still holds loaded directories. If a sentinel directory is a namespace package, the current check allows the cross-tree conflict to pass; treat a namespace import as an error or inspect__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/cutedsl_megamoe/shim/_paths.py` around lines 41 - 53, Update the sentinel validation loop around _SENTINEL_MODULES so modules without a __file__ cannot bypass the guard: when a sentinel module is loaded, treat a missing __file__ as a conflict or inspect its __path__ entries and reject any directory outside src_dir. Preserve the existing RuntimeError behavior and message for cross-tree imports.
| if cutlass.const_expr(self.overlapping_accum and unroll_tile_cnt == 0): | ||
| release_after_ldtm = True | ||
| else: | ||
| release_after_ldtm = False | ||
| for i in cutlass.range(remain_subtile_cnt, unroll=1): | ||
| # for i in cutlass.range_constexpr(remain_subtile_cnt): | ||
| real_i = i + unroll_tile_cnt | ||
| if cutlass.const_expr(self.overlapping_accum): | ||
| subtile_idx = ( | ||
| cutlass.Int32(real_i + self.subtile_cnt) - is_odd_turn | ||
| ) % cutlass.Int32(self.subtile_cnt) | ||
| else: | ||
| subtile_idx = cutlass.Int32(real_i) | ||
|
|
||
| if subtile_idx * cutlass.Int32(self._EpilogueTokenTileSize) < valid_tokens: | ||
| self.run_subtile( | ||
| subtile_idx=subtile_idx, | ||
| tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx], | ||
| preload_acc=None, | ||
| fc2_output_router=fc2_output_router, | ||
| alpha_val=alpha_val, | ||
| release_after_ldtm=release_after_ldtm, | ||
| acc_pipeline=acc_pipeline, | ||
| acc_consumer_state=acc_consumer_state, | ||
| ) | ||
| release_after_ldtm = False | ||
|
|
||
| # Non-overlap-path release: at the natural task-tile boundary. | ||
| if cutlass.const_expr(not self.overlapping_accum): | ||
| cute.arch.fence_view_async_tmem_load() | ||
| acc_pipeline.consumer_release(acc_consumer_state) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
The accumulator is never released when every fc2 subtile is skipped.
In the overlapping_accum path with unroll_tile_cnt == 0, acc_pipeline.consumer_release runs only inside run_subtile (Line 2055-2057), which executes only when subtile_idx * _EpilogueTokenTileSize < valid_tokens. If valid_tokens is 0 for a work tile, no subtile runs, the release never fires, and the block at Line 2032 is const-expr disabled. The producer then blocks on producer_acquire for that stage.
epilogue.py::_run_fc2_bulk_task_tile guards exactly this case (it releases when the first subtile index is past valid_tokens); this refactored path dropped that guard.
🔒️ Proposed fix: release the accumulator when no subtile ran
if cutlass.const_expr(self.overlapping_accum and unroll_tile_cnt == 0):
release_after_ldtm = True
else:
release_after_ldtm = False
for i in cutlass.range(remain_subtile_cnt, unroll=1):
# for i in cutlass.range_constexpr(remain_subtile_cnt):
real_i = i + unroll_tile_cnt
if cutlass.const_expr(self.overlapping_accum):
subtile_idx = (
cutlass.Int32(real_i + self.subtile_cnt) - is_odd_turn
) % cutlass.Int32(self.subtile_cnt)
else:
subtile_idx = cutlass.Int32(real_i)
if subtile_idx * cutlass.Int32(self._EpilogueTokenTileSize) < valid_tokens:
self.run_subtile(
...
)
release_after_ldtm = False
+ # Overlap path with no unroll: if every subtile was skipped the release
+ # inside run_subtile never fired; release here so the mma producer is
+ # not left blocked on this acc stage.
+ if cutlass.const_expr(self.overlapping_accum and unroll_tile_cnt == 0):
+ if release_after_ldtm:
+ cute.arch.fence_view_async_tmem_load()
+ acc_pipeline.consumer_release(acc_consumer_state)
+
# Non-overlap-path release: at the natural task-tile boundary.
if cutlass.const_expr(not self.overlapping_accum):
cute.arch.fence_view_async_tmem_load()
acc_pipeline.consumer_release(acc_consumer_state)📝 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 cutlass.const_expr(self.overlapping_accum and unroll_tile_cnt == 0): | |
| release_after_ldtm = True | |
| else: | |
| release_after_ldtm = False | |
| for i in cutlass.range(remain_subtile_cnt, unroll=1): | |
| # for i in cutlass.range_constexpr(remain_subtile_cnt): | |
| real_i = i + unroll_tile_cnt | |
| if cutlass.const_expr(self.overlapping_accum): | |
| subtile_idx = ( | |
| cutlass.Int32(real_i + self.subtile_cnt) - is_odd_turn | |
| ) % cutlass.Int32(self.subtile_cnt) | |
| else: | |
| subtile_idx = cutlass.Int32(real_i) | |
| if subtile_idx * cutlass.Int32(self._EpilogueTokenTileSize) < valid_tokens: | |
| self.run_subtile( | |
| subtile_idx=subtile_idx, | |
| tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx], | |
| preload_acc=None, | |
| fc2_output_router=fc2_output_router, | |
| alpha_val=alpha_val, | |
| release_after_ldtm=release_after_ldtm, | |
| acc_pipeline=acc_pipeline, | |
| acc_consumer_state=acc_consumer_state, | |
| ) | |
| release_after_ldtm = False | |
| # Non-overlap-path release: at the natural task-tile boundary. | |
| if cutlass.const_expr(not self.overlapping_accum): | |
| cute.arch.fence_view_async_tmem_load() | |
| acc_pipeline.consumer_release(acc_consumer_state) | |
| if cutlass.const_expr(self.overlapping_accum and unroll_tile_cnt == 0): | |
| release_after_ldtm = True | |
| else: | |
| release_after_ldtm = False | |
| for i in cutlass.range(remain_subtile_cnt, unroll=1): | |
| # for i in cutlass.range_constexpr(remain_subtile_cnt): | |
| real_i = i + unroll_tile_cnt | |
| if cutlass.const_expr(self.overlapping_accum): | |
| subtile_idx = ( | |
| cutlass.Int32(real_i + self.subtile_cnt) - is_odd_turn | |
| ) % cutlass.Int32(self.subtile_cnt) | |
| else: | |
| subtile_idx = cutlass.Int32(real_i) | |
| if subtile_idx * cutlass.Int32(self._EpilogueTokenTileSize) < valid_tokens: | |
| self.run_subtile( | |
| subtile_idx=subtile_idx, | |
| tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[None, None, subtile_idx], | |
| preload_acc=None, | |
| fc2_output_router=fc2_output_router, | |
| alpha_val=alpha_val, | |
| release_after_ldtm=release_after_ldtm, | |
| acc_pipeline=acc_pipeline, | |
| acc_consumer_state=acc_consumer_state, | |
| ) | |
| release_after_ldtm = False | |
| # Overlap path with no unroll: release if every subtile was skipped. | |
| if cutlass.const_expr(self.overlapping_accum and unroll_tile_cnt == 0): | |
| if release_after_ldtm: | |
| cute.arch.fence_view_async_tmem_load() | |
| acc_pipeline.consumer_release(acc_consumer_state) | |
| # Non-overlap-path release: at the natural task-tile boundary. | |
| if cutlass.const_expr(not self.overlapping_accum): | |
| cute.arch.fence_view_async_tmem_load() | |
| acc_pipeline.consumer_release(acc_consumer_state) |
🤖 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 2004 - 2034, Update the accumulator-release logic in the refactored
epilogue task-tile loop so the overlapping_accum path releases
acc_consumer_state when no subtile passes the valid_tokens check, including the
unroll_tile_cnt == 0 case. Track whether run_subtile executed or apply the
equivalent first-subtile guard used by _run_fc2_bulk_task_tile, while preserving
the existing release_after_ldtm behavior and non-overlap release path.
| new_params = new_from_mlir_values(self.params, values[idx : idx + 3]) | ||
| idx += 3 |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🔴 Critical | ⚡ Quick win
Fix the hardcoded params length in the static scheduler deserializer.
MoEStaticSchedulerParams.__extract_mlir_values__ is type-discriminated. It emits one value per field only when that field is Int32. When expert_cnt, intermediate, or hidden is a Python int, it emits fewer than 3 values. This deserializer always consumes 3. Every following slice then reads the wrong values, and the assert idx == len(values) inside the params deserializer fires or silently mis-binds offs.
The dynamic scheduler already uses a length-derived _take helper. Use the same pattern here.
🐛 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📝 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.
| 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 1207 - 1208, Update the static scheduler deserialization around
MoEStaticSchedulerParams and new_from_mlir_values to consume the parameter count
derived from the emitted MLIR fields rather than always slicing three values.
Mirror the dynamic scheduler’s _take helper pattern, advance idx by the consumed
length, and preserve correct subsequent slicing for offs and remaining values.
| TensormapDescBytes = 128 | ||
| TensormapDescBytes = 64 # {$nv-internal-release} |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🔴 Critical | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Compare TensormapDescBytes across vendored kernel trees and find leftover single-line internal markers.
set -euo pipefail
rg -nP --type=py -C3 'TensormapDescBytes'
rg -n 'nv-internal-release' --type=py | head -50Repository: flashinfer-ai/flashinfer
Length of output: 2070
🌐 Web query:
CUDA tma tensormap descriptor size 128 bytes CUtensorMap
💡 Result:
The CUtensorMap struct is a 128-byte opaque data structure used as a descriptor for the Tensor Memory Accelerator (TMA) on Hopper (SM90) and newer GPU architectures [1][2][3]. This descriptor is required to be 128-byte aligned and contains the necessary metadata for TMA hardware to perform efficient, asynchronous multidimensional data transfers between global and shared memory [1][2][4]. Key characteristics of the CUtensorMap include: - Size and Alignment: It is exactly 128 bytes in size and must be aligned to a 128-byte boundary [2][4][5]. - Purpose: It acts as a reusable metadata object that stores tensor information—such as the base pointer, data type, tensor shape, strides, swizzling patterns, and memory layout—enabling the TMA hardware to execute complex memory copy operations without needing individual pointers for every access [2][4][6]. - Usage: The descriptor is typically created on the host using the CUDA Driver API (e.g., cuTensorMapEncodeTiled) [2][4]. Once created, it is passed to a CUDA kernel as a grid-constant argument or shared memory object so the TMA hardware can read the configuration directly during execution [7][2][6]. In practice, because it is an opaque structure, developers often define a minimal version of the struct in their kernels to satisfy compiler alignment requirements without requiring the full driver header [7]. For example: struct align(128) CUtensorMap { unsigned long long opaque; }; This alignment and size ensure compatibility with the TMA unit's requirements for hardware-accelerated memory movement [1][2][5].
Citations:
- 1: https://docs.nvidia.com/cuda/cuda-driver-api/structCUtensorMap.html
- 2: https://cudacourseh100.github.io/pages/lesson-4.html
- 3: https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.TensorMapDescriptor.html
- 4: https://cudacourseh100.github.io/H100-Course/slides/4.%20cuTensorMap.pdf
- 5: https://docs.rs/cudarc/latest/cudarc/driver/sys/type.CUtensorMap.html
- 6: https://docs.nvidia.com/cutlass/4.2.1/media/docs/cpp/cute/0z_tma_tensors.html
- 7: https://github.com/NVIDIA/cuda-python/blob/2a2186c8/cuda_core/examples/tma_tensor_map.py
Use the public 128-byte tensormap descriptor size.
Keep TensormapDescBytes = 128 and remove the 64 assignment with the {$nv-internal-release} marker. cuda.core.TensorMapDescriptor/CUDA TMA descriptors are 128-byte aligned, while this file currently uses 64 bytes to compute descriptor strides and get_workspace_size, which can make cpasync.copy_tensormap overwrite the next descriptor slot and under-size the workspace.
🤖 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_utils.py`
around lines 49 - 50, Restore TensormapDescBytes to 128 in moe_utils.py and
remove the 64-byte {$nv-internal-release} assignment. Ensure descriptor stride
calculations and get_workspace_size use the restored public CUDA TMA descriptor
size.
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/backends/mega/kernel/sm100/fp8_fp4_bf16_deepgemm/__init__.py`:
- Around line 2-7: Preserve the deprecated DeepGemmMegaMoeConfig compatibility
alias in the module namespace by aliasing it to
Sm100_Fp8_Fp4_Bf16_Deepgemm_MegaMoeConfig, and include the deprecated name in
__all__. Keep the existing deep_gemm_mega backend alias support unchanged.
🪄 Autofix
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: 0acfa5c0-5a3d-4564-a025-e85b63c7f62e
📒 Files selected for processing (16)
docs/design_docs/moe_ep_architecture.mddocs/design_docs/moe_ep_runbook.mdflashinfer/moe_ep/__init__.pyflashinfer/moe_ep/backends/mega/kernel/sm100/__init__.pyflashinfer/moe_ep/backends/mega/kernel/sm100/fp8_fp4_bf16_deepgemm/__init__.pyflashinfer/moe_ep/backends/mega/kernel/sm100/fp8_fp4_bf16_deepgemm/backend.pyflashinfer/moe_ep/backends/mega/kernel/sm100/fp8_fp4_bf16_deepgemm/config.pyflashinfer/moe_ep/backends/mega/kernel/sm100/fp8_fp4_bf16_deepgemm/staging.pyflashinfer/moe_ep/backends/mega/kernel/sm100/fp8_fp4_bf16_deepgemm/weights.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/TUNING.mdtests/moe_ep/test_deep_gemm_mega_kernel_vs_reference.pytests/moe_ep/test_fused_quant_stage.pytests/moe_ep/test_layer_factory.pytests/moe_ep/test_mega_layer_validation.pytests/moe_ep/test_moe_ep_deep_gemm_mega_multirank.pytests/moe_ep/test_moe_ep_deep_gemm_skew_determinism.py
🚧 Files skipped from review as they are similar to previous changes (10)
- tests/moe_ep/test_moe_ep_deep_gemm_skew_determinism.py
- flashinfer/moe_ep/init.py
- tests/moe_ep/test_moe_ep_deep_gemm_mega_multirank.py
- tests/moe_ep/test_layer_factory.py
- docs/design_docs/moe_ep_runbook.md
- tests/moe_ep/test_fused_quant_stage.py
- docs/design_docs/moe_ep_architecture.md
- flashinfer/moe_ep/kernel_src/cutedsl_megamoe/TUNING.md
- tests/moe_ep/test_mega_layer_validation.py
- flashinfer/moe_ep/backends/mega/kernel/sm100/init.py
| from .config import Sm100_Fp8_Fp4_Bf16_Deepgemm_MegaMoeConfig | ||
| from .weights import TransformedMegaWeights, preprocess_mega_weights | ||
|
|
||
| __all__ = [ | ||
| "DeepGemmMegaKernelBackend", | ||
| "DeepGemmMegaMoeConfig", | ||
| "Sm100_Fp8_Fp4_Bf16_Deepgemm_MegaMoeConfig", |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Preserve the deprecated configuration alias.
This change removes DeepGemmMegaMoeConfig from the module namespace and __all__. Existing callers that import the previously exported symbol will fail at import time. Keep an alias to Sm100_Fp8_Fp4_Bf16_Deepgemm_MegaMoeConfig while the deep_gemm_mega backend alias remains supported.
Proposed compatibility alias
from .config import Sm100_Fp8_Fp4_Bf16_Deepgemm_MegaMoeConfig
+DeepGemmMegaMoeConfig = Sm100_Fp8_Fp4_Bf16_Deepgemm_MegaMoeConfig
@@
"Sm100_Fp8_Fp4_Bf16_Deepgemm_MegaMoeConfig",
+ "DeepGemmMegaMoeConfig",📝 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 .config import Sm100_Fp8_Fp4_Bf16_Deepgemm_MegaMoeConfig | |
| from .weights import TransformedMegaWeights, preprocess_mega_weights | |
| __all__ = [ | |
| "DeepGemmMegaKernelBackend", | |
| "DeepGemmMegaMoeConfig", | |
| "Sm100_Fp8_Fp4_Bf16_Deepgemm_MegaMoeConfig", | |
| from .config import Sm100_Fp8_Fp4_Bf16_Deepgemm_MegaMoeConfig | |
| DeepGemmMegaMoeConfig = Sm100_Fp8_Fp4_Bf16_Deepgemm_MegaMoeConfig | |
| from .weights import TransformedMegaWeights, preprocess_mega_weights | |
| __all__ = [ | |
| "DeepGemmMegaKernelBackend", | |
| "Sm100_Fp8_Fp4_Bf16_Deepgemm_MegaMoeConfig", | |
| "DeepGemmMegaMoeConfig", |
🤖 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/fp8_fp4_bf16_deepgemm/__init__.py`
around lines 2 - 7, Preserve the deprecated DeepGemmMegaMoeConfig compatibility
alias in the module namespace by aliasing it to
Sm100_Fp8_Fp4_Bf16_Deepgemm_MegaMoeConfig, and include the deprecated name in
__all__. Keep the existing deep_gemm_mega backend alias support unchanged.
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/cutedsl_megamoe/VENDOR.md`:
- Around line 29-40: Document the recorded base vendored snapshot before the
ahead-of-drop description, filling in the immutable upstream URL, base commit
SHA, and sync date. Update the provenance section in VENDOR.md while preserving
the existing ahead-of-drop details for inputs_process.py and host_utils.py.
🪄 Autofix
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: e114497b-c20e-43a3-8bfa-52b89fc297a4
📒 Files selected for processing (5)
.pre-commit-config.yamldocs/design_docs/moe_ep_runbook.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 (2)
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/VENDOR.md (1)
57-60: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winList the BF16 backend as a consumer.
Add
backends/mega/kernel/sm100/bf16_bf16_bf16_cutedsl/to this list. The new BF16 backend consumes this vendored package. The incomplete list weakens re-sync impact review.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 `@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/VENDOR.md` around lines 57 - 60, Add the BF16 backend path backends/mega/kernel/sm100/bf16_bf16_bf16_cutedsl/ to the Consumers list in VENDOR.md, alongside the existing NVFP4 and MXFP8 entries.Source: Coding guidelines
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/common/host_utils.py (1)
83-108: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy liftResolve the mixed vendor revisions before merge.
mxfp8_quantize_per_block_32does not provide themxfp8_quantize_per_block_32_rowsymbol required by the retained ahead-of-dropinputs_process.py. As documented inVENDOR.md, this makes the MXFP8src.inputs_processharness fail withImportError.Use one consistent upstream snapshot. Either revert
inputs_process.pywithhost_utils.pyto the recorded drop, or complete the dependent upstream migration atomically.🤖 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/src/common/host_utils.py` around lines 83 - 108, Resolve the vendor snapshot mismatch between mxfp8_quantize_per_block_32 in host_utils.py and the dependent src.inputs_process import. Use one consistent upstream revision by either reverting both files to the recorded VENDOR.md drop or completing the migration so the required mxfp8_quantize_per_block_32_row symbol and all related call sites are present atomically.
🤖 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/src/common/host_utils.py`:
- Around line 83-108: Resolve the vendor snapshot mismatch between
mxfp8_quantize_per_block_32 in host_utils.py and the dependent
src.inputs_process import. Use one consistent upstream revision by either
reverting both files to the recorded VENDOR.md drop or completing the migration
so the required mxfp8_quantize_per_block_32_row symbol and all related call
sites are present atomically.
In `@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/VENDOR.md`:
- Around line 57-60: Add the BF16 backend path
backends/mega/kernel/sm100/bf16_bf16_bf16_cutedsl/ to the Consumers list in
VENDOR.md, alongside the existing NVFP4 and MXFP8 entries.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 32d7b52c-4953-4212-8658-df09bfa4e5aa
📒 Files selected for processing (3)
flashinfer/moe_ep/backends/mega/kernel/sm100/bf16_bf16_bf16_cutedsl/backend.pyflashinfer/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/backends/mega/kernel/sm100/bf16_bf16_bf16_cutedsl/backend.py
- test_deep_gemm_mega_kernel_vs_reference: skip cleanly under plain pytest when the torchrun rendezvous env (MASTER_ADDR/MASTER_PORT) is absent instead of failing in dist.init_process_group. - Propagate bf16 through the drop docs the port missed: SKILL.md layer isolation + shim-audit tables + what-not-to-update list gain the bf16_bf16_bf16_cutedsl backend and shim/bf16.py entries; TUNING.md knob-system section documents the single bf16 default profile; shim/tuner.py with_knobs docstring covers the BF16 config. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- test_deep_gemm_mega_kernel_vs_reference: skip cleanly under plain pytest when the torchrun rendezvous env (MASTER_ADDR/MASTER_PORT) is absent instead of failing in dist.init_process_group. - Propagate bf16 through the drop docs the port missed: SKILL.md layer isolation + shim-audit tables + what-not-to-update list gain the bf16_bf16_bf16_cutedsl backend and shim/bf16.py entries; TUNING.md knob-system section documents the single bf16 default profile; shim/tuner.py with_knobs docstring covers the BF16 config. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
fa4970f to
4922cd9
Compare
…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>
- test_deep_gemm_mega_kernel_vs_reference: skip cleanly under plain pytest when the torchrun rendezvous env (MASTER_ADDR/MASTER_PORT) is absent instead of failing in dist.init_process_group. - Propagate bf16 through the drop docs the port missed: SKILL.md layer isolation + shim-audit tables + what-not-to-update list gain the bf16_bf16_bf16_cutedsl backend and shim/bf16.py entries; TUNING.md knob-system section documents the single bf16 default profile; shim/tuner.py with_knobs docstring covers the BF16 config. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
4b6b28e to
6d9d79b
Compare
|
scheduler
|
…r-ai#4386 Lists the sm100 bf16 mega backend (PR flashinfer-ai#4386, sequenced to merge ahead of this PR) so the layout line already matches the post-merge union and the rebase over flashinfer-ai#4386 resolves mechanically. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
@flashinfer-bot run |
|
/bot run tests/moe_ep |
|
[FAILED] Pipeline #63146628 — 26/30 executed test jobs passed Compared with nightly #63077496. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 5/6 passed
Failure detailsNew relative to nightly (attribution uncertain)
Timeouts, infrastructure, or incomplete jobs
|
…r-ai#4120) Port of upstream draft PR flashinfer-ai#4120 (BF16 MegaMOE integration) onto the restructured taxonomy layout: - kernel drop: kernel_src/cutedsl_megamoe/src/moe_bf16_glu/ + shim/bf16.py, plus the PR's bf16-enabling generalizations of the shared mxfp8/nvfp4/src kernel sources (epilogue fc1_output width guards, epi_flag_batch as a (fc1, fc2) pair, TopkReduce sm_arch parameter, iket ranges). - backend: backends/mega/kernel/sm100/bf16_bf16_bf16_cutedsl/ with taxonomy config Sm100_Bf16_Bf16_Bf16_Cutedsl_MegaMoeConfig, kernel_name sm100_bf16_bf16_bf16_cutedsl and deprecated alias bf16_cutedsl. - tests: bf16 oracle + config + mega multirank wired into run_tests.sh. Deviations from the PR (reviewed intentionally): - PR tree was unformatted; content re-formatted with the repo-pinned ruff 0.12.8 before merging so vendored-file diffs stay semantic. - kept our validated Triton reference helpers in moe_nvfp4_swapab (runner_common/mega_runner) instead of the PR's cute_ref_ops.py rewrite; bf16 does not use cute_ref_ops, so the file is not imported. - kept requirements.txt floors (cutlass-dsl>=4.5.0 for the 4.5.2 WAR chain, tvm-ffi>=0.1.6); the PR bumped both. - dropped PR's stale reverts (zip strict=False, exception chaining) and fixed two latent error-path bugs (valid_ab_tuple NameError, undefined 'testing' module in mega_reference_bf16). PR flashinfer-ai#4120 validated as-is beforehand: 13 pytest + 40/40 functional/mega harness cases green on 4x GB200 with nvidia-cutlass-dsl 4.6.1. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
validate_init passed the deep_gemm-oriented alignment=128 default to validate_mega_fleet_params, rejecting shapes the bf16 kernel supports — the drop's own shim bound (shim/bf16.py) is hidden % 32 / intermediate % 64, which the backend already enforces explicitly. Pass alignment=32 so the shared gate matches; the stricter intermediate bound stays. Unlocks gpt-oss-120b-class geometry (hidden = inter = 2880). Validated on 8x B200 job 2384696: bf16 config + oracle + 4-rank multirank green, and the 2880/2880/128e/top4 EP8 geometry runs at 0.289-0.290% rel-L2 vs the dense bf16 reference (pure bf16 rounding) at 8/512/8192 tok/rank. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- test_deep_gemm_mega_kernel_vs_reference: skip cleanly under plain pytest when the torchrun rendezvous env (MASTER_ADDR/MASTER_PORT) is absent instead of failing in dist.init_process_group. - Propagate bf16 through the drop docs the port missed: SKILL.md layer isolation + shim-audit tables + what-not-to-update list gain the bf16_bf16_bf16_cutedsl backend and shim/bf16.py entries; TUNING.md knob-system section documents the single bf16 default profile; shim/tuner.py with_knobs docstring covers the BF16 config. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
6d9d79b to
250bec4
Compare
|
/bot run tests/moe_ep |
|
[SUCCESS] Pipeline #63378304: 16/16 executed test jobs passed |
…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>
|
@flashinfer-bot run |
… (MXFP8 + NVFP4) (#4531) ## Summary Fixes a livelock in the MXFP8 and NVFP4 CuTeDSL MegaMoE kernels when `in_kernel_fc2_reduce` (IKR) is enabled and a rank receives zero tokens for a launch: the reduce path spins waiting for FC2 tiles that will never be produced, hanging the fleet. Found while integrating moe_ep with SGLang, where empty-rank launches occur routinely under real routing distributions. Also guards the autotuner so a tuned `token_back_mode` can no longer conflict with the IKR setting. Standalone repro scripts and multirank regression tests (with routing jitter to reproduce the pytest-vs-script scheduling gap) are included for both dtypes. ## Directories affected - `flashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/` — `mxfp8.py`, `nvfp4.py` (functional fixes in the shim layer; the vendored `src/` tree is untouched) - `tests/moe_ep/` — MXFP8 and NVFP4 mega multirank regression tests - `tests/` — standalone `repro_ikr_zero_token_idle{,_nvfp4}.py` artifacts 6 files changed, +824 / −7. ## Changes - `shim/mxfp8.py`, `shim/nvfp4.py`: on zero-token launches the IKR path is bypassed so the kernel completes and the combine step sees an empty contribution instead of spinning; the tuned-knob resolution no longer lets a cached `token_back_mode` enable a reduce mode that conflicts with the active IKR configuration. - `tests/moe_ep/test_moe_ep_mxfp8_cutedsl_mega_multirank.py`, `test_moe_ep_nvfp4_cutedsl_mega_multirank.py`: regression cases that drive a rank to zero tokens and assert completion; routing jitter added because the deterministic pytest distribution masked the livelock that the standalone scripts reproduced. - `tests/repro_ikr_zero_token_idle.py`, `tests/repro_ikr_zero_token_idle_nvfp4.py`: self-contained repro artifacts documenting the failure mode outside pytest. ## Testing Reproduced and verified fixed on 8x B200 via the standalone repros and the new multirank tests for both MXFP8 and NVFP4. The zero-token case livelocks deterministically before the fix and completes after. ## Notes for reviewers - The fix lives entirely in the shim layer, not in the vendored kernel drop, so no provenance update is needed. - Touches the same shim files area as the pending BF16 drop PR (#4386) only at the directory level; no file overlap — whichever lands second should still re-run the mega multirank suites. AI-assisted (Claude Code). 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **Bug Fixes** - Improved distributed Mixture-of-Experts processing when some devices receive zero tokens. - Prevented potential stalls or livelocks during mixed zero-token and real-token workloads. - Ensured all devices participate consistently in dispatch, reduction, cleanup, and synchronization. - Resolved conflicting token-back configuration when in-kernel reduction is enabled. - **Tests** - Added multi-device regression coverage for MXFP8 and NVFP4 zero-token scenarios. - Added standalone diagnostic scripts with progress monitoring and stall detection. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Md Saidul Hoque Anik <mhoqueanik@login-preos01.a51.clusters.nvidia.com> Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
3 participants
feat(moe_ep): SM100 BF16 CuTeDSL MegaMoE kernel
Summary
Adds an unquantized BF16 MegaMoE backend to
flashinfer.moe_ep— bf16 weights, bf16 activations, bf16 combine, no quantization anywhere in the pipeline — as a fused mega kernel (dispatch + grouped GEMM + combine in one launch) for Blackwell SM100. Port of draft PR #4120 (BF16 MegaMOE integration) onto the restructured taxonomy layout.Two roles:
MoEEpLayermega path as the fp8/fp4 kernels.What's included
kernel_src/cutedsl_megamoe/src/moe_bf16_glu/+shim/bf16.py, plus the bf16-enabling generalizations of the shared mxfp8/nvfp4 kernel sources (epiloguefc1_outputwidth guards,epi_flag_batchas a(fc1, fc2)pair,TopkReducesm_archparameter, iket ranges).backends/mega/kernel/sm100/bf16_bf16_bf16_cutedsl/with taxonomy configSm100_Bf16_Bf16_Bf16_Cutedsl_MegaMoeConfig, kernel namesm100_bf16_bf16_bf16_cutedsl, and deprecated aliasbf16_cutedsl.run_tests.sh.benchmarks/bench_bf16_cutedsl_megamoe.py.Usage
There is no pre-quantized activation path (
MegaConfig.quantize_input=Trueis required); input staging is a plain bf16 copy into the symmetric buffer.Performance
8x B200, EP8 (DP8/EP8/TP1), model-shape sweep
Standalone microbenchmark (
moe_ep_benchmark, jobs 2384005-2384012, 2026-08-10),e2e_pipelinedp50 µs, warmup 20 / iters 50, staging and weight preprocessing excluded from the timed region. cutlass-dsl 4.6.1. bf16 is the baseline; brackets are the quantized kernels' speedup vs bf16.deepseek_v4_flash (hidden 4096, inter 2048, 256 experts, top-6)
deepseek_v4_pro (hidden 7168, inter 3072, 384 experts, top-6)
deepseek_v3 (hidden 7168, inter 2048, 256 experts, top-8)
kimi_k2_6 (hidden 7168, inter 2048, 384 experts, top-8)
qwen3_5_397b (hidden 4096, inter 1024, 512 experts, top-10)
2x B200, EP2 (DP2/EP2/TP1) reference (from @djns99)
Same harness, deepseek_v3 geometry (256 experts, top-8, hidden 7168, inter 2048).
e2e_pipelinedp50 µs, speedup vs bf16 in brackets:TOKENS=8 single-point,
MEGA_TIMING=kernel(tester-parity bare launch): bf16 1279.0 µs, mxfp8 533.5 µs, nvfp4 289.8 µs.Absolute latencies do not transfer across EP sizes (per-rank weight bytes scale with
num_experts/world_size, and the small-batch mega kernels are weight-bandwidth bound — EP2 holds 4x the local experts of EP8). The quantities that reproduce across the EP2 and EP8 runs are the speedup ratios (mxfp8/bf16 ~2.0-2.5x, decaying toward ~2.0x at large batch) and the accuracy losses.Accuracy (% rel-L2 vs fp32 dense MoE reference, all-rank)
Flat across tokens/rank and identical on EP2 and EP8:
Constraints
Testing
run_tests.sh(bf16 oracle + config + mega multirank).Deviations from PR #4120 (reviewed intentionally)
moe_nvfp4_swapab(runner_common/mega_runner) instead of the PR'scute_ref_ops.pyrewrite; bf16 does not usecute_ref_ops, so the file is not imported.requirements.txtfloors (cutlass-dsl>=4.5.0for the 4.5.2 WAR chain,tvm-ffi>=0.1.6); the PR bumped both.strict=False, exception chaining) and fixed two latent error-path bugs (valid_ab_tupleNameError, undefinedtestingmodule inmega_reference_bf16).Summary by CodeRabbit
New Features
Compatibility
Documentation
Tests