Conversation
djns99
requested review from
Anerudhan,
IwakuraRein,
aleozlx,
bkryu,
cyx-6,
dhiraj113,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
July 23, 2026 23:59
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Contributor
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughAdds a BF16 CuTeDSL MegaMoE backend with fused Blackwell kernels, distributed token communication, weight preprocessing, runtime integration, public exports, references, benchmarks, functional runners, and configuration and multirank tests. ChangesBF16 CuTeDSL MegaMoE
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related issues
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 9
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/common/moe_utils.py (1)
1-1: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winExtraneous
fprefix on string literals without placeholders (Ruff F541). Same fix pattern (drop thefprefix) applies at all sites below.
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/common/moe_utils.py#L113-160: infmin/fmax, drop thefprefix fromptx_instr = f"min.NaN.f32 $0, $1, $2;"/f"min.f32 $0, $1, $2;"/f"max.NaN.f32 $0, $1, $2;"/f"max.f32 $0, $1, $2;"(lines 122, 124, 147, 149) and from the"=f,f,f"constraint-string arguments (lines 130, 155).flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/common/host_utils.py#L155-155: changeprint(f"\n")toprint("\n").🤖 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/moe_utils.py` at line 1, Remove unnecessary f-string prefixes from the literal PTX instruction and constraint strings in fmin and fmax within moe_utils.py, and change print(f"\n") to print("\n") in host_utils.py. Preserve the existing string contents and behavior.Source: Linters/SAST tools
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_mxfp8_glu/kernel_mxfp8_glu_fc12.py (1)
100-113: 🎯 Functional Correctness | 🔴 Critical | ⚡ Quick winNameError masks the intended
ValueErroron invalidab_dtype/sf_vec_sizepairing.
valid_ab_tupleis never defined — the validated tuple is bound tovalid_ab(line 102). Any caller that hits this branch (invalidab_dtypefor the givensf_vec_size) gets aNameError: name 'valid_ab_tuple' is not definedinstead of the intended descriptiveValueError.🐛 Proposed fix
raise ValueError( f"ab_dtype={ab_dtype.__name__} is not valid for " f"sf_vec_size={sf_vec_size}. " - f"Expected one of: {[t.__name__ for t in valid_ab_tuple]}." + f"Expected one of: {[t.__name__ for t in valid_ab]}." )🤖 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_mxfp8_glu/kernel_mxfp8_glu_fc12.py` around lines 100 - 113, In the validation branch of the constructor’s ab_dtype/sf_vec_size pairing check, replace the undefined valid_ab_tuple reference in the ValueError message with the already-bound valid_ab collection. Preserve the existing validation logic and ensure invalid pairings raise the intended descriptive ValueError.Source: Linters/SAST tools
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/epilogue.py (1)
3059-3078: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winSemicolon-joined statements trigger Ruff E702. Lines 3059 and 3074-3078 pack two calls per line (
second_t.bot.r1_load(); second_t.top.r1_load()). See consolidated note (also confirms whethersrc/is in the lint scope).flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/src/ptx_helpers.py (1)
511-519: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winDocstring says "system scope" but the fence is now explicitly device-scoped.
_fence_rel_gpuwas changed to passsyncscope="device", but its docstring still says "acquire-release semantics at system scope." Given the function name and the explicit scope argument, "device" is almost certainly correct — please update the docstring so callers don't assume cross-device (system-scope) visibility from this fence.📝 Suggested fix
""" - Fence operation with acquire-release semantics at system scope. + Fence operation with acquire-release semantics at device (GPU) scope.🤖 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/src/ptx_helpers.py` around lines 511 - 519, Update the docstring for _fence_rel_gpu to describe release semantics at device scope, matching the llvm.fence call’s syncscope="device" argument; remove the incorrect system-scope and acquire-release wording.
🧹 Nitpick comments (3)
flashinfer/moe_ep/backends/mega/kernel/__init__.py (1)
5-5: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSort
__all__to satisfy RUF022.Proposed fix
-__all__ = ["deep_gemm_mega", "bf16_cutedsl", "mxfp8_cutedsl", "nvfp4_cutedsl"] +__all__ = ["bf16_cutedsl", "deep_gemm_mega", "mxfp8_cutedsl", "nvfp4_cutedsl"]🤖 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/__init__.py` at line 5, Sort the entries in __all__ alphabetically in the module so the export list satisfies RUF022, preserving all existing exported names.Source: Linters/SAST tools
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_bf16_glu/kernel_bf16_glu_fc12.py (1)
465-507: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove the leftover debug
Lines 498–505 emit an unconditional
_compute_stagesis annotated-> Tuple[int, int, int]but returns a 4-tuple(num_acc_stage, num_a_stage, num_b_stage, num_sched_stages).♻️ Suggested changes
- ) -> Tuple[int, int, int]: + ) -> Tuple[int, int, int, int]:- print( - f"[fc12 stages] num_ab_stage={num_a_stage, num_b_stage} " - f"num_acc_stage={num_acc_stage} " - f"misc_budget={self._smem_misc_budget_bytes()} " - f"c_bytes_total={c_bytes_total} " - f"smem_cap={smem_capacity} " - f"unused_smem={unused_smem}" - )🤖 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 465 - 507, Remove the unconditional debug print from _compute_stages to avoid compile/build log spam, and update its return annotation from a three-element tuple to a four-element integer tuple matching the returned stage values.flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/simulate_fc1_fc2_sched.py (1)
33-34: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider
itertools.pairwisefor the non-decreasing check.Ruff flags this
zip()for missingstrict=(B905) and suggestsitertools.pairwise(RUF007). Behavior is correct as-is (the two iterables are intentionally different lengths), but switching removes the lint noise and self-documents intent.♻️ Suggested refactor
- if any(b < a for a, b in zip(offsets, offsets[1:])): + if any(b < a for a, b in itertools.pairwise(offsets)): raise ValueError("offsets must be non-decreasing (it's a cumsum)")(requires
from itertools import pairwiseorimport itertools)🤖 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/simulate_fc1_fc2_sched.py` around lines 33 - 34, Update the non-decreasing validation near the offsets check to use itertools.pairwise instead of zip(offsets, offsets[1:]), adding the required import and preserving the existing ValueError condition and message.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@flashinfer/moe_ep/backends/mega/kernel/bf16_cutedsl/backend.py`:
- Around line 99-102: Update the configuration construction in the
Bf16CutedslMegaMoe backend to merge Bf16CutedslMegaMoeConfig.load_balance_mode
into the knobs passed to get_symm_buffer_for_bf16_mega_moe(), while preserving
existing user-provided knob values. Ensure the same merged load_balance_mode is
retained through the auto-tune path so selecting "atomic_counter" reaches the
frontend configuration and scheduler.
In `@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/autotune.py`:
- Around line 336-343: Update the BF16 autotuning wrapper around autotune_knobs
to validate caller-supplied candidates with is_valid_bf16(), while retaining
bf16_candidates() for the default path. Pass only valid BF16 candidates to
autotune_knobs and raise an error when filtering leaves no candidates.
In `@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/SKILL.md`:
- Around line 63-64: Update the parenthetical in the replacement step to say
“copy only these five dirs,” matching the five kernel packages and directories
enumerated by the surrounding instructions.
In
`@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_bf16_glu/mega_reference_bf16.py`:
- Around line 1244-1256: Add the CuTe DSL testing import required by the
validation code, binding it as testing before the function that raises
testing.CantImplementError and catches it in can_implement. Do not alter the
existing dtype validation behavior.
In
`@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/contract.py`:
- Line 85: Restore explicit strict=False on the zip() calls at
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/contract.py
lines 85-85, 290-290, and 356-356. Update the loops over coord_tuple/self.sizes,
domain.names/domain_coord, and self.table/other.table respectively; preserve the
post-loop length-mismatch handling for the table comparison.
In
`@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/epilogue_refactor.py`:
- Around line 886-893: Remove the semicolon-joined statements in
from_r1_perm_until_last_store by placing each self.bot and self.top operation on
its own line, preserving the existing call order and behavior.
In
`@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/mega_runner.py`:
- Around line 2041-2042: In the reverse-check skipped message near the
fc1_output peek handling, remove the unnecessary f-string prefixes from both
adjacent string literals because they contain no interpolation. Preserve the
existing concatenated message text and formatting.
- Line 305: Replace the lambda assigned to dtype_name with a local named
function that returns the final component of str(t), preserving the existing
behavior while satisfying Ruff E731.
In
`@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_fc12_common.py`:
- Line 124: Update the kind handling in the runner, including the interleave
selection, scale-vector sizing, compute_reference interleave branch, and
_partition_workspace. Handle nvfp4, mxfp8, bf16, fp8_e4m3, and fp8_e5m2
explicitly, using the correct interleave granularity and element width; in
particular, size BF16 workspace elements as 16 bits rather than falling through
to MXFP8’s 8-bit path. Preserve the existing behavior for NVFP4 and MXFP8 while
ensuring dtype-aware workspace reshaping matches each kind.
---
Outside diff comments:
In `@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/common/moe_utils.py`:
- Line 1: Remove unnecessary f-string prefixes from the literal PTX instruction
and constraint strings in fmin and fmax within moe_utils.py, and change
print(f"\n") to print("\n") in host_utils.py. Preserve the existing string
contents and behavior.
In
`@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_mxfp8_glu/kernel_mxfp8_glu_fc12.py`:
- Around line 100-113: In the validation branch of the constructor’s
ab_dtype/sf_vec_size pairing check, replace the undefined valid_ab_tuple
reference in the ValueError message with the already-bound valid_ab collection.
Preserve the existing validation logic and ensure invalid pairings raise the
intended descriptive ValueError.
In `@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/src/ptx_helpers.py`:
- Around line 511-519: Update the docstring for _fence_rel_gpu to describe
release semantics at device scope, matching the llvm.fence call’s
syncscope="device" argument; remove the incorrect system-scope and
acquire-release wording.
---
Nitpick comments:
In `@flashinfer/moe_ep/backends/mega/kernel/__init__.py`:
- Line 5: Sort the entries in __all__ alphabetically in the module so the export
list satisfies RUF022, preserving all existing exported names.
In
`@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_bf16_glu/kernel_bf16_glu_fc12.py`:
- Around line 465-507: Remove the unconditional debug print from _compute_stages
to avoid compile/build log spam, and update its return annotation from a
three-element tuple to a four-element integer tuple matching the returned stage
values.
In
`@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/simulate_fc1_fc2_sched.py`:
- Around line 33-34: Update the non-decreasing validation near the offsets check
to use itertools.pairwise instead of zip(offsets, offsets[1:]), adding the
required import and preserving the existing ValueError condition and message.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 8b638279-b541-45bf-9753-8569ea1edf1b
📥 Commits
Reviewing files that changed from the base of the PR and between 766f94b and 313489ce8b13e3643bbb15ec697be63eb410371a.
📒 Files selected for processing (79)
benchmarks/bench_bf16_cutedsl_megamoe.pydocs/design_docs/moe_ep_architecture.mddocs/design_docs/moe_ep_runbook.mdflashinfer/moe_ep/__init__.pyflashinfer/moe_ep/backends/mega/kernel/__init__.pyflashinfer/moe_ep/backends/mega/kernel/bf16_cutedsl/__init__.pyflashinfer/moe_ep/backends/mega/kernel/bf16_cutedsl/backend.pyflashinfer/moe_ep/backends/mega/kernel/bf16_cutedsl/config.pyflashinfer/moe_ep/backends/mega/kernel/bf16_cutedsl/staging.pyflashinfer/moe_ep/backends/mega/kernel/bf16_cutedsl/weights.pyflashinfer/moe_ep/core/runtime/__init__.pyflashinfer/moe_ep/core/runtime/bootstrap.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/ACKNOWLEDGEMENT.mdflashinfer/moe_ep/kernel_src/cutedsl_megamoe/SKILL.mdflashinfer/moe_ep/kernel_src/cutedsl_megamoe/__init__.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/__init__.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/autotune.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/bf16.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/kernel_helpers.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/tuner.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/common/__init__.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/common/host_utils.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/common/megamoe_constants.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/common/moe_utils.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_bf16_glu/__init__.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_bf16_glu/custom_ext_bf16.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_bf16_glu/epilogue_bf16.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_bf16_glu/kernel_bf16_glu_fc12.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_bf16_glu/mega_reference_bf16.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_bf16_glu/mega_runner.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_bf16_glu/megamoe_kernel_bf16.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_bf16_glu/run_functional_tests.shflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_bf16_glu/run_mega_tests.shflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_bf16_glu/runner_common.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_bf16_glu/runner_fc12.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_mxfp8_glu/epilogue_mxfp8.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_mxfp8_glu/kernel_mxfp8_glu_fc12.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_mxfp8_glu/mega_reference_mxfp8.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_mxfp8_glu/mega_runner.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_mxfp8_glu/megamoe_kernel_mxfp8.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_mxfp8_glu/runner_common.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_mxfp8_glu/runner_fc12.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/benchmark_p2p.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/contract.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/custom_ext.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/cute_ref_ops.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/dynamic_mainloop.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/epilogue.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/epilogue_refactor.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/fc1_fc2_fuse_sched.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/kernel_fc12.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/mega_reference.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/mega_runner.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/megamoe_kernel.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/moe_persistent_scheduler.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/moe_utils.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/run_functional_tests.shflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/run_mega_tests.shflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_common.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_fc12.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_fc12_common.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/simulate_fc1_fc2_sched.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/topk_reduce.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/src/bootstrap.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/src/cleanup_kernel.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/src/config.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/src/flag_batch.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/src/grid_sync.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/src/iket_compat.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/src/inputs_process.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/src/ptx_helpers.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/src/reference.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/src/sym_buffer.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/src/token_comm.pyrequirements.txttests/moe_ep/run_tests.shtests/moe_ep/test_bf16_cutedsl_kernel_vs_reference.pytests/moe_ep/test_bf16_cutedsl_mega_config.pytests/moe_ep/test_moe_ep_bf16_cutedsl_mega_multirank.py
Comment on lines
+99
to
+102
| gate_up_clamp=_clamp(config), | ||
| in_kernel_fc2_reduce=config.in_kernel_fc2_reduce, | ||
| token_back_mode=config.token_back_mode, | ||
| knobs=config.knobs if isinstance(config.knobs, dict) else None, |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Forward load_balance_mode into the frontend configuration.
Bf16CutedslMegaMoeConfig.load_balance_mode is never passed to get_symm_buffer_for_bf16_mega_moe(), whose factory only receives it through knobs. Selecting "atomic_counter" therefore has no effect and runs the static scheduler instead. Merge this field into the workspace knobs and preserve it through the auto-tune 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/backends/mega/kernel/bf16_cutedsl/backend.py` around lines
99 - 102, Update the configuration construction in the Bf16CutedslMegaMoe
backend to merge Bf16CutedslMegaMoeConfig.load_balance_mode into the knobs
passed to get_symm_buffer_for_bf16_mega_moe(), while preserving existing
user-provided knob values. Ensure the same merged load_balance_mode is retained
through the auto-tune path so selecting "atomic_counter" reaches the frontend
configuration and scheduler.
Comment on lines
+336
to
+343
| return autotune_knobs( | ||
| symm_buffer._frontend, | ||
| launch, | ||
| bf16_candidates() if candidates is None else candidates, | ||
| label="bf16_mega", | ||
| warmup_iters=warmup_iters, | ||
| timed_iters=timed_iters, | ||
| ) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Reject non-BF16 candidate geometries before autotuning.
This wrapper forwards caller-supplied candidates directly to autotune_knobs() without is_valid_bf16(). Generic validity accepts geometries outside BF16’s fixed (256, 256, 64) / (2, 1, 1) contract, so a custom candidate can select an unsupported kernel configuration. Filter with is_valid_bf16() and raise if none remain.
🤖 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/autotune.py` around lines
336 - 343, Update the BF16 autotuning wrapper around autotune_knobs to validate
caller-supplied candidates with is_valid_bf16(), while retaining
bf16_candidates() for the default path. Pass only valid BF16 candidates to
autotune_knobs and raise an error when filtering leaves no candidates.
Comment on lines
+63
to
64
| 1. **Replace `src/` verbatim** with the drop's five kernel packages — no injected | ||
| files, no edits (the drop is a full repo; copy only these four dirs): |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Stale count: "four dirs" should be "five". The step now lists five packages and the rm -rf/cp -r commands enumerate five dirs (common,src,moe_bf16_glu,moe_mxfp8_glu,moe_nvfp4_swapab), but the parenthetical still says "copy only these four dirs".
📝 Proposed fix
-1. **Replace `src/` verbatim** with the drop's five kernel packages — no injected
- files, no edits (the drop is a full repo; copy only these four dirs):
+1. **Replace `src/` verbatim** with the drop's five kernel packages — no injected
+ files, no edits (the drop is a full repo; copy only these five dirs):📝 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.
Suggested change
| 1. **Replace `src/` verbatim** with the drop's five kernel packages — no injected | |
| files, no edits (the drop is a full repo; copy only these four dirs): | |
| 1. **Replace `src/` verbatim** with the drop's five kernel packages — no injected | |
| files, no edits (the drop is a full repo; copy only these five dirs): |
🤖 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 63 - 64,
Update the parenthetical in the replacement step to say “copy only these five
dirs,” matching the five kernel packages and directories enumerated by the
surrounding instructions.
Comment on lines
+1244
to
+1256
| valid_ab_dtypes = { | ||
| cutlass.Float16, | ||
| cutlass.BFloat16, | ||
| cutlass.TFloat32, | ||
| cutlass.Uint8, | ||
| cutlass.Int8, | ||
| cutlass.Float8E4M3FN, | ||
| cutlass.Float8E5M2, | ||
| } | ||
| if a_dtype not in valid_ab_dtypes or b_dtype not in valid_ab_dtypes: | ||
| raise testing.CantImplementError( | ||
| f"Unsupported AB dtype: {a_dtype} and {b_dtype}" | ||
| ) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
cutlass DSL CantImplementError testing module import path
💡 Result:
In the context of the NVIDIA CUTLASS CuTe DSL, CantImplementError is a custom exception defined within the cutlass.cute.testing module [1][2]. It is typically raised when the DSL encounters a configuration or hardware requirement—such as specific alignment needs, unsupported data types, or experimental features—that cannot be satisfied or has not yet been validated for the target environment [1]. If you are encountering this error, it generally indicates a deliberate "fail-fast" mechanism in the DSL's preflight checks rather than an issue with your Python module import paths [1]. Regarding import paths for testing and running CuTe DSL examples: 1. Module Structure: The CuTe DSL source code is organized within the cutlass Python package, specifically under cutlass/cute/ [2]. The testing utilities reside in cutlass.cute.testing [1][2]. 2. Handling Path Issues: If you are running examples or tests from the CUTLASS repository, import failures (e.g., ModuleNotFoundError) are often caused by improper sys.path configuration or stale installations [3][4]. The CUTLASS test suite (e.g., test/examples/CuTeDSL/conftest.py) manages paths by dynamically appending the appropriate example and utility directories to sys.path [4]. 3. Common Pitfalls: - Environment Pollution: Upgrading or downgrading the nvidia-cutlass-dsl package can leave behind stale files or missing metadata in your site-packages, leading to unexpected runtime errors or import issues [3]. If you suspect an installation conflict, manual cleanup of the site-packages/nvidia_cutlass_dsl/ directory may be required [3]. - Namespace Conflicts: Be cautious when placing custom code in directories that might conflict with namespace packages (like jax) [4]. It is standard practice to import the main cutlass package before modifying sys.path to ensure the environment is correctly initialized [4]. If you are receiving a CantImplementError, examine the error message carefully; it usually contains specific details about which alignment or architectural constraint was violated (e.g., "[alignment]...") [1]. If the error is unexpected, ensure you are using a version of the CUTLASS DSL compatible with your installed CUDA toolkit version [3].
Citations:
- 1: https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/cutlass/utils/blackwell_helpers.py
- 2: https://github.com/NVIDIA/cutlass/blob/d4e16f5d/python/CuTeDSL/cutlass/cute/__init__.py
- 3: [BUG] pip upgrade/downgrade breaks nvidia-cutlass-dsl NVIDIA/cutlass#3132
- 4: https://github.com/NVIDIA/cutlass/blob/0ce648f5/test/examples/CuTeDSL/conftest.py
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_bf16_glu/mega_reference_bf16.py"
echo "== file exists and line count =="
wc -l "$file"
echo
echo "== imports/top-level names related to cutlass/testing =="
sed -n '1,120p' "$file"
echo
echo "== references to unsupported AB dtype and CantImplementError =="
rg -n "valid_ab_dtypes|CantImplementError|check_supported_dtypes|check_mma_tiler_and_cluster_shape|check_tensor_alignment|check_epilog_store_option|can_implement|except testing" "$file"
echo
echo "== imports in module =="
sed -n '1,80p' "$file" | rg -n "^[[:space:]]*(import|from)\b|name" || trueRepository: flashinfer-ai/flashinfer
Length of output: 7622
Import CuTe DSL testing utilities before using testing.CantImplementError.
The validation helpers raise testing.CantImplementError and can_implement catches it, but this module has no testing import. Add the appropriate CuTe DSL import, e.g. import cutlass.cute.testing as testing, so these validation paths can fail fast instead of raising NameError.
🧰 Tools
🪛 Ruff (0.15.21)
[error] 1254-1254: Undefined name testing
(F821)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In
`@flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_bf16_glu/mega_reference_bf16.py`
around lines 1244 - 1256, Add the CuTe DSL testing import required by the
validation code, binding it as testing before the function that raises
testing.CantImplementError and catches it in can_implement. Do not alter the
existing dtype validation behavior.
| @@ -81,7 +82,7 @@ def linearize(self, coord: Sequence[int]) -> int: | |||
|
|
|||
| linear = 0 | |||
| stride = 1 | |||
| for axis, (idx, size) in enumerate(zip(coord_tuple, self.sizes, strict=False)): | |||
| for axis, (idx, size) in enumerate(zip(coord_tuple, self.sizes)): | |||
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Re-add explicit strict= to satisfy Ruff B905. This PR dropped strict=False from three zip() calls; Ruff B905 requires the parameter to be explicit. Restore explicit strict=False to keep behavior identical (do not silently switch to strict=True).
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/contract.py#L85-L85:zip(coord_tuple, self.sizes, strict=False)(lengths already validated equal).flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/contract.py#L290-L290:zip(domain.names, domain_coord, strict=False).flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/contract.py#L356-L356:zip(self.table, other.table, strict=False)—strict=Falseis required here because the length-mismatch case is handled after the loop (Line 361);strict=Truewould raise instead.
🧰 Tools
🪛 Ruff (0.15.21)
[warning] 85-85: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
📍 Affects 1 file
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/contract.py#L85-L85(this comment)flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/contract.py#L290-L290flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/contract.py#L356-L356
🤖 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/contract.py`
at line 85, Restore explicit strict=False on the zip() calls at
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/contract.py
lines 85-85, 290-290, and 356-356. Update the loops over coord_tuple/self.sizes,
domain.names/domain_coord, and self.table/other.table respectively; preserve the
post-loop length-mismatch handling for the table comparison.
Source: Linters/SAST tools
Comment on lines
+886
to
+893
| self.bot = TmemTranspose16x32(tmem_ptr, Region.Bottom, reg_tensor=reg_tensor_bot) | ||
|
|
||
| def from_r1_perm_until_last_store(self) -> Tuple[cute.Tensor, cute.Tensor]: | ||
| self.bot.r1_perm() | ||
| self.top.r1_perm() | ||
| self.bot.r1_store() | ||
| self.top.r1_store() | ||
| self.bot.r1_perm(); self.top.r1_perm() | ||
| self.bot.r1_store(); self.top.r1_store() | ||
|
|
||
| self.bot.r2_load() | ||
| self.top.r2_load() | ||
| self.top.r2_store() | ||
| self.bot.r2_store() | ||
| self.bot.r2_load(); self.top.r2_load() | ||
| self.top.r2_store(); self.bot.r2_store() |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Semicolon-joined statements trigger Ruff E702. Lines 889-893 pack two calls per line in from_r1_perm_until_last_store. See consolidated note.
🧰 Tools
🪛 Ruff (0.15.21)
[error] 889-889: Multiple statements on one line (semicolon)
(E702)
[error] 890-890: Multiple statements on one line (semicolon)
(E702)
[error] 892-892: Multiple statements on one line (semicolon)
(E702)
[error] 893-893: Multiple statements on one line (semicolon)
(E702)
🤖 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/epilogue_refactor.py`
around lines 886 - 893, Remove the semicolon-joined statements in
from_r1_perm_until_last_store by placing each self.bot and self.top operation on
its own line, preserving the existing call order and behavior.
| @@ -399,15 +297,19 @@ def __post_init__(self) -> None: | |||
| ) | |||
| if self.gate_up_clamp is not None and self.gate_up_clamp < 0.0: | |||
| raise ValueError( | |||
| f"gate_up_clamp must be None or non-negative, got {self.gate_up_clamp}." | |||
| f"gate_up_clamp must be None or non-negative, got " | |||
| f"{self.gate_up_clamp}." | |||
| ) | |||
|
|
|||
| def __str__(self) -> str: | |||
| dtype_name = lambda t: str(t).split(".")[-1] | |||
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Replace the lambda assignment with a local function.
Ruff E731 rejects assigning lambda to dtype_name.
Proposed fix
- dtype_name = lambda t: str(t).split(".")[-1]
+ def dtype_name(dtype: torch.dtype) -> str:
+ return str(dtype).split(".")[-1]📝 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.
Suggested change
| dtype_name = lambda t: str(t).split(".")[-1] | |
| def dtype_name(dtype: torch.dtype) -> str: | |
| return str(dtype).split(".")[-1] |
🧰 Tools
🪛 Ruff (0.15.21)
[error] 305-305: Do not assign a lambda expression, use a def
Rewrite dtype_name as a def
(E731)
🤖 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/mega_runner.py`
at line 305, Replace the lambda assigned to dtype_name with a local named
function that returns the final component of str(t), preserving the existing
behavior while satisfying Ruff E731.
Source: Linters/SAST tools
Comment on lines
+2041
to
+2042
| f"---- expected-vs-actual reverse check skipped " | ||
| f"(fc1_output peek unavailable) ----" |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Remove the unnecessary f-string prefixes.
These strings have no interpolation, triggering Ruff F541.
🧰 Tools
🪛 Ruff (0.15.21)
[error] 2041-2041: f-string without any placeholders
Remove extraneous f prefix
(F541)
[error] 2042-2042: f-string without any placeholders
Remove extraneous f prefix
(F541)
🤖 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/mega_runner.py`
around lines 2041 - 2042, In the reverse-check skipped message near the
fc1_output peek handling, remove the unnecessary f-string prefixes from both
adjacent string literals because they contain no interpolation. Preserve the
existing concatenated message text and formatting.
Source: Linters/SAST tools
| if self.kind == "nvfp4" | ||
| else Mxfp8Fc1GateUpInterleave | ||
| ) | ||
| _interleave = Nvfp4Fc1GateUpInterleave if self.kind == "nvfp4" else Mxfp8Fc1GateUpInterleave |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
repo_root="$(pwd)"
echo "repo_root=$repo_root"
echo "git status short:"
git status --short
echo
echo "Find relevant files:"
fd -a 'runner_fc12_common.py|runner_fc12.py|.*moe.*' . | sed 's#^\./##' | head -200
echo
echo "Outline target file sections:"
target="$(fd 'runner_fc12_common.py' . | head -1)"
if [ -n "$target" ]; then
echo "target=$target"
wc -l "$target"
echo "--- class-like/func names ---"
ast-grep outline "$target" --view expanded | sed -n '1,240p' || true
echo "--- relevant line ranges ---"
for range in '1,150' '200,250' '940,985' '1080,1125'; do
echo "### $range"
sed -n "${range},${range}p;${range},${range}p" "$target" | cat -n | sed 's/^//'
done
fi
echo
echo "Search for ProblemDesc/Fc12TesterBase and kind values:"
rg -n "ProblemDesc|Fc12TesterBase|kind=.*(nvfp4|mxfp8|fp8_e4m3|fp8_e5m2|bf16)|_interleave|Mxfp8Fc1GateUpInterleave|__elem_bits|elem_bits|fc1_output_byte_count|data_dtype" .Repository: flashinfer-ai/flashinfer
Length of output: 20501
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast, pathlib
for p in pathlib.Path('.').rglob('runner_fc12_common.py'):
print(f"FILE {p}")
s=p.read_text()
tree=ast.parse(s)
for node in ast.walk(tree):
if isinstance(node, ast.FunctionDef) and node.name == '_partition_workspace':
print(f'{p}:{node.lineno}')
for sub in ast.walk(node):
if isinstance(sub, ast.Assign):
for tgt in sub.targets:
if isinstance(tgt, ast.Name) and tgt.id in {"is_nvfp4","elem_bits"}:
print(f' assign line {sub.lineno}: {ast.get_source_segment(s, sub)}')
if isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name) and sub.func.id=='is_nvfp4':
print(f' call line {sub.lineno}: {ast.get_source_segment(s, sub)}')
print()
PYRepository: flashinfer-ai/flashinfer
Length of output: 464
🏁 Script executed:
#!/bin/bash
set -euo pipefail
target="./flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_fc12_common.py"
echo "### target sections"
for range in "80,140" "180,250" "454,600" "1040,1125"; do
echo "### $range"
sed -n "$range,$range p" "$target" | cat -n | sed -n '1,90p'
done
echo
echo "### bf16/fp8 kind construction references"
rg -n "kind=['\"](?:nf4|nvfp4|mxfp8|fp8_e4m3|fp8_e5m2|bf16)['\"]|ProblemDesc\\(|Fc12TesterBase\\(|class .*Fc12|class .*Tester|runner_fc12_common\\." flashinfer tests benchmarks -SRepository: flashinfer-ai/flashinfer
Length of output: 245
🏁 Script executed:
#!/bin/bash
set -euo pipefail
target="./flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/runner_fc12_common.py"
echo "### target sections"
for range in "80,140" "180,250" "454,600" "1040,1125"; do
start="${range%,*}"
end="${range#*,}"
echo "### $range"
sed -n "${start},${end}p" "$target" | cat -n | sed -n '1,90p'
done
echo
echo "### bf16/fp8 kind construction references"
rg -n "kind=['\"](?:nf4|nvfp4|mxfp8|fp8_e4m3|fp8_e5m2|bf16)['\"]|ProblemDesc\\(|Fc12TesterBase\\(|class .*Fc12|class .*Tester|runner_fc12_common\\." flashinfer tests benchmarks -SRepository: flashinfer-ai/flashinfer
Length of output: 19962
Handle the new bf16/fp8_e4m3/fp8_e5m2 kinds explicitly
The ProblemDesc literal now accepts BF16/FP8, but the interleave checks, scale vector sizing, compute_reference interleave branch, and especially _partition_workspace’s elem_bits = 4 if is_nvfp4 else 8 still fall back to MXFP8 behavior. This can pick the wrong interleave granularity and size the fc1_output workspace with 8-bit elements for BF16 data, causing a later dtype-aware reshape/view to fail or produce an incorrect layout.
🤖 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_fc12_common.py`
at line 124, Update the kind handling in the runner, including the interleave
selection, scale-vector sizing, compute_reference interleave branch, and
_partition_workspace. Handle nvfp4, mxfp8, bf16, fp8_e4m3, and fp8_e5m2
explicitly, using the correct interleave granularity and element width; in
particular, size BF16 workspace elements as 16 bits rather than falling through
to MXFP8’s 8-bit path. Preserve the existing behavior for NVFP4 and MXFP8 while
ensuring dtype-aware workspace reshaping matches each kind.
Collaborator
|
Thanks for the bf16 support @djns99 ! Could you add some perf numbers using this repo? The microbenchmark should be sufficient for now. https://github.com/[mhoqueanik/moe_ep_benchmark](https://github.com/mhoqueanik/moe_ep_benchmark) |
Contributor
Author
Here is a 2xB200 run Details
Setup
TOKENS=8 —
|
| Backend | p50 (µs) | min | max | tok/s | acc_loss % |
|---|---|---|---|---|---|
| bf16_cutedsl | 1322.1 | 1308.9 | 1470.1 | 12 102 | 0.286 |
| mxfp8_cutedsl | 557.3 | 547.0 | 705.4 | 28 708 | 6.371 |
| nvfp4_cutedsl | 326.2 | 322.5 | 581.8 | 49 044 | 23.324 |
TOKENS=8 — MEGA_TIMING=kernel (tester-parity bare launch)
| Backend | p50 (µs) | min | max | tok/s | acc_loss % |
|---|---|---|---|---|---|
| bf16_cutedsl | 1279.0 | 1262.6 | 1312.0 | 12 510 | 0.286 |
| mxfp8_cutedsl | 533.5 | 521.2 | 568.2 | 29 993 | 6.371 |
| nvfp4_cutedsl | 289.8 | 287.6 | 313.3 | 55 218 | 23.324 |
Sweep — MEGA_TIMING=e2e_pipelined
Latency p50 (µs)
| tok/rank | bf16_cutedsl | mxfp8_cutedsl | nvfp4_cutedsl |
|---|---|---|---|
| 8 | 1290.2 | 535.6 | 293.9 |
| 64 | 2675.3 | 1075.3 | 529.3 |
| 256 | 2749.5 | 1139.7 | 551.9 |
| 1024 | 2940.4 | 1285.2 | 656.3 |
| 4096 | 4493.4 | 2259.0 | 1139.1 |
Throughput (tok/s)
| tok/rank | bf16_cutedsl | mxfp8_cutedsl | nvfp4_cutedsl |
|---|---|---|---|
| 8 | 12 402 | 29 876 | 54 448 |
| 64 | 47 846 | 119 033 | 241 809 |
| 256 | 186 219 | 449 230 | 927 778 |
| 1024 | 696 515 | 1 593 506 | 3 120 505 |
| 4096 | 1 823 122 | 3 626 345 | 7 191 516 |
Speedup vs bf16 (latency)
| tok/rank | mxfp8 / bf16 | nvfp4 / bf16 |
|---|---|---|
| 8 | 2.41× | 4.39× |
| 64 | 2.49× | 5.05× |
| 256 | 2.41× | 4.98× |
| 1024 | 2.29× | 4.48× |
| 4096 | 1.99× | 3.94× |
Accuracy loss (% rel-L2 vs bf16 dense)
| tok/rank | bf16_cutedsl | mxfp8_cutedsl | nvfp4_cutedsl |
|---|---|---|---|
| 8 | 0.286 | 6.371 | 23.324 |
| 64 | 0.288 | 6.357 | 23.144 |
| 256 | 0.288 | 6.361 | 23.192 |
| 1024 | 0.288 | 6.358 | 23.177 |
| 4096 | 0.288 | 6.358 | 23.171 |
Full sweep detail (p50 / min / max µs)
| tok/rank | Backend | p50 | min | max | tok/s | acc_loss % |
|---|---|---|---|---|---|---|
| 8 | bf16_cutedsl | 1290.2 | 1273.9 | 1322.0 | 12 402 | 0.286 |
| 8 | mxfp8_cutedsl | 535.6 | 524.3 | 537.8 | 29 876 | 6.371 |
| 8 | nvfp4_cutedsl | 293.9 | 291.7 | 296.4 | 54 448 | 23.324 |
| 64 | bf16_cutedsl | 2675.3 | 2605.0 | 2693.1 | 47 846 | 0.288 |
| 64 | mxfp8_cutedsl | 1075.3 | 1056.9 | 1095.7 | 119 033 | 6.357 |
| 64 | nvfp4_cutedsl | 529.3 | 527.2 | 600.1 | 241 809 | 23.144 |
| 256 | bf16_cutedsl | 2749.5 | 2698.3 | 2767.9 | 186 219 | 0.288 |
| 256 | mxfp8_cutedsl | 1139.7 | 1134.6 | 1157.1 | 449 230 | 6.361 |
| 256 | nvfp4_cutedsl | 551.9 | 547.8 | 554.8 | 927 778 | 23.192 |
| 1024 | bf16_cutedsl | 2940.4 | 2825.3 | 2953.3 | 696 515 | 0.288 |
| 1024 | mxfp8_cutedsl | 1285.2 | 1263.6 | 1349.8 | 1 593 506 | 6.358 |
| 1024 | nvfp4_cutedsl | 656.3 | 650.3 | 664.7 | 3 120 505 | 23.177 |
| 4096 | bf16_cutedsl | 4493.4 | 4418.6 | 4544.4 | 1 823 122 | 0.288 |
| 4096 | mxfp8_cutedsl | 2259.0 | 2002.8 | 2615.2 | 3 626 345 | 6.358 |
| 4096 | nvfp4_cutedsl | 1139.1 | 1132.5 | 1181.6 | 7 191 516 | 23.171 |
djns99
marked this pull request as draft
July 27, 2026 00:48
4 tasks
djns99
force-pushed
the
djns99/bf16_mega_moe_integration
branch
from
August 10, 2026 03:54
313489c to
de2d893
Compare
mhoqueanik
added a commit
to mhoqueanik/flashinfer-moe_ep
that referenced
this pull request
Aug 13, 2026
…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>
mhoqueanik
added a commit
to mhoqueanik/flashinfer-moe_ep
that referenced
this pull request
Aug 15, 2026
…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>
mhoqueanik
added a commit
to mhoqueanik/flashinfer-moe_ep
that referenced
this pull request
Aug 18, 2026
…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>
mhoqueanik
added a commit
that referenced
this pull request
Aug 19, 2026
# feat(moe_ep): SM100 BF16 CuTeDSL MegaMoE kernel > Note: This PR is a rebased version of #4120 from @djns99 ## 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: 1. **A backend for serving bf16 MoE checkpoints** through the same `MoEEpLayer` mega path as the fp8/fp4 kernels. 2. **The accuracy baseline for the quantized backends**: it runs real bf16 model math (~0.29% rel-L2 vs the fp32 dense reference, i.e. bf16 rounding only), so quantized-kernel speedups and quality losses can both be quoted against it. ## What's included - **Kernel drop**: `kernel_src/cutedsl_megamoe/src/moe_bf16_glu/` + `shim/bf16.py`, plus the bf16-enabling generalizations of the shared mxfp8/nvfp4 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`. - **Benchmark**: `benchmarks/bench_bf16_cutedsl_megamoe.py`. ## Usage ```python from flashinfer.moe_ep import ( BootstrapConfig, FleetParams, MegaConfig, MoEEpLayer, MoEEpTensors, MoEWeightPack, Sm100_Bf16_Bf16_Bf16_Cutedsl_MegaMoeConfig, ) layer = MoEEpLayer( BootstrapConfig(world_size=world_size, rank=rank), FleetParams(num_experts=256, max_tokens_per_rank=tokens, token_hidden_size=7168), weights=MoEWeightPack(w13=w13_bf16, w2=w2_bf16), # plain bf16, no scales backend=MegaConfig( megakernel=Sm100_Bf16_Bf16_Bf16_Cutedsl_MegaMoeConfig(intermediate_size=2048, top_k=8), ), ) out = layer.forward(MoEEpTensors(hidden_states=x_bf16, topk_ids=ids, topk_weights=w)) ``` There is no pre-quantized activation path (`MegaConfig.quantize_input=True` is 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_pipelined` p50 µ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) | tok/rank | bf16 p50 (µs) | bf16 tok/s | mxfp8 (vs bf16) | deepgemm fp8/fp4 (vs bf16) | |---|---|---|---|---| | 8 | 336.9 | 189,942 | 171.0 (1.97x) | 108.6 (3.10x) | | 64 | 392.2 | 1,305,483 | 197.6 (1.98x) | 125.0 (3.14x) | | 512 | 433.7 | 9,444,752 | 261.2 (1.66x) | 155.1 (2.80x) | | 2048 | 822.2 | 19,926,056 | 435.2 (1.89x) | 380.9 (2.16x) | | 8192 | 2705.9 | 24,219,775 | 1357.3 (1.99x) | 1325.2 (2.04x) | **deepseek_v4_pro** (hidden 7168, inter 3072, 384 experts, top-6) | tok/rank | bf16 p50 (µs) | bf16 tok/s | mxfp8 (vs bf16) | deepgemm fp8/fp4 (vs bf16) | |---|---|---|---|---| | 8 | 1032.7 | 61,972 | 472.2 (2.19x) | 258.1 (4.00x) | | 64 | 1449.0 | 353,353 | 659.6 (2.20x) | 327.7 (4.42x) | | 512 | 1580.0 | 2,592,484 | 758.4 (2.08x) | 376.8 (4.19x) | | 2048 | 2434.1 | 6,731,085 | 1155.6 (2.11x) | 905.2 (2.69x) | | 8192 | 7117.0 | 9,208,426 | 3601.4 (1.98x) | 3193.4 (2.23x) | **deepseek_v3** (hidden 7168, inter 2048, 256 experts, top-8) | tok/rank | bf16 p50 (µs) | bf16 tok/s | mxfp8 (vs bf16) | deepgemm fp8/fp4 (vs bf16) | |---|---|---|---|---| | 8 | 633.9 | 100,959 | 275.5 (2.30x) | 171.5 (3.70x) | | 64 | 732.3 | 699,194 | 306.2 (2.39x) | 184.4 (3.97x) | | 512 | 805.9 | 5,082,693 | 445.4 (1.81x) | 278.6 (2.89x) | | 2048 | 1931.7 | 8,481,596 | 1000.4 (1.93x) | 807.4 (2.39x) | | 8192 | 6504.4 | 10,075,567 | 3288.5 (1.98x) | 3098.7 (2.10x) | **kimi_k2_6** (hidden 7168, inter 2048, 384 experts, top-8) | tok/rank | bf16 p50 (µs) | bf16 tok/s | mxfp8 (vs bf16) | deepgemm fp8/fp4 (vs bf16) | |---|---|---|---|---| | 8 | 832.5 | 76,880 | 363.6 (2.29x) | 207.0 (4.02x) | | 64 | 1047.6 | 488,759 | 431.2 (2.43x) | 252.9 (4.14x) | | 512 | 1145.9 | 3,574,470 | 529.3 (2.16x) | 315.4 (3.63x) | | 2048 | 2164.7 | 7,568,590 | 1099.2 (1.97x) | 819.3 (2.64x) | | 8192 | 6678.0 | 9,813,764 | 3441.2 (1.94x) | 3153.5 (2.12x) | **qwen3_5_397b** (hidden 4096, inter 1024, 512 experts, top-10) | tok/rank | bf16 p50 (µs) | bf16 tok/s | mxfp8 (vs bf16) | deepgemm fp8/fp4 (vs bf16) | |---|---|---|---|---| | 8 | 349.2 | 183,259 | 181.2 (1.93x) | 112.6 (3.10x) | | 64 | 414.7 | 1,234,520 | 212.0 (1.96x) | 131.2 (3.16x) | | 512 | 453.6 | 9,029,983 | 313.3 (1.45x) | 194.6 (2.33x) | | 2048 | 1023.9 | 16,001,750 | 525.1 (1.95x) | 546.9 (1.87x) | | 8192 | 3782.8 | 17,324,880 | 1776.6 (2.13x) | 2026.5 (1.87x) | ### 2x B200, EP2 (DP2/EP2/TP1) reference (from @djns99) Same harness, deepseek_v3 geometry (256 experts, top-8, hidden 7168, inter 2048). `e2e_pipelined` p50 µs, speedup vs bf16 in brackets: | tok/rank | bf16 | mxfp8 (vs bf16) | nvfp4 (vs bf16) | |---|---|---|---| | 8 | 1290.2 | 535.6 (2.41x) | 293.9 (4.39x) | | 64 | 2675.3 | 1075.3 (2.49x) | 529.3 (5.05x) | | 256 | 2749.5 | 1139.7 (2.41x) | 551.9 (4.98x) | | 1024 | 2940.4 | 1285.2 (2.29x) | 656.3 (4.48x) | | 4096 | 4493.4 | 2259.0 (1.99x) | 1139.1 (3.94x) | 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: | backend | acc loss % | |---|---| | bf16_cutedsl | **0.286-0.288** (bf16 rounding only — no quantization error) | | mxfp8_cutedsl | 6.36-6.37 | | nvfp4_cutedsl | 23.1-23.3 | ## Constraints - SM100 only. - Hidden size must be a multiple of 128 (shared mega fleet validation; e.g. gpt-oss-120b's hidden 2880 is rejected), intermediate size a multiple of 64. - Validated against nvidia-cutlass-dsl 4.6.1 (4.7.0 is known-broken for the mega multirank path; see the test-container recipe pin). ## Testing - PR #4120 validated as-is beforehand: 13 pytest + 40/40 functional/mega harness cases green on 4x GB200 with nvidia-cutlass-dsl 4.6.1. - Port re-wired into `run_tests.sh` (bf16 oracle + config + mega multirank). - Microbenchmark sweep above ran clean on 8x B200 (5 shapes x 5 token counts, plus deep_gemm/mxfp8 in-session controls agreeing with prior measurements within ~1%). ## Deviations from PR #4120 (reviewed intentionally) - PR tree was unformatted; content re-formatted with the repo-pinned ruff 0.12.8 so vendored-file diffs stay semantic. - Kept the 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 the 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`). <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added BF16 MegaMoE support for Blackwell GPUs. * Added Hopper SM90 FP8 MegaMoE support with per-tensor and blockwise scaling. * Added architecture-specific configuration, preprocessing, tuning, and benchmarking APIs. * **Compatibility** * Preserved deprecated backend aliases during the naming transition. * Added safeguards for incompatible GPU kernel environments. * **Documentation** * Expanded architecture, tuning, benchmarking, provenance, and operational guidance. * **Tests** * Added broad single-GPU, multi-rank, oracle, validation, and architecture-specific coverage. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
5 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
📌 Description
🔍 Related Issues
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit