feat(moe): sync SM12x W4A16 fused MoE family to b12x HEAD

[yichengj0]

📌 Description

Updates the SM120/SM121 W4A16 (NVFP4 weights, bf16 activations) fused-MoE family to current b12x upstream (cc9b476).

Changes:

• Cooperative persistent launches for the fused FC1/FC2 kernel, with occupancy-aware tile selection and launch bounds.
• A tensor-core decode path for small batches on the packed serving weights, folding the top-k sum into the FC2 store epilogue and removing a separate launch.
• Shape-stable route packing: the triton packing kernels specialize on a power-of-two capacity instead of the exact token count, so decode batch-size changes no longer recompile. The W4A16 workspaces are sized to the same capacity.
• A new fp4_e8m0_k32 (MXFP4 K/32) weight source format, including scale-tail handling for TP shards that are not 128-aligned.

Public API and behavior changes:

• b12x_fused_moe's source_format parameter accepts the new fp4_e8m0_k32 value.
• W4A16 workspaces are now sized to the power-of-two route capacity, slightly larger than before. This is what lets decode batch-size changes reuse one compiled route-packing kernel instead of recompiling.
• Kernel launches now dispatch through torch custom ops, registered at module import. This brings the family in line with the library-wide convention it predated.

Upstream features with no FlashInfer target checkpoint are not included:

• NF3 (3-bit normal-float) weights. Their kernel branches stay: they compile out and keep future sync diffs clean.
• The NVFP4+NF3 hybrid entry points.
• The native-ModelOpt small-batch micro kernel.
• The SiTU and swiglu-oai activations.

📊 Performance

Speedup over the previous kernels at the MoE shape of a Nemotron variant, with relu2 and silu activations, run under CUDA graphs, median of three runs. Each cell reads RTX 5080 / RTX Pro 6000 Server Edition / GB10:

Activation	m=1	m=4	m=16	m=64	m=2048
relu2	1.07x / 1.15x / 1.11x	1.03x / 1.07x / 1.04x	1.00x / 1.00x / 1.02x	0.99x / 1.00x / 1.03x	12.91x / 18.52x / 1.00x
silu	0.98x / 1.23x / 1.10x	0.99x / 1.05x / 1.02x	0.95x / 0.98x / 1.01x	0.95x / 0.98x / 1.00x	11.70x / 16.88x / 1.02x

The m=2048 cells fix a stall rather than measure a general kernel win: the fused kernel uses grid-wide barriers but was not launched cooperatively, so at large m CTAs could spin waiting for peers that were never co-scheduled. GB10 kept the whole grid resident and never hit the stall, so its m=2048 cells have no win to inherit.

🔍 Related Issues

#4223 (item 3).

🚀 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

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

🧪 Tests

• tests/moe/test_b12x_fused_moe.py (142 cases) passes on SM121, both on the default dispatch and with the W4A16 path forced.
• Added tests/moe/test_b12x_w4a16_route_pack.py (64 cases): route packing against a host reference, covering expert_map, invalid expert ids, preallocated buffers, and the workspace capacity contract.
• Weight preparation is bit-identical to the previous implementation across the supported source formats and activations.

Reviewer Notes

• Most of the diff is upstream code taken as is. FlashInfer-local changes are limited to the import remaps, the two support modules, the API-boundary checks, the workspace sizing in moe_dispatch, and the decode-path cap below.
• One tuning deviation from upstream: the tensor-core decode path is selected up to m = 4 instead of upstream's m = 8. The crossover to the route-packed GEMM varies by card, so the cap conservatively stops where the decode path wins or ties on every card measured.
• The fused kernel compiles at DSL optimizer level 3 instead of upstream's 2. Level 3 generates a faster mainloop on the DSL version pinned here, at the cost of a longer compile for this kernel.
• The kernels import activation metadata and a kernel compile cache from upstream's library. The two new modules, moe_w4a16_activations.py and moe_w4a16_compiler.py, are small local implementations of those two pieces.

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features

  • Expanded Blackwell W4A16 fused MoE support for NF3, NVFP4, ModelOpt, CompressedTensors, and FP4 E8M0 formats.
  • Added configurable activations, SwiGLU parameters, routing, calibration, and launch options.
  • Improved route packing for large workloads, dynamic capacities, expert remapping, and invalid route IDs.
  • Added optimized quantization, dequantization, conversion, and matrix-operation support.
  • Added compile caching to reduce repeated kernel compilation.

