MoE monokernel Bug fix, barrrier remove and kernel rewrite. - #4027
Merged
Merged
Conversation
yugong333
requested review from
IwakuraRein,
aleozlx,
bkryu,
cyx-6,
dhiraj113,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
qiching,
qsang-nv,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
July 17, 2026 16:44
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughMonoMoe is consolidated into a fixed E=256, N=512, K=2048, BS≤8 FP8 kernel. The API adds expert bias and routed scaling, while routing, projection synchronization, TMA/PTX primitives, validation, and documentation are updated. ChangesMonoMoe fixed-shape kernel
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant mono_moe
participant monomoe_topk
participant UpProjection
participant DownProjection
participant Output
mono_moe->>monomoe_topk: validate and launch fixed-shape kernel
monomoe_topk->>UpProjection: select experts and run up projection
UpProjection->>DownProjection: publish FP8 activations and scale readiness
DownProjection->>Output: accumulate results and publish down-ready state
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 14
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
csrc/fused_moe/monomoe/monomoe_wrapper.cuh (2)
138-153: 🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy liftReinitialize persistent handoff state when the scratchpad changes configuration.
The statics are per-
Dims. After A initializes a buffer, B zeroes and rewrites that same buffer using a differentMoEGemmSpeclayout; returning to A skips initialization because A still remembers the pointer. Stalelaunch_flip, scale sentinels, ordown_readycounters can cause premature reads or hangs.Track the active layout per scratchpad and reset on transitions, or make the persistent state layout configuration-independent.
🤖 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 `@csrc/fused_moe/monomoe/monomoe_wrapper.cuh` around lines 138 - 153, Update the first-use initialization block around _zeroed_ptr, _zeroed_size, and _zeroed_dev so it also tracks the active MoEGemmSpec layout for each scratchpad, not just pointer, size, and device. Trigger cudaMemsetAsync whenever the layout changes, including when returning from B’s layout to A’s, so persistent launch_flip, scale sentinels, and down_ready state are reset before reuse.
44-58: 🩺 Stability & Availability | 🔴 Critical | ⚡ Quick winValidate caller-provided output and scratchpad extents before launch.
outmay be smaller than[num_tokens, K], andscratchpadmay be smaller thansizeof(MoEGemmSpec<Dims>); both produce out-of-bounds device accesses.Proposed checks
const uint32_t num_tokens = activations_in.size(0); + TVM_FFI_ICHECK(activations_out.ndim() == 2 && + activations_out.size(0) == num_tokens && + activations_out.size(1) == Dims::HIDDEN_STATES) + << "activations_out must have shape [M, K]."; const size_t shmem_size = get_moe_shmem_size<Dims>(); const size_t scratchpad_size = static_cast<size_t>(scratchpad.numel()) * get_element_size(scratchpad); + TVM_FFI_ICHECK(scratchpad_size >= sizeof(MoEGemmSpec<Dims>)) + << "scratchpad is too small for the selected MonoMoe configuration.";Also applies to: 84-90
🤖 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 `@csrc/fused_moe/monomoe/monomoe_wrapper.cuh` around lines 44 - 58, Add extent validation in the wrapper’s input-checking flow for activations_out and scratchpad before launching the kernel. Require activations_out to provide storage for [num_tokens, K] elements and scratchpad to be at least sizeof(MoEGemmSpec<Dims>), using the existing shape/size validation helpers and preserving the current dtype checks.csrc/fused_moe/monomoe/src/moe.cuh (1)
64-76: 🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy liftDo not rely on timing to finish accumulator zeroing before Phase 4.
A fast block can reach
atomicAddwhile another block has not yet zeroed the samedown_partial_outregion, losing the partial sum. Use a stream-ordered, graph-capturable memset before launch or an explicit correctness-preserving handoff.🤖 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 `@csrc/fused_moe/monomoe/src/moe.cuh` around lines 64 - 76, Replace the in-kernel zeroing block for spec->down_partial_out with a stream-ordered, graph-capturable memset issued before the MoE kernel launch, or add an explicit synchronization/handoff that guarantees all accumulator elements are initialized before any Phase-4 atomicAdd. Remove the timing-based assumption and preserve the existing partial_n sizing across Dims::BS * Dims::HIDDEN_STATES.
🤖 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 `@csrc/fused_moe/monomoe/monomoe_binding.cu`:
- Around line 140-154: Resolve a single effective configuration for BS16 before
preprocessing, scratchpad sizing, and kernel dispatch. In
csrc/fused_moe/monomoe/monomoe_binding.cu:140-154, reject unavailable BS16
config IDs or expose the config-0 fallback ID; in
flashinfer/fused_moe/monomoe.py:655-669, use that effective BS-specific ID for
_config_uch, allocation, and dispatch; and in
csrc/fused_moe/monomoe/monomoe_binding.cu:171-196, size the actual Base16
fallback or reject the request before sizing.
In `@csrc/fused_moe/monomoe/shapes.json`:
- Around line 2-9: The registry’s max_token_expert_pairs constraint is too low
for the BS16 path with runtime top_k=8. Update the relevant shape entries,
including entries 19–21, so their constraint permits 128 token-expert pairs, or
implement a shape-specific cap while preserving 64 for shapes that do not
support the larger workload.
In `@csrc/fused_moe/monomoe/src/moe.cuh`:
- Around line 199-204: Update the sentinel-based BS16 documentation across
csrc/fused_moe/monomoe/src/moe.cuh (199-204, 287-290),
csrc/fused_moe/monomoe/src/moe_internal.h (250-254, 532-538, 598-601, 655-656),
and csrc/fused_moe/monomoe/src/moe_routing.cuh (473-477): describe +0.0f/nonzero
sentinel handoffs and reference moe.cuh, replace obsolete software-barrier and
Phase-3→4 claims with actual phase sequencing and per-block union-reuse
transitions, remove obsolete site-2/barrier-separation statements, and change
the routing pair bound from 64 to 128.
In `@csrc/fused_moe/monomoe/tools/_tune_helpers.py`:
- Around line 97-110: Update the output accumulation in the FP8 reference path
around the expert loop so `out` remains fp32 and each `block_wise_gemm` result
is accumulated without an intermediate bf16 cast. Convert to bf16 only once when
returning the final result, matching the kernel’s fp32 accumulation and
writeback behavior.
In `@csrc/fused_moe/monomoe/tools/enum_configs.py`:
- Around line 171-177: Replace the hardcoded SHAPES table in enum_configs.py
with entries loaded from shapes.json, resolving both canonical shape keys and
registered aliases through the registry. Update the --shape validation and
related lookup paths to use this single source of truth, while preserving
explicit --N/--K/--E/--bs enumeration for unlisted shapes.
- Around line 420-448: Update GDims and the verification probe to instantiate
the enumerated expert count and UP_COL_HALVES dimensions instead of hard-coding
NUM_EXPERTS to 256 or omitting UP_COL_HALVES. Thread both E and up_col_halves
through GDims and row_line, and emit them as c["E"] and c["up_col_halves"] so
--verify evaluates the same Dims type as each candidate.
In `@csrc/fused_moe/monomoe/tools/gen_shapes.py`:
- Around line 78-94: Replace every assert-based validation in the shapes loop
with explicit conditional checks that raise ValueError, preserving each existing
condition and descriptive message. Update the checks for dimensions, top_k,
configs, config IDs, duplicate shape keys, and duplicate (E,N,K) tuples so
validation remains active under python -O.
In `@csrc/fused_moe/monomoe/tools/tune_monomoe.py`:
- Line 132: Update the initial make_weights call in the tuning flow to use the
first requested value from seeds instead of the hardcoded 42, and ensure
subsequent evaluation continues from seeds[1:] without skipping or duplicating
that first seed.
- Around line 128-154: Update the tuning setup around run and the reference
computation to use each registry entry’s routing metadata instead of forcing
softmax. Build routing inputs and invoke mono_moe with the entry’s scoring_func,
expert bias, routed scaling factor, and renormalization settings, ensuring the
reference and kernel use the same configured routing behavior.
In `@docs/design_docs/monomoe_kernel.md`:
- Line 187: Remove the obsolete site-#2/site-#3 grid-barrier descriptions from
the documentation, including the passages around the referenced locations. Align
the affected sections with the sentinel and arrival-counter handoff model
defined in the existing authoritative sections, without introducing grid-barrier
behavior.
- Around line 115-119: Update the down_partial_out initialization and Phase-4
pipeline described in the design document to establish an explicit cross-block
handoff: publish zero-fill completion before allowing any Phase-4 atomicAdd
producer to update the scratchpad. Remove reliance on relative phase duration,
co-residency, or launch timing, while preserving Phase 5’s complete output
writes and avoiding an additional output pre-zero pass.
In `@flashinfer/fused_moe/monomoe.py`:
- Around line 324-330: Update the tuning selection logic around tuned_ms and
best so that when no tuned batch size M' <= m exists, it does not fall back to
tuned_ms[0]. Respect MONOMOE_REQUIRE_TUNED by falling through to the next
candidate or untuned path, while preserving the existing config lookup for valid
tuned buckets.
In `@flashinfer/trace/templates/moe.py`:
- Line 2609: Update flashinfer/trace/templates/moe.py lines 2609-2609 so seq_len
permits 16, then regenerate the trace output. In
csrc/fused_moe/monomoe/src/moe_tma.h lines 140-163, describe the activation tile
as Dims::BS × 128 with a BS-dependent byte count; update
docs/design_docs/monomoe_kernel.md lines 528-543 to document 8/16-row activation
boxes; regenerate tests/trace/fi_trace_out/mono_moe_topk8_h2048_i512.json lines
12-15 so its axis matches the template.
In `@tests/moe/test_monomoe.py`:
- Around line 225-235: Isolate the configuration identity test and the
workspace-mutating coverage around _default_tuned_dir by patching it to
tmp_path. Clear environment-based configuration pins and tuned JSON sources
before invoking the default resolution, so the config_id == 0 comparison
deterministically uses config 0 without reading or modifying the user’s
workspace.
---
Outside diff comments:
In `@csrc/fused_moe/monomoe/monomoe_wrapper.cuh`:
- Around line 138-153: Update the first-use initialization block around
_zeroed_ptr, _zeroed_size, and _zeroed_dev so it also tracks the active
MoEGemmSpec layout for each scratchpad, not just pointer, size, and device.
Trigger cudaMemsetAsync whenever the layout changes, including when returning
from B’s layout to A’s, so persistent launch_flip, scale sentinels, and
down_ready state are reset before reuse.
- Around line 44-58: Add extent validation in the wrapper’s input-checking flow
for activations_out and scratchpad before launching the kernel. Require
activations_out to provide storage for [num_tokens, K] elements and scratchpad
to be at least sizeof(MoEGemmSpec<Dims>), using the existing shape/size
validation helpers and preserving the current dtype checks.
In `@csrc/fused_moe/monomoe/src/moe.cuh`:
- Around line 64-76: Replace the in-kernel zeroing block for
spec->down_partial_out with a stream-ordered, graph-capturable memset issued
before the MoE kernel launch, or add an explicit synchronization/handoff that
guarantees all accumulator elements are initialized before any Phase-4
atomicAdd. Remove the timing-based assumption and preserve the existing
partial_n sizing across Dims::BS * Dims::HIDDEN_STATES.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: e575ab9a-a5ee-48e3-94ab-3facbbfc3388
⛔ Files ignored due to path filters (2)
csrc/fused_moe/monomoe/generated/configs_generated.incis excluded by!**/generated/**csrc/fused_moe/monomoe/generated/dims_generated.incis excluded by!**/generated/**
📒 Files selected for processing (27)
CLAUDE.mdcsrc/fused_moe/monomoe/monomoe_binding.cucsrc/fused_moe/monomoe/monomoe_wrapper.cuhcsrc/fused_moe/monomoe/shapes.jsoncsrc/fused_moe/monomoe/src/moe.cuhcsrc/fused_moe/monomoe/src/moe_down_projection.cuhcsrc/fused_moe/monomoe/src/moe_grid_barrier.hcsrc/fused_moe/monomoe/src/moe_interface.hcsrc/fused_moe/monomoe/src/moe_internal.hcsrc/fused_moe/monomoe/src/moe_routing.cuhcsrc/fused_moe/monomoe/src/moe_scale_inputs.cuhcsrc/fused_moe/monomoe/src/moe_tma.cucsrc/fused_moe/monomoe/src/moe_tma.hcsrc/fused_moe/monomoe/src/moe_up_projection.cuhcsrc/fused_moe/monomoe/src/ptx_utils.hcsrc/fused_moe/monomoe/tools/_tune_helpers.pycsrc/fused_moe/monomoe/tools/enum_configs.pycsrc/fused_moe/monomoe/tools/gen_shapes.pycsrc/fused_moe/monomoe/tools/tune_monomoe.pydocs/design_docs/monomoe_kernel.mdflashinfer/aot.pyflashinfer/fused_moe/monomoe.pyflashinfer/jit/monomoe.pyflashinfer/trace/templates/moe.pytests/moe/test_monomoe.pytests/trace/example.pytests/trace/fi_trace_out/mono_moe_topk8_h2048_i512.json
💤 Files with no reviewable changes (1)
- csrc/fused_moe/monomoe/src/moe_grid_barrier.h
Comment on lines
+324
to
+330
| pick = None | ||
| for tm in tuned_ms: | ||
| if tm <= m: | ||
| pick = tm | ||
| if pick is None: | ||
| pick = tuned_ms[0] | ||
| return int(best[str(pick)]["config_id"]) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Do not use a larger-M tuning record for a smaller untuned batch.
When no tuned M' <= m exists, selecting tuned_ms[0] contradicts the documented bucketing rule and bypasses MONOMOE_REQUIRE_TUNED. Fall through to the next candidate or the untuned path instead.
Proposed fix
if pick is None:
- pick = tuned_ms[0]
+ continue
return int(best[str(pick)]["config_id"])📝 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
| pick = None | |
| for tm in tuned_ms: | |
| if tm <= m: | |
| pick = tm | |
| if pick is None: | |
| pick = tuned_ms[0] | |
| return int(best[str(pick)]["config_id"]) | |
| pick = None | |
| for tm in tuned_ms: | |
| if tm <= m: | |
| pick = tm | |
| if pick is None: | |
| continue | |
| return int(best[str(pick)]["config_id"]) |
🤖 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/fused_moe/monomoe.py` around lines 324 - 330, Update the tuning
selection logic around tuned_ms and best so that when no tuned batch size M' <=
m exists, it does not fall back to tuned_ms[0]. Respect MONOMOE_REQUIRE_TUNED by
falling through to the next candidate or untuned path, while preserving the
existing config lookup for valid tuned buckets.
Comment on lines
+225
to
+235
| if config_id == 0: | ||
| # Explicit config 0 and the default resolution (also config 0 when | ||
| # untuned) run the SAME instantiated kernel (bare Base — the config-0 | ||
| # identity is structural in the binding). They are not bit-for-bit | ||
| # across launches because Phase-4's cross-block atomicAdd reduces in | ||
| # nondeterministic order (same ULP jitter as running any config twice), | ||
| # so the check is the same >0.9999 cosine used by the reuse test — a | ||
| # DIFFERENT kernel would diverge far more. | ||
| out0, _ = _run_and_compare(x, logits, weights, N, K, top_k=8, config_id=0) | ||
| out_default, _ = _run_and_compare(x, logits, weights, N, K, top_k=8) | ||
| assert H.cosine(out0, out_default) > 0.9999, ( |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Isolate configuration tests from the user's environment and workspace.
The default call is not guaranteed to resolve config 0 when env pins or tuned JSON exist, and Lines 345-364 mutate the real workspace. Patch _default_tuned_dir to tmp_path and clear configuration sources for the identity test to prevent flakes and concurrent-file clobbering.
Also applies to: 345-364
🤖 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/test_monomoe.py` around lines 225 - 235, Isolate the configuration
identity test and the workspace-mutating coverage around _default_tuned_dir by
patching it to tmp_path. Clear environment-based configuration pins and tuned
JSON sources before invoking the default resolution, so the config_id == 0
comparison deterministically uses config 0 without reading or modifying the
user’s workspace.
yugong333
force-pushed
the
megakernel_BS16
branch
from
July 21, 2026 20:55
88e857c to
336fa21
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/moe/test_monomoe.py (1)
104-115: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winKeep the reference in fp32 until the final cast.
_fp8_moe_run_expertsaccumulates the down-projection in fp32 and returns bf16 at the end; this test roundssiluand each expert contribution to bf16 earlier, which adds avoidable drift to the baseline.🤖 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/test_monomoe.py` around lines 104 - 115, Update the test reference loop around _block_wise_gemm to keep silu and the accumulated out tensor in fp32 throughout, removing intermediate bfloat16 casts from activation and expert contributions. Cast the completed output to bfloat16 only once at the end, matching _fp8_moe_run_experts.
🧹 Nitpick comments (1)
csrc/fused_moe/monomoe/src/ptx_utils.h (1)
21-24: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winFunction name
cp_async_cg_4contradicts the emitted.caqualifier.The helper emits
cp.async.ca.shared.global(cache-all, the only valid form for a 4-byte copy —.cgis 16-byte-only), but the name suggests.cg(L2-only) cache behavior. Rename tocp_async_ca_4to avoid implying a cache policy the instruction doesn't use.Proposed rename
-__device__ static inline void cp_async_cg_4(void* smem_dst, const void* gmem_src) { +__device__ static inline void cp_async_ca_4(void* smem_dst, const void* gmem_src) {🤖 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 `@csrc/fused_moe/monomoe/src/ptx_utils.h` around lines 21 - 24, Rename the helper function cp_async_cg_4 to cp_async_ca_4 so its name matches the emitted cp.async.ca.shared.global instruction. Update every reference to the helper consistently, preserving its 4-byte copy behavior.
🤖 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 `@csrc/fused_moe/monomoe/src/moe_interface.h`:
- Around line 96-102: Update the documentation comment for the W8A8 MoE kernel
to describe both supported scoring modes, softmax and sigmoid, instead of
claiming softmax-only scoring. Keep the existing statements about top-K routing,
renormalization, quantization, shape, and expert count unchanged.
In `@docs/design_docs/monomoe_kernel.md`:
- Around line 426-441: Correct the scratchpad documentation and API references
around MoEGemmSpec: describe omitted-scratchpad allocation through mono_moe and
alloc_scratchpad using get_scratchpad_size_bytes(), without requiring caller
allocation via get_moe_max_scratchpad_size(). Update the layout-invariant note
to identify monomoe_binding.cu as the location of the enforcing static_assert,
while preserving the requirement that new fields follow temp_fp8.
- Around line 3-9: Scope docs/design_docs/monomoe_kernel.md to the BS8
configuration, or clearly document the BS16 companion and dispatch boundary.
Update docs/design_docs/monomoe_kernel.md lines 115-118, 151-159, 202-205,
213-215, and 250-255 to remove or explicitly label alternate routing-window
counts, UP_GROUPS, down-grid/column-tile values, down K-tile counts, and Phase-5
writer counts; ensure all unlabeled scheduling geometry matches the same BS8
configuration table.
In `@tests/moe/test_monomoe.py`:
- Around line 126-128: Update the tests in tests/moe/test_monomoe.py to use the
declared SHAPES values instead of hard-coded E256/N512/K2048 dimensions, and
parameterize the active cases over the registered configurations. Extend the M
coverage to include at least M=9 and M=16 so both BS8 and BS16 companion paths
are exercised.
---
Outside diff comments:
In `@tests/moe/test_monomoe.py`:
- Around line 104-115: Update the test reference loop around _block_wise_gemm to
keep silu and the accumulated out tensor in fp32 throughout, removing
intermediate bfloat16 casts from activation and expert contributions. Cast the
completed output to bfloat16 only once at the end, matching
_fp8_moe_run_experts.
---
Nitpick comments:
In `@csrc/fused_moe/monomoe/src/ptx_utils.h`:
- Around line 21-24: Rename the helper function cp_async_cg_4 to cp_async_ca_4
so its name matches the emitted cp.async.ca.shared.global instruction. Update
every reference to the helper consistently, preserving its 4-byte copy behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 4009a2b4-6819-4804-aef2-99667b4d92ca
📥 Commits
Reviewing files that changed from the base of the PR and between 88e857c and 336fa21dc2ff12ae97c9d69985cdcd5839e2ba7e.
📒 Files selected for processing (18)
csrc/fused_moe/monomoe/monomoe_binding.cucsrc/fused_moe/monomoe/monomoe_wrapper.cuhcsrc/fused_moe/monomoe/src/moe.cuhcsrc/fused_moe/monomoe/src/moe_down_projection.cuhcsrc/fused_moe/monomoe/src/moe_grid_barrier.hcsrc/fused_moe/monomoe/src/moe_interface.hcsrc/fused_moe/monomoe/src/moe_internal.hcsrc/fused_moe/monomoe/src/moe_routing.cuhcsrc/fused_moe/monomoe/src/moe_scale_inputs.cuhcsrc/fused_moe/monomoe/src/moe_tma.cucsrc/fused_moe/monomoe/src/moe_tma.hcsrc/fused_moe/monomoe/src/moe_up_projection.cuhcsrc/fused_moe/monomoe/src/ptx_utils.hdocs/design_docs/monomoe_kernel.mdflashinfer/aot.pyflashinfer/fused_moe/monomoe.pyflashinfer/jit/monomoe.pytests/moe/test_monomoe.py
💤 Files with no reviewable changes (1)
- csrc/fused_moe/monomoe/src/moe_grid_barrier.h
🚧 Files skipped from review as they are similar to previous changes (2)
- flashinfer/aot.py
- csrc/fused_moe/monomoe/src/moe_internal.h
Update the block-FP8 top-K MoE monokernel (Hopper SM90a, single E256/N512/K2048 shape, BS <= 8) with the corrected pipeline, without the BS16 companion or the multi-shape/tunable-config machinery (those land in follow-up PRs). Core change — replace the software grid barriers with sentinel/flag cross-block handoffs: - Delete moe_grid_barrier.h and the grid/partial-barrier counter protocol. - Phase 3->4 handoff (site flashinfer-ai#2): each temp_act_scale cell is release-published after its fp8 payload segment and doubles as the readiness flag; the down-projection polls exactly the cells it consumes (per-expert granularity), so no barrier-counter rendezvous is needed. - Phase 4->5 handoff (site flashinfer-ai#3): each block bumps its column stripe's parity-selected arrival counter; only that stripe's Phase-5 writer polls it. - The launch-parity double-buffer makes the handoff state self-maintaining across CUDA-graph replays; the kernel stays capturable with a plain launch. - Co-residency invariant (GRID_SIZE <= SM count, one block/SM) is what keeps the flag spins deadlock-free. Also folds in the BS8 correctness/perf fixes from the updated kernel: routing-time precompute of the (expert, token) top-k rank (removing the per-expert rescan in both epilogues), the v2 gate/up pair-layout interleave (interleave_for_tma_wgmma_up), and reworked block-FP8 scale handling — with CuTe-delegated WGMMA control ops replaced by explicit inline PTX. Tests: tests/moe/test_monomoe.py gains scratchpad-reuse (no cross-launch contamination) and small-scale correctness coverage on top of the accuracy sweep (M in {1,2,8} x top_k {1,8}); all pass on H200. AI-assisted (Claude Code). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
yugong333
force-pushed
the
megakernel_BS16
branch
from
July 22, 2026 20:35
336fa21 to
4020a0f
Compare
yugong333
changed the title
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
docs/design_docs/monomoe_kernel.md (1)
34-39: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winClarify warp 8’s TMA and prefetch responsibilities.
This section names warp 8 lane 0 as the sole TMA launcher, while later sections assign warp 8 lanes 0–31/PF0 scale-prefetch work. State explicitly which threads issue tensor TMAs and mbarrier operations, and which perform prefetch work, so the role table matches the phase descriptions.
As per coding guidelines, documentation must stay synchronized with infrastructure and execution conventions.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/design_docs/monomoe_kernel.md` around lines 34 - 39, Clarify the role table around the TMA launcher and prefetch warps: state that only warp 8 lane 0 issues tensor TMAs and arms mbarriers, while warp 8’s remaining lanes perform PF0 scale-prefetch work as assigned in later phase descriptions. Ensure the wording distinguishes the single launcher thread from warp-level prefetch responsibilities and remains consistent with the phase descriptions.Source: Coding guidelines
🤖 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 `@csrc/fused_moe/monomoe/monomoe_wrapper.cuh`:
- Around line 77-84: Update the expert_bias handling in the monomoe wrapper to
validate that the tensor is one-dimensional and contains exactly NUM_EXPERTS
elements before obtaining expert_bias_ptr. Keep the existing device and float32
checks, and reject malformed shapes before the kernel launch.
- Around line 103-113: In the use_tma<Dims>::value path, validate that
num_tokens does not exceed the specialization capacity Dims::BS before creating
the activations descriptor. Reject oversized inputs using the module’s existing
validation or error-handling mechanism, and only call
create_activations_tma_desc for supported token counts.
In `@csrc/fused_moe/monomoe/src/moe_interface.h`:
- Around line 40-48: Add a materialized BS16 MonoMoe configuration and dispatch
path, extending the existing Dims/launcher export for the documented
E=256/N=512/K=2048 shape while preserving BS8 behavior. Update
csrc/fused_moe/monomoe/src/moe_interface.h lines 40-48 with the BS16 Dims/config
definition, monomoe_wrapper.cuh lines 171-178 to select and launch it for token
counts 9–16, and flashinfer/jit/monomoe.py lines 88-91 to generate/export the
matching BS16 variant; also update the related fused_moe caps and binding/device
checks so M through 16 is accepted and routed to the new launcher.
In `@docs/design_docs/monomoe_kernel.md`:
- Around line 354-356: Update the unrouted-expert explanation near the
“routed_count == 0” case to remove the claim that E4M3 lacks NaN encodings.
State the safety guarantee solely in terms of rank-filtered results for unrouted
experts never being accumulated or read.
---
Nitpick comments:
In `@docs/design_docs/monomoe_kernel.md`:
- Around line 34-39: Clarify the role table around the TMA launcher and prefetch
warps: state that only warp 8 lane 0 issues tensor TMAs and arms mbarriers,
while warp 8’s remaining lanes perform PF0 scale-prefetch work as assigned in
later phase descriptions. Ensure the wording distinguishes the single launcher
thread from warp-level prefetch responsibilities and remains consistent with the
phase descriptions.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: f79c3ba1-329e-47f1-ba89-e38494d90101
📥 Commits
Reviewing files that changed from the base of the PR and between 336fa21dc2ff12ae97c9d69985cdcd5839e2ba7e and 4020a0f.
📒 Files selected for processing (16)
csrc/fused_moe/monomoe/monomoe_binding.cucsrc/fused_moe/monomoe/monomoe_wrapper.cuhcsrc/fused_moe/monomoe/src/moe.cuhcsrc/fused_moe/monomoe/src/moe_down_projection.cuhcsrc/fused_moe/monomoe/src/moe_grid_barrier.hcsrc/fused_moe/monomoe/src/moe_interface.hcsrc/fused_moe/monomoe/src/moe_internal.hcsrc/fused_moe/monomoe/src/moe_routing.cuhcsrc/fused_moe/monomoe/src/moe_scale_inputs.cuhcsrc/fused_moe/monomoe/src/moe_up_projection.cuhcsrc/fused_moe/monomoe/src/ptx_utils.hdocs/design_docs/monomoe_kernel.mdflashinfer/aot.pyflashinfer/fused_moe/monomoe.pyflashinfer/jit/monomoe.pytests/moe/test_monomoe.py
💤 Files with no reviewable changes (1)
- csrc/fused_moe/monomoe/src/moe_grid_barrier.h
🚧 Files skipped from review as they are similar to previous changes (10)
- flashinfer/aot.py
- csrc/fused_moe/monomoe/monomoe_binding.cu
- csrc/fused_moe/monomoe/src/moe_scale_inputs.cuh
- tests/moe/test_monomoe.py
- csrc/fused_moe/monomoe/src/moe_routing.cuh
- flashinfer/fused_moe/monomoe.py
- csrc/fused_moe/monomoe/src/moe_down_projection.cuh
- csrc/fused_moe/monomoe/src/moe_internal.h
- csrc/fused_moe/monomoe/src/moe_up_projection.cuh
- csrc/fused_moe/monomoe/src/ptx_utils.h
added 2 commits
July 22, 2026 22:56
jhaotingc
reviewed
Jul 23, 2026
Contributor
There was a problem hiding this comment.
- Description says it adds cuda graph replay, adding a test would be nice
- test for added bias
- consider splitting this PR into (1) bug fix (sync rewrite) (2) new routing API
- Some local benchmark/acc test that shows this PR resolve acc issue would be more direct, e.g. has this accuracy test been taken care of by test_monomoe_accuracy? why is test_monomoe_accuracy only testing 1 case? maybe H200 latency before/after for M={1,2,4,8}, top_k={1,8}
Thank you!
Contributor
Author
Thank you @jhaotingc ! I have removed the additional routing method support and focus on the bug fix only now. Thank you for you comments! |
Member
|
/bot run tests/moe |
Collaborator
Collaborator
|
[SUCCESS] Pipeline #60899356: 18/18 executed test jobs passed |
aleozlx
added a commit
to aleozlx/flashinfer
that referenced
this pull request
Aug 7, 2026
The rewrite in flashinfer-ai#4027 changed tma_load_2d from .shared::cluster to .shared::cta, which requires PTX ISA 8.6 / CUDA 12.8+. Revert to .shared::cluster (dropping the erroneous .tile qualifier that neither CUTLASS nor the original intention uses), matching the fallback form in cute/arch/copy_sm90_tma.hpp. This restores CUDA 12.0+ compatibility for the sm_90a monomoe AOT build. Reverts the CUDA >= 12.8 workaround gate added in the previous commit. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This was referenced Aug 7, 2026
aleozlx
added a commit
that referenced
this pull request
Aug 7, 2026
## 📌 Description The rewrite in #4027 changed `tma_load_2d` in `csrc/fused_moe/monomoe/src/ptx_utils.h` from `.shared::cluster` to `.shared::cta`. The `.shared::cta` form requires PTX ISA 8.6 / CUDA 12.8+, breaking the `aot-build-import (cu126)` CI job with `ptxas: State space incorrect for instruction 'cp.async.bulk.tensor'`. Fix: revert to `.shared::cluster` and remove the erroneous `.tile` qualifier (`.tile` is only valid with `.shared::cta` and was never correct on the cluster form). This matches the SM90 fallback in CUTLASS `cute/arch/copy_sm90_tma.hpp`, which uses `shared::cluster` without `.tile` for all pre-SM120 targets. Restores CUDA 12.0+ compatibility. No performance impact: monomoe launches with cluster size 1, so `.shared::cluster` and `.shared::cta` are semantically identical for this kernel. The `.tile` qualifier is a PTX syntax disambiguator only; hardware behavior is unchanged. The first commit added a CUDA >= 12.8 version gate as an initial workaround; the second commit replaces it with this root-cause fix and reverts the gate. ## 🔍 Related Issues Reported in #4048 (comment): `ptxas monomoe_binding.ptx: State space incorrect for instruction 'cp.async.bulk.tensor'` on cu126 AOT build. ## 🚀 Pull Request Checklist - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. ## 🧪 Tests - [ ] `aot-build-import (x64, cu126)` and `aot-build-import (arm64, cu126)` pass in CI - [ ] `aot-build-import (x64/arm64, cu128/cu129/cu130)` continue to pass ## Reviewer Notes The PTX change is a one-liner: `shared::cta.global.tile` → `shared::cluster.global` (dropping `.tile`). CUTLASS's `cute/arch/copy_sm90_tma.hpp` uses exactly this form for the SM90 path. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved compatibility for tensor loading across supported PTX and CUDA versions. * Updated tensor addressing behavior for single-CTA kernel launches. * Improved support across compatible hardware and toolchains. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Find two bugs in previous design:
Fixes:
cross-block handoffs:
Summary by CodeRabbit
New Features
Bug Fixes
Documentation
Tests