• Bug Fixes

  • Improved workspace sizing and safe handling of zero-valued scales.

• Tests

  • Added coverage for route grouping, padding, capacity fallbacks, large shapes, and expert mapping.

[coderabbitai]

📝 Walkthrough

Walkthrough

SM12x W4A16 gains activation normalization, compilation caching, FP4/NF3 support, additional weight formats, expanded kernel execution paths, and capacity-aware route packing with CUDA-gated coverage.

Changes

SM12x W4A16 implementation

Layer / File(s)	Summary
Activation and compilation contracts flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_activations.py, flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_compiler.py, flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_host.py	Adds activation normalization, SwiGLU validation, compile specifications, structural cache keys, thread-safe compilation caching, and shared activation validation.
Weight format preparation flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_prepare.py	Adds ModelOpt, CompressedTensors, and E8M0 K/32 preparation with layout normalization, reusable buffers, scale packing, and W13 rotation.
FP4, FP8, and PTX primitives flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_fp4_helpers.py	Adds quantization, dequantization, memory, reduction, conversion, MMA, asynchronous-copy, and fused activation helpers.
Kernel compilation and execution flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_kernel.py	Adds additional layouts, scales, routing modes, active-row scheduling, TC-decode, activation calibration, SwiGLU variants, and cached compilation.
Route packing and workspace capacity flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_route_pack.py, flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_host.py, flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py, tests/moe/*w4a16*	Adds power-of-two route capacity sizing, parallel prefix construction, live-count handling, workspace fallback behavior, and route-packing tests.

Estimated code review effort: 5 (Critical) | ~120 minutes

Possibly related issues

• [Tracking][SM12x] Port latest b12x MoE and GEMM kernels for NVFP4 model targets #4223 — Covers the SM12x W4A16 MoE and NVFP4 activation and routing work extended by this PR.

Possibly related PRs

• flashinfer-ai/flashinfer#3336 — Modifies the related SM12x W4A16 route-packing, workspace, preparation, and kernel paths.
• flashinfer-ai/flashinfer#4285 — Modifies related SM12x W4A16 activation, routing, and dispatch internals.

Suggested labels: op: gemm

Suggested reviewers: yzh119, nv-yunzheq, aleozlx

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 9.72% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the SM12x W4A16 fused-MoE synchronization as the primary change.
Description check	✅ Passed	The description covers the change, related issue, checklist, tests, performance, exclusions, and reviewer notes.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[bkryu]

/bot run tests/moe

[flashinfer-bot]

GitLab MR !1086 has been created, and the CI pipeline #60392848 is currently running. I'll report back once the pipeline job completes.

[coderabbitai]

Actionable comments posted: 3

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_kernel.py (1)

6640-6659: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win

Propagate the stream into route packing or forbid non-default streams.

pack_topk_routes_by_expert accepts stream, but forwards the pack into Triton kernels that do not take any stream argument. Passing a different current stream (the API is also invoked from run_w4a16_moe with the launch stream) won’t order the route pack relative to the fused MoE launches. Thread the stream through the packing implementation, or explicitly reject non-default streams instead of dropping the argument.

🤖 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/cute_dsl/blackwell_sm12x/moe_w4a16_kernel.py` around
lines 6640 - 6659, Fix pack_topk_routes_by_expert so its stream argument is not
silently discarded: thread stream through _pack_topk_routes_by_expert and the
underlying Triton launches, ensuring route packing is ordered on the requested
stream, or validate and reject non-default streams. Remove the unconditional del
stream while preserving existing validation and output behavior.

🧹 Nitpick comments (5)

flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_host.py (1)

163-171: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

route_pack_token_capacity takes a required topk it immediately discards.

It's now public (__all__), so the signature is a contract. Either give it a default (topk: int = 1) or drop the parameter unless a future topk-dependent capacity is planned.

🤖 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/cute_dsl/blackwell_sm12x/moe_w4a16_host.py` around lines
163 - 171, Update the public route_pack_token_capacity signature to make topk
optional with a default of 1, preserving the current capacity calculation and
compatibility with existing callers.

flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_route_pack.py (2)

307-334: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Collapse the redundant re-assignments.

max_packed_routes/max_route_blocks are assigned from the capacity values and then immediately re-clamped; fold the clamp into a single assignment.

♻️ Proposed cleanup

-    max_packed_routes = capacity_packed_routes
-    max_route_blocks = capacity_route_blocks
-    max_packed_routes = max(max_packed_routes, 1)
-    max_route_blocks = max(max_route_blocks, 1)
+    max_packed_routes = max(capacity_packed_routes, 1)
+    max_route_blocks = max(capacity_route_blocks, 1)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_route_pack.py` around
lines 307 - 334, In the capacity setup, simplify the assignments to
max_packed_routes and max_route_blocks by applying the minimum-value clamp
directly when assigning from capacity_packed_routes and capacity_route_blocks.
Remove the redundant intermediate assignments while preserving the existing
minimum of 1 behavior.

---

375-475: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low value

Consider reusing a workspace buffer for expert_counts.

The eager large-shape path allocates a fresh torch.zeros(num_experts) per call. It's caching-allocator cheap and the comment justifies why capture never reaches here, but the packing entry point already accepts preallocated workspaces — threading one more slice through would keep the prefill path allocation-free.

🤖 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/cute_dsl/blackwell_sm12x/moe_w4a16_route_pack.py` around
lines 375 - 475, Reuse a caller-provided workspace for expert_counts in the
eager large-shape path instead of allocating torch.zeros inside the packing
flow. Thread an appropriately sized expert-count slice through the packing entry
point and use it in _w4a16_route_count_kernel, preserving the existing
captured-path fallback and count initialization semantics.

flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_prepare.py (1)

95-146: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Deduplicate the two padding branches in _permute_packed_scales.

The group_size < size_k and single-scale branches differ only in the permutation vector; the padding/validation/trim logic is copied verbatim.

♻️ Suggested restructure

 def _permute_packed_scales(
     scales: torch.Tensor,
     *,
     size_k: int,
     size_n: int,
     group_size: int,
     output_size_n: int | None = None,
 ) -> torch.Tensor:
     scale_perm, scale_perm_single = _scale_perms()
-    if group_size < size_k and group_size != -1:
-        block = len(scale_perm)
-        ...
-    else:
-        block = len(scale_perm_single)
-        ...
+    perm = (
+        scale_perm if (group_size < size_k and group_size != -1) else scale_perm_single
+    )
+    block = len(perm)
+    if output_size_n is not None or int(size_n) % block != 0:
+        padded_n = (
+            ((int(size_n) + block - 1) // block) * block
+            if output_size_n is None
+            else int(output_size_n)
+        )
+        if padded_n < int(size_n) or padded_n % block != 0:
+            raise ValueError(
+                f"output_size_n must be a multiple of {block} and >= size_n"
+            )
+        padded = scales.new_zeros((int(scales.shape[0]), padded_n))
+        padded[:, : int(size_n)] = scales
+        rows = int(scales.shape[0])
+        scales = padded.reshape((rows, -1, block))[:, :, perm].reshape((rows, padded_n))
+        if output_size_n is None:
+            scales = scales[:, : int(size_n)]
+        return scales.contiguous()
+    scales = scales.reshape((-1, block))[:, perm]
     return scales.reshape((-1, size_n)).contiguous()

🤖 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/cute_dsl/blackwell_sm12x/moe_w4a16_prepare.py` around
lines 95 - 146, Deduplicate the repeated padding, validation, permutation, and
trimming logic in _permute_packed_scales by selecting the appropriate
permutation vector and block size first, then applying one shared transformation
path. Preserve the existing branch condition that chooses scale_perm versus
scale_perm_single and all output_size_n validation and return-shape behavior.

flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_kernel.py (1)

112-121: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoff

All nf3_2p1 / e4m3_k32 code is unreachable — consider dropping or gating it explicitly.

_WEIGHT_LAYOUTS is {"packed", "modelopt"} and _SCALE_FORMATS only maps e4m3_k16/e8m0_k32, and every entry point validates against them (W4A16GemmKernel.__init__ Line 568, W4A16FusedMoeKernel.__init__ Line 4134, compile_w4a16_fused_moe Line 5252, run_w4a16_moe Line 6707). So weight_layout == "nf3_2p1" can never hold, which makes a large amount of added code dead: the NF3 register table (Lines 195-197), B-stage geometry (Lines 744-754), _load_b_scale_registers NF3 arm (Lines 2522-2544), _scaled_dequant_b_fragment_nf3, the NF3 flat-span staging (Lines 3057-3086), NF3 fake buffers (Lines 5573-5584), and the prepared-tile check (Lines 7004-7016). The nf3_2p1 branch at Lines 577-583 is doubly unreachable since it demands scale_format == "e4m3_k32", which _normalize_scale_format rejects. Same for self.scale_k32's e4m3_k32 term (Line 654) and _scale_group_size (Line 5017).

This matches the PR statement that NF3 is not included, but as-is it is untestable dead weight. Either remove it from this PR or add nf3_2p1/e4m3_k32 to the allow-lists behind an explicit feature gate so the paths are reachable and covered.

🤖 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/cute_dsl/blackwell_sm12x/moe_w4a16_kernel.py` around
lines 112 - 121, The NF3 and e4m3_k32 implementation is unreachable because
validation excludes both formats. Either remove the related NF3/e4m3_k32 code
paths and supporting constants, including branches in W4A16GemmKernel,
_scale_group_size, staging, and prepared-tile handling, or explicitly
feature-gate and add both formats to the validation allow-lists in
W4A16GemmKernel.__init__, W4A16FusedMoeKernel.__init__, compile_w4a16_fused_moe,
and run_w4a16_moe so the existing paths become reachable and testable.

🤖 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/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_compiler.py`:
- Around line 94-104: Update the compilation-cache flow around _COMPILE_CACHE
and clear_compile_cache() to serialize in-flight compilation per cache_key and
track cache generation across clears. Ensure concurrent misses for the same key
share one compilation, and only insert a compiled result when it belongs to the
current generation, so a pre-clear compilation cannot repopulate the cache after
invalidation.

In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_fp4_helpers.py`:
- Around line 296-299: Discard the unused first return value from
pow2_ceil_ue8m0_torch in this grouped quantization path by binding it with the
project’s underscore convention, while retaining byte for
_ue8m0_output_scale_torch.

In `@tests/moe/test_b12x_w4a16_route_pack.py`:
- Around line 303-325: Add a CUDA/SM12x availability skip marker to
test_w4a16_workspace_route_buffers_cover_route_pack_capacity, matching the
module’s existing guards and using the appropriate flashinfer.utils architecture
check or backend capability marker. Keep the test body unchanged.

---

Outside diff comments:
In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_kernel.py`:
- Around line 6640-6659: Fix pack_topk_routes_by_expert so its stream argument
is not silently discarded: thread stream through _pack_topk_routes_by_expert and
the underlying Triton launches, ensuring route packing is ordered on the
requested stream, or validate and reject non-default streams. Remove the
unconditional del stream while preserving existing validation and output
behavior.

---

Nitpick comments:
In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_host.py`:
- Around line 163-171: Update the public route_pack_token_capacity signature to
make topk optional with a default of 1, preserving the current capacity
calculation and compatibility with existing callers.

In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_kernel.py`:
- Around line 112-121: The NF3 and e4m3_k32 implementation is unreachable
because validation excludes both formats. Either remove the related NF3/e4m3_k32
code paths and supporting constants, including branches in W4A16GemmKernel,
_scale_group_size, staging, and prepared-tile handling, or explicitly
feature-gate and add both formats to the validation allow-lists in
W4A16GemmKernel.__init__, W4A16FusedMoeKernel.__init__, compile_w4a16_fused_moe,
and run_w4a16_moe so the existing paths become reachable and testable.

In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_prepare.py`:
- Around line 95-146: Deduplicate the repeated padding, validation, permutation,
and trimming logic in _permute_packed_scales by selecting the appropriate
permutation vector and block size first, then applying one shared transformation
path. Preserve the existing branch condition that chooses scale_perm versus
scale_perm_single and all output_size_n validation and return-shape behavior.

In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_route_pack.py`:
- Around line 307-334: In the capacity setup, simplify the assignments to
max_packed_routes and max_route_blocks by applying the minimum-value clamp
directly when assigning from capacity_packed_routes and capacity_route_blocks.
Remove the redundant intermediate assignments while preserving the existing
minimum of 1 behavior.
- Around line 375-475: Reuse a caller-provided workspace for expert_counts in
the eager large-shape path instead of allocating torch.zeros inside the packing
flow. Thread an appropriately sized expert-count slice through the packing entry
point and use it in _w4a16_route_count_kernel, preserving the existing
captured-path fallback and count initialization semantics.

🪄 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: 8822f6f9-4333-47d2-be3e-39f99ae6e9b5

📥 Commits

Reviewing files that changed from the base of the PR and between 43f12df and da1edde.

📒 Files selected for processing (10)

• flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py
• flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_activations.py
• flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_compiler.py
• flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_fp4_helpers.py
• flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_host.py
• flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_kernel.py
• flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_prepare.py
• flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_route_pack.py
• tests/moe/test_b12x_fused_moe.py
• tests/moe/test_b12x_w4a16_route_pack.py

[coderabbitai]

🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift

Make cache lookup, compilation, and invalidation atomic per key.

Line 94 releases the lock before compilation, so concurrent misses compile the same specialization repeatedly. More importantly, clear_compile_cache() can clear the cache while this work is in progress, then Line 103 repopulates it with a pre-clear result. Track in-flight compilations and a cache generation (or equivalent) so one caller compiles per key and pre-clear results cannot be inserted afterward.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_compiler.py` around
lines 94 - 104, Update the compilation-cache flow around _COMPILE_CACHE and
clear_compile_cache() to serialize in-flight compilation per cache_key and track
cache generation across clears. Ensure concurrent misses for the same key share
one compilation, and only insert a compiled result when it belongs to the
current generation, so a pre-clear compilation cannot repopulate the cache after
invalidation.

[yichengj0]

Compiling the same kernel twice gives the same result, so a concurrent miss only costs a wasted compile. clear_compile_cache() is a test helper. Keeping the simple lock.

[coderabbitai]

@yichengj0, understood. Given duplicate compilation is deterministic and clear_compile_cache() is test-only, the current simple locking approach is appropriate here. Withdrawing this finding.

🐇

---

✏️ Learnings added

Learnt from: yichengj0
Repo: flashinfer-ai/flashinfer PR: 4255
File: flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_compiler.py:94-104
Timestamp: 2026-07-31T01:07:30.288Z
Learning: In `flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_compiler.py`, concurrent cache misses for the same SM12x W4A16 CuTe-DSL kernel specialization are intentionally allowed to compile redundantly because compilation is deterministic and the impact is only wasted compile work. `clear_compile_cache()` is a test helper, so cache-generation protection against in-flight pre-clear compilations is not required.

Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.

_{You are interacting with an AI system.}

[bkryu]

Thanks @yichengj0 left a few comments!

[bkryu]

/bot run tests/moe

[flashinfer-bot]

GitLab MR !1086 has been updated with latest changes, and the CI pipeline #60398942 is currently running. I'll report back once the pipeline job completes.

[flashinfer-bot]

[FAILED] Pipeline #60398942 — 6/18 executed test jobs passed

Compared with nightly #60101730.

Unit Tests

GPU	CUDA 12.9	CUDA 13.0	Notes
5090	❔ Unknown	❔ Unknown	Not compared: tests.moe.test_cute_dsl_fused_moe.TestAutotuneReplayMemsetContract (4 failures; CUDA 12.9, CUDA 13.0)
B300	❌ New	❌ New	New: tests.moe.test_cute_dsl_fused_moe.TestAutotuneReplayMemsetContract (4 failures; CUDA 12.9, CUDA 13.0)
GB200	❌ New	❌ New	New: tests.moe.test_cute_dsl_fused_moe.TestAutotuneReplayMemsetContract (4 failures; CUDA 12.9, CUDA 13.0)
GB300	❌ New	❌ New	New: tests.moe.test_cute_dsl_fused_moe.TestAutotuneReplayMemsetContract (4 failures; CUDA 12.9, CUDA 13.0)
H100	❌ New	❌ New	New: tests.moe.test_cute_dsl_fused_moe.TestAutotuneReplayMemsetContract (4 failures; CUDA 12.9, CUDA 13.0)
RTX Pro 6000 Blackwell	❌ New	❔ Unknown	New: tests.moe.test_cute_dsl_fused_moe.TestAutotuneReplayMemsetContract (2 failures; CUDA 12.9) Not compared: tests.moe.test_cute_dsl_fused_moe.TestAutotuneReplayMemsetContract (2 failures; CUDA 13.0)

✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · ⚠️ Infrastructure · ❔ Unknown or unclassified · — Not run

Multi-GPU and Multi-Node Tests — 6/6 passed

GPU	CUDA 12.9	CUDA 13.0	Notes
B300 (multi-GPU)	✅ Pass	✅ Pass	—
GB200 (multi-node)	✅ Pass	✅ Pass	—
GB300 (multi-node)	✅ Pass	✅ Pass	—

Failure details

New relative to nightly (attribution uncertain)

• tests.moe.test_cute_dsl_fused_moe.TestAutotuneReplayMemsetContract — 18 failures on B300 / CUDA 12.9, B300 / CUDA 13.0, GB200 / CUDA 12.9, GB200 / CUDA 13.0, GB300 / CUDA 12.9, GB300 / CUDA 13.0, H100 / CUDA 12.9, H100 / CUDA 13.0, RTX Pro 6000 Blackwell / CUDA 12.9
  • ValueError: Expected a cuda device, but got: cpu

Could not compare

• tests.moe.test_cute_dsl_fused_moe.TestAutotuneReplayMemsetContract — 6 failures on 5090 / CUDA 12.9, 5090 / CUDA 13.0, RTX Pro 6000 Blackwell / CUDA 13.0
  • ValueError: Expected a cuda device, but got: cpu

[bkryu]

/bot run tests/moe

[flashinfer-bot]

GitLab MR !1086 has been updated with latest changes, and the CI pipeline #60881596 is currently running. I'll report back once the pipeline job completes.

[coderabbitai]

🧹 Nitpick comments (4)

flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_fp4_helpers.py (4)

1747-1759: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Align has_side_effects across the shared-memory load helpers.

ld_shared_f32_offset sets has_side_effects=True, while ld_shared_u16_offset (Line 823) and ld_shared_u8_offset (Line 854) set False for the same class of operation. The flag controls whether LLVM may reorder or CSE the load across barriers. Pick one policy for all shared loads, and record the reason in a comment. If the intent is to keep the loads pinned relative to barrier calls, False on the u8/u16 variants is the risky side.

🤖 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/cute_dsl/blackwell_sm12x/moe_w4a16_fp4_helpers.py`
around lines 1747 - 1759, Align the has_side_effects policy across
ld_shared_f32_offset, ld_shared_u16_offset, and ld_shared_u8_offset so all
shared-memory load helpers use the same value, preserving barrier ordering; add
a concise comment documenting the chosen policy and its relationship to barrier
calls.

---

5586-5589: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Record why silu_mul_32 uses rcp_approx for the sigmoid.

The device path computes the sigmoid with cute.arch.rcp_approx, while the Torch reference silu_mul_quantize_grouped_mxfp8_torch uses the exact F.silu. The two therefore differ by the rcp_approx error. Add a short comment that states the accuracy tradeoff and the precise alternative, so a later reader does not treat the difference as a defect when comparing against the reference.

As per coding guidelines, "For performance-critical hot paths, comment the justification for special algorithmic choices and relevant alternatives."

🤖 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/cute_dsl/blackwell_sm12x/moe_w4a16_fp4_helpers.py`
around lines 5586 - 5589, Add a concise comment immediately before the sigmoid
computation in silu_mul_32 explaining that rcp_approx is used for performance in
this hot path, with a small accuracy tradeoff versus the exact sigmoid, and
identify F.silu in silu_mul_quantize_grouped_mxfp8_torch as the precise
reference alternative.

Source: Coding guidelines

---

2075-2092: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Wrap val_f32 in Float32 and forward loc/ip.

Two deviations from the pattern used by every other helper in this file:

• val_f32 has no annotation and is not wrapped, so the helper fails if a caller passes a Python float or a non-Float32 DSL value. Compare red_add_global_f32 at Line 1414.
• llvm.inline_asm does not receive loc=loc, ip=ip, so the emitted asm loses source attribution.

♻️ Proposed fix

-def scatter_add_bf16(addr: Int64, val_f32, *, loc=None, ip=None):
+def scatter_add_bf16(addr: Int64, val_f32: Float32, *, loc=None, ip=None):
     """BF16 atomic reduction add to global memory.
 
     Converts the f32 input to bf16 inside PTX and atomically accumulates it into
     one bf16 output lane. This is intended for opt-in approximate reductions.
     """
     llvm.inline_asm(
         None,
         [
             Int64(addr).ir_value(loc=loc, ip=ip),
-            val_f32.ir_value(loc=loc, ip=ip),
+            Float32(val_f32).ir_value(loc=loc, ip=ip),
         ],
         "{ .reg .b16 packed; cvt.rn.satfinite.bf16.f32 packed, $1; red.relaxed.gpu.global.add.noftz.bf16 [$0], packed; }",
         "l,f",
         has_side_effects=True,
         is_align_stack=False,
         asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
     )

🤖 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/cute_dsl/blackwell_sm12x/moe_w4a16_fp4_helpers.py`
around lines 2075 - 2092, Update scatter_add_bf16 to type and normalize val_f32
as Float32, matching red_add_global_f32 so Python floats and other compatible
DSL values are accepted. Also pass loc=loc and ip=ip to llvm.inline_asm while
preserving the existing atomic BF16 reduction behavior.

---

3394-3404: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Unify the bid/tid encoding between the two block-scaled MMA helpers.

This helper embeds bid_a/tid_a/bid_b/tid_b as literal immediates in the asm template. The sibling mxfp8_mma_m16n8k32_f32_e4m3 (Lines 3305-3325) builds four i16 MLIR constants and passes them through h constraints for the same operand positions. Both forms are valid, and the literal form here is simpler. Convert the e4m3 variant to the same form so one pattern covers both, or add a comment that explains why the two differ.

🤖 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/cute_dsl/blackwell_sm12x/moe_w4a16_fp4_helpers.py`
around lines 3394 - 3404, Unify the operand encoding in the e4m3 block-scaled
MMA helper around the asm construction near `mxfp8_mma_m16n8k32_f32_e4m3`:
either replace the literal `bid_a`/`tid_a`/`bid_b`/`tid_b` immediates with four
`i16` MLIR constants passed through `h` constraints like the sibling helper, or
document the intentional difference directly beside this template.

🤖 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.

Nitpick comments:
In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_fp4_helpers.py`:
- Around line 1747-1759: Align the has_side_effects policy across
ld_shared_f32_offset, ld_shared_u16_offset, and ld_shared_u8_offset so all
shared-memory load helpers use the same value, preserving barrier ordering; add
a concise comment documenting the chosen policy and its relationship to barrier
calls.
- Around line 5586-5589: Add a concise comment immediately before the sigmoid
computation in silu_mul_32 explaining that rcp_approx is used for performance in
this hot path, with a small accuracy tradeoff versus the exact sigmoid, and
identify F.silu in silu_mul_quantize_grouped_mxfp8_torch as the precise
reference alternative.
- Around line 2075-2092: Update scatter_add_bf16 to type and normalize val_f32
as Float32, matching red_add_global_f32 so Python floats and other compatible
DSL values are accepted. Also pass loc=loc and ip=ip to llvm.inline_asm while
preserving the existing atomic BF16 reduction behavior.
- Around line 3394-3404: Unify the operand encoding in the e4m3 block-scaled MMA
helper around the asm construction near `mxfp8_mma_m16n8k32_f32_e4m3`: either
replace the literal `bid_a`/`tid_a`/`bid_b`/`tid_b` immediates with four `i16`
MLIR constants passed through `h` constraints like the sibling helper, or
document the intentional difference directly beside this template.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: d7110405-96db-4e2b-bda1-576a7b578c62

📥 Commits

Reviewing files that changed from the base of the PR and between adb976a and f9afbf0.

📒 Files selected for processing (3)

• flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py
• flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_fp4_helpers.py
• tests/moe/test_b12x_fused_moe.py

🚧 Files skipped from review as they are similar to previous changes (2)

• tests/moe/test_b12x_fused_moe.py
• flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py