Conversation
Route gated activations through the optimized branch-paired dynamic kernel while preserving the generic fallback for non-gated activations. Support SiLU, GELU-tanh, and SwiGLU-OAI without an environment toggle.
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
📝 WalkthroughWalkthroughThe PR adds generic and optimized gated dynamic SM12x NVFP4 MoE kernels. It adds runtime task scheduling, routed-input quantization, activation-specific GEMM pipelines, weighted output scattering, capability-based dispatch, and regression tests. ChangesDynamic NVFP4 MoE execution
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Host as MoEDynamicKernel
participant Dispatch as moe_dispatch
participant Kernel as Selected dynamic kernel
participant Queue as Task queue
participant Output as Weighted scatter
Host->>Dispatch: provide model dimensions and top-k
Dispatch->>Kernel: select generic or gated implementation
Kernel->>Queue: publish and claim routed expert tasks
Queue->>Kernel: provide task metadata
Kernel->>Output: scatter weighted FC2 results
Possibly related PRs
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}
There was a problem hiding this comment.
Actionable comments posted: 9
🧹 Nitpick comments (19)
flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/generic.py (3)
748-748: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefix the unused
bidxbinding.Ruff reports RUF059 for the unpacked variable
bidx.♻️ Proposed fix
- bidx, _, bidz = cute.arch.block_idx() + _bidx, _, bidz = cute.arch.block_idx()🤖 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_dynamic/generic.py` at line 748, Update the block index unpacking in the relevant kernel function to prefix the unused bidx binding with an underscore, while preserving the existing block_idx() unpacking and the bidz value.Source: Linters/SAST tools
214-234: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRename
_spin_wait_global_eq_i32to match its behavior.The inline PTX branches back to
spin_loopwhile the loaded value equals$1. The helper therefore waits until the value differs fromexpected. Both call sites depend on that behavior: Line 442 waits while the epoch staysold_epoch, and Line 1793 waits whiletask_readystays0. The current name states the opposite condition and can cause an incorrect edit later.♻️ Proposed rename
-def _spin_wait_global_eq_i32(addr, expected, *, loc=None, ip=None): +def _spin_wait_global_ne_i32(addr, while_value, *, loc=None, ip=None): + """Spin until the value at ``addr`` differs from ``while_value``."""Update both call sites accordingly.
🤖 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_dynamic/generic.py` around lines 214 - 234, Rename the helper _spin_wait_global_eq_i32 to reflect that it spins while the loaded value equals expected and returns only after it differs. Update both call sites at the epoch wait and task_ready wait to use the new name, preserving their existing arguments and behavior.
961-961: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove or gate the dead full-tile publish path.
full_tile_publish_enabledis a compile-time constantInt32(0). Everyfull_tile_publish_enabled > Int32(0)block is therefore unreachable, including thetile_write_countzeroing, both incremental publish blocks, theproducers_done_countflush, and the CAS-based consumer claim at Lines 1784-1815._publish_ready_tasks,_atomic_cas_global_i32,tile_write_count, andproducers_done_countbecome dead as a result.Two options:
- Delete the unreachable blocks and the now-unused helpers.
- Convert
full_tile_publish_enabledinto acutlass.Constexprconstruction option and select the path withcutlass.const_expr, so the intent stays explicit and the dead branch is not emitted.The current form keeps a second, untested control plane in the file and hides which queue protocol the kernel actually uses.
Also applies to: 1003-1007, 1264-1301, 1395-1422, 1487-1520, 1784-1815
🤖 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_dynamic/generic.py` at line 961, Remove the unreachable full-tile publish control plane rooted at full_tile_publish_enabled = Int32(0), including its guarded tile_write_count, producers_done_count, incremental publish, flush, and CAS consumer-claim blocks; then remove now-unused helpers such as _publish_ready_tasks and _atomic_cas_global_i32. Alternatively, make full_tile_publish_enabled a cutlass.Constexpr construction option and gate every listed branch with cutlass.const_expr so only the selected protocol is emitted.flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py (16)
1063-1133: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueShare the descriptor-write block across the three publishers.
publish_ready_tasks,publish_uniform_deferred_tasks, andpublish_variable_deferred_tasksrepeat the same loop body. Only thestartvalue, theslice_chunkderivation, and the trailingtask_readyrelease differ. Extract the common write loop into one helper that takesstart,num_groups, andslice_chunk. That removes two copies of the slot arithmetic and keeps any future bounds check in one place.🤖 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_dynamic/gated.py` around lines 1063 - 1133, Extract the repeated descriptor-write loop from publish_ready_tasks, publish_uniform_deferred_tasks, and publish_variable_deferred_tasks into a shared helper accepting start, num_groups, and slice_chunk, while preserving each publisher’s existing task_ready release behavior. Replace all three loop bodies with calls to the helper so slot arithmetic and descriptor assignments have one implementation.
3082-3091: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winThe prefix-sum fast path also assumes
num_mma_warps == 8.The condition tests
num_experts == Int32(256), but the implementation also requires exactly 256 participating threads: it maps one expert per thread (rows = row_counts[tidx],expert_tile_base[tidx] = ...), scans 5 shuffle stages for 32 lanes, and combines exactlyself.num_mma_warpswarp subtotals. That holds only whileself.num_mma_warps * self.num_threads_per_warp == 256.
num_mma_warpsis set to 8 in the constructor. A future change to 4 or 16 would leave thenum_experts == 256test passing while the scan silently produces a wrong tile prefix, which then mis-routes every token. Add a compile-time assertion or state the coupling in a comment.🤖 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_dynamic/gated.py` around lines 3082 - 3091, The prefix-sum fast path condition in the gated MoE implementation must also enforce its 256-thread participation requirement, not only num_experts == Int32(256). Add a compile-time assertion that self.num_mma_warps * self.num_threads_per_warp equals 256, or document and enforce this coupling near the condition, so changes to num_mma_warps cannot silently use the incompatible scan.
2569-2579: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDrop the unused
acc_shapeanddown_alpha_valueparameters.
fc2_epilogue_to_sCreads neither parameter.down_alpha_valueis misleading here, because the expert down-alpha is applied later insidescatter_add_weighted_bf16x8_packed_alpha(lines 2681, 2723). A reader who checks whether alpha is applied twice must read both functions to rule it out. Remove both parameters and their arguments at the call site (lines 4943-4951).♻️ Proposed signature change
def fc2_epilogue_to_sC( self, - acc_shape, - down_alpha_value, down_acc, sC,self.fc2_epilogue_to_sC( - acc_shape, - down_alpha_value, down_acc, sC,🤖 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_dynamic/gated.py` around lines 2569 - 2579, Remove the unused acc_shape and down_alpha_value parameters from fc2_epilogue_to_sC, then remove the corresponding arguments from its call site around the fc2 epilogue invocation. Leave alpha application in scatter_add_weighted_bf16x8_packed_alpha unchanged.
3717-3734: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExplain the task-splitting constants.
This block chooses
split_tile_countfromnum_tokensthresholds 256, 2048, and 4096, a target of4 * gdim_z, and the expression(Int32(125) * Int32(gdim_z) + Int32(31)) // Int32(32). The origin of125and the intent of the 2048 threshold are not recorded. A reader cannot tell which values are measured tuning points and which are structural.Add two or three lines that state what each threshold targets, for example the intended tasks-per-SM ratio. That keeps the heuristic reproducible when the occupancy or tile shape changes.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py` around lines 3717 - 3734, Document the task-splitting heuristic beside the threshold logic in the dynamic MoE scheduling block, identifying 256, 2048, and 4096 as tuning boundaries, explaining that 4 * gdim_z and the 125/32 expression target specific tasks-per-SM occupancy, and distinguishing measured tuning values from structural calculations. Keep the existing split_tile_count behavior unchanged.
1381-1385: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the dead initial assignments before the stage selection.
Line 1381 assigns
csSFB_up_pand lines 1382-1385 immediately overwrite it on both branches. The same dead assignment appears at line 1484, line 1550, line 1809, line 1913, and line 1980. Deleting them makes theab_storage_stagealiasing rule easier to follow, because each variable then has exactly one definition per path.♻️ Proposed cleanup
- csSFB_up_p = csSFB_up_fc1_half[None, None, None, Int32(0)] if cons_state.index < Int32(self.ab_storage_stage): csSFB_up_p = csSFB_up_fc1_half[None, None, None, cons_state.index] else: csSFB_up_p = csSFB_up_fc1_extra_half- csSFB_up_cur = csSFB_up_p if cons_state.index < Int32(self.ab_storage_stage):🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py` around lines 1381 - 1385, Remove the redundant initial assignments to csSFB_up_p before the cons_state.index stage-selection branches, keeping only the branch-specific definitions. Apply the same cleanup to the corresponding dead assignments at the other identified locations, while preserving the existing ab_storage_stage selection logic.
1688-1734: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftMerge the two FC1 functions behind a compile-time predicate flag.
fc1_gate_up_swiglu_to_sC_tailduplicates about 430 lines offc1_gate_up_swiglu_to_sC. The only differences are the extrawarp_m_coordparameter and thevalid_rows > Int32(_mt * 64) + warp_m_coord * Int32(16)guard around each MMA group, the hold flush, and the activation loop. Every pipeline, staging, and aliasing detail is copied.Two copies of this logic will drift. A bug fix in the aliasing rule or the SFB half selection must be applied twice, in two places that a reviewer cannot diff easily. Take one function with a
cutlass.Constexprpredicate flag, so the compiler still generates the predicate-free full-tile variant:`@cute.jit` def fc1_gate_up_swiglu_to_sC(self, ..., warp_m_coord: Int32, *, predicated: cutlass.Constexpr): ... for _mt in cutlass.range_constexpr(fc1_m_tiles): if cutlass.const_expr(predicated): row_active = valid_rows > Int32(_mt * 64) + warp_m_coord * Int32(16) ...The same argument applies to
fc2_accumulate_sliceandfc2_accumulate_slice_tail(lines 2390-2567), which duplicate about 90 lines under the identical predicate.If the duplication is intentional for tuning stability in this release, add a comment that says so, so a future reader does not merge them by accident.
🤖 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_dynamic/gated.py` around lines 1688 - 1734, Merge fc1_gate_up_swiglu_to_sC_tail into fc1_gate_up_swiglu_to_sC using a cutlass.Constexpr predicated flag, preserving the existing valid_rows guard only when predicated is enabled while keeping the full-tile path predicate-free. Apply the same consolidation to fc2_accumulate_slice and fc2_accumulate_slice_tail using the identical compile-time predicate approach. If retaining separate implementations for tuning stability, add an explicit comment documenting that intent instead.
2804-2812: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winThe
pass_gate_barrierarrival balance depends on reaching pipeline stage 2.The TMA warp consumes one
pass_gate_barrierarrival per slice only whenprod_state.indexreachesab_storage_stage. Each half issuesfc1_k_tile_cntstages andprod_stateis reset per slice, so stage 2 is reached only when2 * fc1_k_tile_cnt >= 3.The math warps call
pass_gate_barrier.arrive_unaligned()once per slice unconditionally (line 4854), and the TMA warp performs one finalwait_unaligned()at line 5063. Iffc1_k_tile_cnt == 1, the arrivals and waits no longer balance across slices.
fc1_k_tile_cnt = hidden_size // tile_shape_mnk[2], so this needshidden_size >= 256. Real configurations satisfy it, but the protocol depends on it silently. Record the precondition in a comment, or validatehidden_sizein_setup_attributes.🤖 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_dynamic/gated.py` around lines 2804 - 2812, Document the required protocol precondition near the stage-2 wait in the k_tile loop: hidden_size must produce fc1_k_tile_cnt >= 2 so prod_state.index reaches ab_storage_stage and balances pass_gate_barrier arrivals and waits. Prefer validating this in _setup_attributes using the existing hidden_size and tile_shape_mnk values; otherwise add a precise comment explaining the invariant and the failure mode when fc1_k_tile_cnt == 1.
2290-2296: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winConfirm the deferred slot order matches between quantize and flush.
quantize_q1_sC_to_sA_sSFAincrementsdeferred_slotinside theepi_mloop, so the slot order followsepi_mfirst and thenquant_idxwithinepi_rows * sf_blocks_per_row.flush_deferred_q1_aandflush_deferred_q1_sfare-derive the slot order from one flat range overvalid_rows * sf_blocks_per_row.The two orders agree only while
epi_rest_m == 1, which holds becauseepi_tile[0] == mma_tiler_mn[0] == tile_shape_mnk[0]. Ifepi_tileis ever decoupled from the CTA tile, the flush reads the wrong deferred slot for each block and the FC2 input becomes wrong without any bounds error. Add a comment that records this coupling, or assertepi_rest_m == 1where the deferred path is selected.🤖 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_dynamic/gated.py` around lines 2290 - 2296, Document the deferred-slot ordering invariant at the deferred quantize/flush path, anchored to quantize_q1_sC_to_sA_sSFA and the flush_deferred_q1_a/flush_deferred_q1_sfa consumers: the flat flush order is valid only when epi_rest_m == 1, which currently follows from epi_tile[0] == mma_tiler_mn[0] == tile_shape_mnk[0]. Add either a clear comment recording this coupling or an assertion enforcing epi_rest_m == 1 when selecting the deferred path.
2663-2680: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the shared-memory swizzle transform into one helper.
The expression
offset ^ ((offset & Int32(0x1C0)) >> Int32(3))appears four times: lines 2184-2186, 2196-2198, 2674-2676, and 2716-2718. It encodes the epilogueS<3,4,3>swizzle in BF16 element units. The quantize path and the scatter path must agree on it exactly. If they diverge, both read valid shared memory at the wrong addresses, and the result is wrong values with no fault.Add one small helper and call it from all four sites.
♻️ Proposed helper
+def _apply_epi_swizzle(element_offset: Int32) -> Int32: + """Apply the sC S<3,4,3> swizzle in BF16 element units. + + ``sC.layout`` returns the unswizzled offset. A raw shared pointer does not + retain CuTe's swizzle transform, so apply it explicitly. + """ + return element_offset ^ ((element_offset & Int32(0x1C0)) >> Int32(3))- sc_element_offset = sc_element_offset ^ ( - (sc_element_offset & Int32(0x1C0)) >> Int32(3) - ) + sc_element_offset = _apply_epi_swizzle(sc_element_offset)🤖 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_dynamic/gated.py` around lines 2663 - 2680, Extract the repeated epilogue S<3,4,3> shared-memory offset transform into a single helper near the existing utilities, preserving Int32 arithmetic and the exact mask/shift operation. Replace the four inline transforms in the quantize and scatter paths, including the code around sc_element_offset and the corresponding sites near lines 2184, 2196, and 2716, with calls to that helper so every path uses identical swizzling.
1-1: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd automated coverage for the split, and document the new N128 tile constraint.
The PR objectives report manual validation on one SM120 GPU. This review set contains no test change. The split introduces two behavioral facts that a test should pin:
MoEDynamicKernelselectsMoEGatedDynamicKernelforsilu,gelu_tanh, andswigluoai_uninterleave, and the generic implementation forrelu2.- The gated implementation now rejects
mma_tiler_mn[1] != 128.Point 2 is a new user-visible constraint. Record it where the dynamic backend is documented, so a caller that tunes tile shapes learns the restriction before hitting the
ValueError.I can draft the parametrized dispatch test and the documentation note. Tell me if you want me to open an issue to track it.
As per coding guidelines: "When adding a new operation, provide a Python API, JIT module generator, tests, AOT registration, package export, and trace integration as applicable" and "Keep documentation synchronized with code changes."
#!/bin/bash # Locate existing tests and docs for the SM120 dynamic MoE backend. fd -t f -i 'moe' tests | head -40 rg -nl --type=md -i 'sm12|blackwell.*moe|dynamic moe' docs 2>/dev/null | head -20🤖 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_dynamic/gated.py` at line 1, Add parametrized automated coverage for MoEDynamicKernel dispatch, asserting silu, gelu_tanh, and swigluoai_uninterleave select MoEGatedDynamicKernel while relu2 selects the generic implementation. Document the gated backend constraint that mma_tiler_mn[1] must equal 128, including that other values are rejected, in the existing dynamic backend documentation.Source: Coding guidelines
152-232: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the unused gated scatter helpers, or document the intended use.
scatter_add_v4_bf16x2inflashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.pyhas no call sites, unlike the same helper in another file.scatter_add_weighted_bf16x8_packedis also unused in the gated scatter path, which calls the alpha helper repeatedly. Drop both helpers unless there is a planned use.🤖 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_dynamic/gated.py` around lines 152 - 232, Remove the unused helper definitions scatter_add_v4_bf16x2 and scatter_add_weighted_bf16x8_packed from the gated module, since the current scatter path does not call either function. Do not alter the existing alpha-helper-based scatter flow.
433-482: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove or wire the unused helpers in
_moe_dynamic.
flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.pydefinesload_global_bf16x16_to_f32x16,_ld_global_u64, and_atomic_cas_global_i32, but onlyload_shared_bf16x16_to_f32x16has call sites in this file. This leaves dead helper definitions in the dynamic 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/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py` around lines 433 - 482, Remove the unused helper definitions load_global_bf16x16_to_f32x16, _ld_global_u64, and _atomic_cas_global_i32 from the _moe_dynamic implementation, since only load_shared_bf16x16_to_f32x16 is referenced. Preserve the existing shared-memory loading path and any helpers with active call sites.
111-114: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the unused constants from
gated.py.
_SF_VEC_SIZE,_PRODUCER_PAIRS_PER_WARP, and_FC2_TILE_RECIP_GS_NUMare no longer referenced in this file, while_TASK_SLICE_CHUNKis used at line 2989. Remove the dead constants or add a local rationale that justifies keeping them.🤖 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_dynamic/gated.py` around lines 111 - 114, Remove the unused module constants _SF_VEC_SIZE, _PRODUCER_PAIRS_PER_WARP, and _FC2_TILE_RECIP_GS_NUM from gated.py, while preserving _TASK_SLICE_CHUNK because it is referenced by the task-slicing logic.
529-548: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRename
_spin_wait_global_eq_i32to describe its actual behavior.The inline asm waits while the loaded value equals
expected, soresident_grid_barrier()exits when the epoch changes. Rename the helper, for example to_spin_wait_global_ne_i32, so this synchronization primitive is not read as “wait until equal to expected”.🤖 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_dynamic/gated.py` around lines 529 - 548, Rename the helper `_spin_wait_global_eq_i32` to `_spin_wait_global_ne_i32` (or an equivalent name indicating it waits while equal and exits when different), and update every call site such as `resident_grid_barrier()` to use the new name.
868-898: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the fixed stage counts and keep the shared-memory budget assertive.
The generic gated path derives
ab_stageand checks it dividesk_tile_cnt; this implementation fixesab_stage = 3,ab_storage_stage = 2, andphase2_stage = 3. Record the FC1 per-slicereset_count()alignment and the staging-choice rationale here. SinceStorageGatednow uses fixed stages, assert the totalsmemallocation againstself.smem_capacityso over-budget variants fail during setup instead of later.🤖 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_dynamic/gated.py` around lines 868 - 898, Update the fixed-stage setup around ab_stage, ab_storage_stage, and phase2_stage to document the FC1 per-slice reset_count() alignment and why the staging choices are used. After StorageGated computes its shared-memory allocation, add an assert comparing the total smem requirement with self.smem_capacity so over-budget configurations fail during setup.Source: Coding guidelines
4119-4124: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove the unselected FC1 dataflow scaffolding.
sequential_branch_compactandfc1_storage_aliasare only read throughgetattr(..., False)so the compact FC1 branch path is unreachable here. Remove the deadup_pipeline/up_pipeline_array,up_prod_state, andup_cons_statewire, the zero-lengthroute_phys_rows,route_expert_ids,scatter_weight_cache, andsB_upfields, and theup_pipeline.producer_tail(up_prod_state)guard. Replace the removed storage regions with comments pointing to the actual storage regions, and keepsSFB_upsized unconditionally. Apply the same parameter removal tofc1_gate_up_swiglu_to_sC,fc1_gate_up_swiglu_to_sC_tail, andload_fc1_tma_slice.🤖 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_dynamic/gated.py` around lines 4119 - 4124, Remove the unreachable compact-FC1 scaffolding in gated.py: delete up_pipeline/up_pipeline_array, up_prod_state, up_cons_state, zero-length route_phys_rows, route_expert_ids, scatter_weight_cache, and sB_up storage, replacing removed regions with comments referencing the actual storage regions; size sSFB_up unconditionally and remove the up_pipeline.producer_tail(up_prod_state) guard. Update fc1_gate_up_swiglu_to_sC, fc1_gate_up_swiglu_to_sC_tail, and load_fc1_tma_slice to remove the corresponding parameters, including the affected sites in gated.py at lines 4119-4124 and 1265-1268.
🤖 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_dynamic/gated.py`:
- Line 4095: Update the block index unpacking in the surrounding function to
bind the unused first value to an underscore instead of bidx, while preserving
bidz for subsequent use.
- Around line 2234-2272: Update the constructor validation for the dynamic MoE
kernel to enforce the supported deferred-buffer geometry: constrain the relevant
mma_tiler_mn[0] and sf_vec_size values so each thread processes at most four
deferred blocks. Anchor the change in the constructor and preserve the existing
deferred_a_words and deferred_sfa_words indexing in the processing logic.
- Around line 748-771: Validate sf_vec_size in the constructor alongside the
existing tile-shape checks, requiring it to equal the hardcoded 16-element
quantization block size. Reject any other value before computing tile_shape_mnk
and related SFB dimensions, while preserving the existing tile and SFB
validation.
- Around line 4270-4275: Update the tensor recasting block to assign each
cute.recast_tensor result back to its corresponding tensor variable, including
sA, sB, sB_phase2_extra, sB_fc1_all, sB_fc1, and sB_up_fc1, so subsequent GPU
operations use the Uint8 views rather than stale dtype views.
- Around line 3506-3523: Add a host-side validation wherever the runtime
num_topk/top_k is established, rejecting values greater than 16 before launching
the kernel. Preserve routing behavior by failing explicitly rather than
clamping, and ensure the validation covers both the route_gs population and
subsequent scale-read paths in the relevant MoE setup flow.
- Around line 3004-3021: Ensure the dynamic launcher’s scatter_output input is
contiguous and 16-byte aligned before passing scatter_output.data_ptr() to the
kernel, either by validating and rejecting invalid tensors or materializing a
contiguous tensor and preserving the expected output behavior. Document this
requirement at the relevant public API, and keep the kernel’s default row-major
[num_tokens, k] layout consistent with the enforced contract.
- Around line 3889-3899: Update MoEDynamicKernel.__call__ to derive a
per-invocation configuration containing a_dtype, b_dtype, sf_dtype, a_layout,
b_layout, c_layout, and the layouts produced by _setup_attributes instead of
mutating shared self state. Ensure compiled-kernel generation and dispatch
consume only that local configuration, so cached kernels remain safe across
differing inputs and concurrent calls.
- Around line 4811-4832: Update publish_variable_deferred_tasks so its
final-group slice_chunk is capped to the same maximum used by the uniform
publisher, keeping slice_idx within the defined 0–3 FC2 stage mappings and
preventing later slices from reusing slice 0 stages.
In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/generic.py`:
- Around line 1-47: Update the module docstring for MoEDynamicKernel to document
all activation modes accepted by __init__: silu, relu2, gelu_tanh, and
swigluoai_uninterleave, including their relevant behavior. Revise or remove the
closing statement that calls the implementation uncompiled or unprofiled so the
documentation reflects the validated implementation.
---
Nitpick comments:
In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py`:
- Around line 1063-1133: Extract the repeated descriptor-write loop from
publish_ready_tasks, publish_uniform_deferred_tasks, and
publish_variable_deferred_tasks into a shared helper accepting start,
num_groups, and slice_chunk, while preserving each publisher’s existing
task_ready release behavior. Replace all three loop bodies with calls to the
helper so slot arithmetic and descriptor assignments have one implementation.
- Around line 3082-3091: The prefix-sum fast path condition in the gated MoE
implementation must also enforce its 256-thread participation requirement, not
only num_experts == Int32(256). Add a compile-time assertion that
self.num_mma_warps * self.num_threads_per_warp equals 256, or document and
enforce this coupling near the condition, so changes to num_mma_warps cannot
silently use the incompatible scan.
- Around line 2569-2579: Remove the unused acc_shape and down_alpha_value
parameters from fc2_epilogue_to_sC, then remove the corresponding arguments from
its call site around the fc2 epilogue invocation. Leave alpha application in
scatter_add_weighted_bf16x8_packed_alpha unchanged.
- Around line 3717-3734: Document the task-splitting heuristic beside the
threshold logic in the dynamic MoE scheduling block, identifying 256, 2048, and
4096 as tuning boundaries, explaining that 4 * gdim_z and the 125/32 expression
target specific tasks-per-SM occupancy, and distinguishing measured tuning
values from structural calculations. Keep the existing split_tile_count behavior
unchanged.
- Around line 1381-1385: Remove the redundant initial assignments to csSFB_up_p
before the cons_state.index stage-selection branches, keeping only the
branch-specific definitions. Apply the same cleanup to the corresponding dead
assignments at the other identified locations, while preserving the existing
ab_storage_stage selection logic.
- Around line 1688-1734: Merge fc1_gate_up_swiglu_to_sC_tail into
fc1_gate_up_swiglu_to_sC using a cutlass.Constexpr predicated flag, preserving
the existing valid_rows guard only when predicated is enabled while keeping the
full-tile path predicate-free. Apply the same consolidation to
fc2_accumulate_slice and fc2_accumulate_slice_tail using the identical
compile-time predicate approach. If retaining separate implementations for
tuning stability, add an explicit comment documenting that intent instead.
- Around line 2804-2812: Document the required protocol precondition near the
stage-2 wait in the k_tile loop: hidden_size must produce fc1_k_tile_cnt >= 2 so
prod_state.index reaches ab_storage_stage and balances pass_gate_barrier
arrivals and waits. Prefer validating this in _setup_attributes using the
existing hidden_size and tile_shape_mnk values; otherwise add a precise comment
explaining the invariant and the failure mode when fc1_k_tile_cnt == 1.
- Around line 2290-2296: Document the deferred-slot ordering invariant at the
deferred quantize/flush path, anchored to quantize_q1_sC_to_sA_sSFA and the
flush_deferred_q1_a/flush_deferred_q1_sfa consumers: the flat flush order is
valid only when epi_rest_m == 1, which currently follows from epi_tile[0] ==
mma_tiler_mn[0] == tile_shape_mnk[0]. Add either a clear comment recording this
coupling or an assertion enforcing epi_rest_m == 1 when selecting the deferred
path.
- Around line 2663-2680: Extract the repeated epilogue S<3,4,3> shared-memory
offset transform into a single helper near the existing utilities, preserving
Int32 arithmetic and the exact mask/shift operation. Replace the four inline
transforms in the quantize and scatter paths, including the code around
sc_element_offset and the corresponding sites near lines 2184, 2196, and 2716,
with calls to that helper so every path uses identical swizzling.
- Line 1: Add parametrized automated coverage for MoEDynamicKernel dispatch,
asserting silu, gelu_tanh, and swigluoai_uninterleave select
MoEGatedDynamicKernel while relu2 selects the generic implementation. Document
the gated backend constraint that mma_tiler_mn[1] must equal 128, including that
other values are rejected, in the existing dynamic backend documentation.
- Around line 152-232: Remove the unused helper definitions
scatter_add_v4_bf16x2 and scatter_add_weighted_bf16x8_packed from the gated
module, since the current scatter path does not call either function. Do not
alter the existing alpha-helper-based scatter flow.
- Around line 433-482: Remove the unused helper definitions
load_global_bf16x16_to_f32x16, _ld_global_u64, and _atomic_cas_global_i32 from
the _moe_dynamic implementation, since only load_shared_bf16x16_to_f32x16 is
referenced. Preserve the existing shared-memory loading path and any helpers
with active call sites.
- Around line 111-114: Remove the unused module constants _SF_VEC_SIZE,
_PRODUCER_PAIRS_PER_WARP, and _FC2_TILE_RECIP_GS_NUM from gated.py, while
preserving _TASK_SLICE_CHUNK because it is referenced by the task-slicing logic.
- Around line 529-548: Rename the helper `_spin_wait_global_eq_i32` to
`_spin_wait_global_ne_i32` (or an equivalent name indicating it waits while
equal and exits when different), and update every call site such as
`resident_grid_barrier()` to use the new name.
- Around line 868-898: Update the fixed-stage setup around ab_stage,
ab_storage_stage, and phase2_stage to document the FC1 per-slice reset_count()
alignment and why the staging choices are used. After StorageGated computes its
shared-memory allocation, add an assert comparing the total smem requirement
with self.smem_capacity so over-budget configurations fail during setup.
- Around line 4119-4124: Remove the unreachable compact-FC1 scaffolding in
gated.py: delete up_pipeline/up_pipeline_array, up_prod_state, up_cons_state,
zero-length route_phys_rows, route_expert_ids, scatter_weight_cache, and sB_up
storage, replacing removed regions with comments referencing the actual storage
regions; size sSFB_up unconditionally and remove the
up_pipeline.producer_tail(up_prod_state) guard. Update fc1_gate_up_swiglu_to_sC,
fc1_gate_up_swiglu_to_sC_tail, and load_fc1_tma_slice to remove the
corresponding parameters, including the affected sites in gated.py at lines
4119-4124 and 1265-1268.
In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/generic.py`:
- Line 748: Update the block index unpacking in the relevant kernel function to
prefix the unused bidx binding with an underscore, while preserving the existing
block_idx() unpacking and the bidz value.
- Around line 214-234: Rename the helper _spin_wait_global_eq_i32 to reflect
that it spins while the loaded value equals expected and returns only after it
differs. Update both call sites at the epoch wait and task_ready wait to use the
new name, preserving their existing arguments and behavior.
- Line 961: Remove the unreachable full-tile publish control plane rooted at
full_tile_publish_enabled = Int32(0), including its guarded tile_write_count,
producers_done_count, incremental publish, flush, and CAS consumer-claim blocks;
then remove now-unused helpers such as _publish_ready_tasks and
_atomic_cas_global_i32. Alternatively, make full_tile_publish_enabled a
cutlass.Constexpr construction option and gate every listed branch with
cutlass.const_expr so only the selected protocol is emitted.
🪄 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: 62e61abc-7cf8-41cc-9951-3c6948252176
📒 Files selected for processing (4)
flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/__init__.pyflashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.pyflashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/generic.pyflashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dynamic_kernel.py
| self.sf_vec_size = sf_vec_size | ||
| self.input_scales_are_reciprocal = input_scales_are_reciprocal | ||
| self.fast_math = fast_math | ||
| self.activation = activation | ||
| self.swiglu_alpha = float(swiglu_alpha) | ||
| self.swiglu_beta = float(swiglu_beta) | ||
| self.swiglu_limit = float(swiglu_limit) if swiglu_limit is not None else None | ||
| self.share_input_across_experts = share_input_across_experts | ||
| tile_k = sf_vec_size * 8 | ||
| self.tile_shape_mnk = (mma_tiler_mn[0], mma_tiler_mn[1], tile_k) | ||
| self.fc1_tile_shape_mnk = ( | ||
| mma_tiler_mn[0], | ||
| mma_tiler_mn[1] // 2, | ||
| tile_k, | ||
| ) | ||
| self.fc1_sfb_tile_shape_nk = ( | ||
| max(128, self.fc1_tile_shape_mnk[1]), | ||
| tile_k, | ||
| ) | ||
| self.fc1_sfb_tiles_per_block = ( | ||
| self.fc1_sfb_tile_shape_nk[0] // self.fc1_tile_shape_mnk[1] | ||
| ) | ||
| if self.fc1_sfb_tiles_per_block != 2: | ||
| raise ValueError("expected exactly two logical N64 tiles per SFB block") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Validate sf_vec_size in the constructor.
The constructor validates the N128 tile and the SFB block split, but it accepts any sf_vec_size. The quantization path hardcodes a 16-element scale block: sf_blocks_per_row = tile_shape_mnk[2] // 16 (line 2145), block_start = sf_block * Int32(16) (line 2169), a 16-element values tensor (line 2209), and the Int32(32 * 4 * 4) scale-layout strides (lines 2255-2260). If a caller passes sf_vec_size != 16, tile_k changes while the block size stays 16, and the scale factors silently no longer match the packed data. Add an explicit check next to the existing tile checks.
🛡️ Proposed guard
if self.fc1_sfb_tiles_per_block != 2:
raise ValueError("expected exactly two logical N64 tiles per SFB block")
+ if sf_vec_size != _SF_VEC_SIZE:
+ raise ValueError(
+ "the gated dynamic kernel hardcodes a 16-element scale block; "
+ f"got sf_vec_size={sf_vec_size}"
+ )📝 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.
| self.sf_vec_size = sf_vec_size | |
| self.input_scales_are_reciprocal = input_scales_are_reciprocal | |
| self.fast_math = fast_math | |
| self.activation = activation | |
| self.swiglu_alpha = float(swiglu_alpha) | |
| self.swiglu_beta = float(swiglu_beta) | |
| self.swiglu_limit = float(swiglu_limit) if swiglu_limit is not None else None | |
| self.share_input_across_experts = share_input_across_experts | |
| tile_k = sf_vec_size * 8 | |
| self.tile_shape_mnk = (mma_tiler_mn[0], mma_tiler_mn[1], tile_k) | |
| self.fc1_tile_shape_mnk = ( | |
| mma_tiler_mn[0], | |
| mma_tiler_mn[1] // 2, | |
| tile_k, | |
| ) | |
| self.fc1_sfb_tile_shape_nk = ( | |
| max(128, self.fc1_tile_shape_mnk[1]), | |
| tile_k, | |
| ) | |
| self.fc1_sfb_tiles_per_block = ( | |
| self.fc1_sfb_tile_shape_nk[0] // self.fc1_tile_shape_mnk[1] | |
| ) | |
| if self.fc1_sfb_tiles_per_block != 2: | |
| raise ValueError("expected exactly two logical N64 tiles per SFB block") | |
| self.sf_vec_size = sf_vec_size | |
| self.input_scales_are_reciprocal = input_scales_are_reciprocal | |
| self.fast_math = fast_math | |
| self.activation = activation | |
| self.swiglu_alpha = float(swiglu_alpha) | |
| self.swiglu_beta = float(swiglu_beta) | |
| self.swiglu_limit = float(swiglu_limit) if swiglu_limit is not None else None | |
| self.share_input_across_experts = share_input_across_experts | |
| tile_k = sf_vec_size * 8 | |
| self.tile_shape_mnk = (mma_tiler_mn[0], mma_tiler_mn[1], tile_k) | |
| self.fc1_tile_shape_mnk = ( | |
| mma_tiler_mn[0], | |
| mma_tiler_mn[1] // 2, | |
| tile_k, | |
| ) | |
| self.fc1_sfb_tile_shape_nk = ( | |
| max(128, self.fc1_tile_shape_mnk[1]), | |
| tile_k, | |
| ) | |
| self.fc1_sfb_tiles_per_block = ( | |
| self.fc1_sfb_tile_shape_nk[0] // self.fc1_tile_shape_mnk[1] | |
| ) | |
| if self.fc1_sfb_tiles_per_block != 2: | |
| raise ValueError("expected exactly two logical N64 tiles per SFB block") | |
| if sf_vec_size != _SF_VEC_SIZE: | |
| raise ValueError( | |
| "the gated dynamic kernel hardcodes a 16-element scale block; " | |
| f"got sf_vec_size={sf_vec_size}" | |
| ) |
🤖 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_dynamic/gated.py` around
lines 748 - 771, Validate sf_vec_size in the constructor alongside the existing
tile-shape checks, requiring it to equal the hardcoded 16-element quantization
block size. Reject any other value before computing tile_shape_mnk and related
SFB dimensions, while preserving the existing tile and SFB validation.
| if defer_a > Int32(0): | ||
| deferred_a_words[deferred_slot * Int32(2)] = Uint32( | ||
| packed64 & Uint64(0xFFFFFFFF) | ||
| ) | ||
| deferred_a_words[deferred_slot * Int32(2) + Int32(1)] = Uint32( | ||
| packed64 >> Uint64(32) | ||
| ) | ||
| else: | ||
| for byte_idx in cutlass.range_constexpr(8): | ||
| src_pcol = packed_base + Int32(byte_idx) | ||
| dst_row = ((src_pcol ^ xor_bits) << Int32(1)) + row_high | ||
| dst_flat = dst_row * packed_cols + dst_pcol | ||
| byte_val = Uint8( | ||
| (packed64 >> Uint64(byte_idx * 8)) & Uint64(0xFF) | ||
| ) | ||
| sA_u8[dst_flat] = byte_val | ||
|
|
||
| outer_m_idx = row % Int32(32) | ||
| inner_m_idx = row // Int32(32) | ||
| inner_k_idx = sf_block % Int32(4) | ||
| k_tile_idx = sf_block // Int32(4) | ||
| sf_raw_idx = ( | ||
| k_tile_idx * Int32(32 * 4 * 4) | ||
| + outer_m_idx * Int32(4 * 4) | ||
| + inner_m_idx * Int32(4) | ||
| + inner_k_idx | ||
| ) | ||
| if defer_sfa > Int32(0): | ||
| deferred_sfa_words[deferred_sfa_slot] = deferred_sfa_words[ | ||
| deferred_sfa_slot | ||
| ] | (Uint32(scale_byte) << Uint32(deferred_slot * Int32(8))) | ||
| else: | ||
| st_shared_u8( | ||
| sfa_base_addr | ||
| + q1_sfa_stage_idx * sfa_stage_elements | ||
| + sf_raw_idx, | ||
| scale_byte, | ||
| ) | ||
| deferred_slot += Int32(1) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
The deferred register buffers assume exactly four blocks per thread.
deferred_a_words holds 8 Uint32 (line 4711) and each deferred block writes 2 words at deferred_slot * 2. deferred_sfa_words holds 2 Uint32 and each deferred block ORs one byte at deferred_slot * 8 bits. Both therefore support at most 4 blocks per thread.
deferred_slot counts every block this thread processes: epi_rows * sf_blocks_per_row blocks, strided by num_mma_warps * num_threads_per_warp = 256. With tile_shape_mnk = (128, 128, 128) and sf_vec_size = 16 that is 128 * 8 / 256 = 4, which fits exactly. The constructor constrains mma_tiler_mn[1] to 128 but places no constraint on mma_tiler_mn[0]. With mma_tiler_mn[0] = 256 the count doubles to 8, deferred_a_words[deferred_slot * 2] indexes past element 7, and the SFA shift reaches 56 bits on a Uint32.
Derive both buffer sizes from the tile shape, or add a constructor check that pins mma_tiler_mn[0] and sf_vec_size to the supported values.
🛡️ Proposed constructor guard
if self.fc1_sfb_tiles_per_block != 2:
raise ValueError("expected exactly two logical N64 tiles per SFB block")
+ # The deferred Q1 register buffers hold exactly four blocks per thread.
+ blocks_per_thread = (
+ mma_tiler_mn[0] * (tile_k // _SF_VEC_SIZE)
+ ) // (8 * 32)
+ if blocks_per_thread != 4:
+ raise ValueError(
+ "the gated dynamic kernel requires four Q1 blocks per thread; "
+ f"got {blocks_per_thread}"
+ )🤖 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_dynamic/gated.py` around
lines 2234 - 2272, Update the constructor validation for the dynamic MoE kernel
to enforce the supported deferred-buffer geometry: constrain the relevant
mma_tiler_mn[0] and sf_vec_size values so each thread processes at most four
deferred blocks. Anchor the change in the constructor and preserve the existing
deferred_a_words and deferred_sfa_words indexing in the processing logic.
| scatter_total_u32 = num_tokens * cols_u32 | ||
| scatter_vecs = scatter_total_u32 // Int32(4) | ||
| zero_u32 = Uint32(0) | ||
| zv = flat_tid | ||
| while zv < scatter_vecs: | ||
| st_global_v4_u32( | ||
| scatter_base + Int64(zv) * Int64(16), | ||
| zero_u32, | ||
| zero_u32, | ||
| zero_u32, | ||
| zero_u32, | ||
| ) | ||
| zv += flat_stride | ||
|
|
||
| j = scatter_vecs * Int32(4) + flat_tid | ||
| while j < scatter_total_u32: | ||
| scatter_output_u32[j // cols_u32, j % cols_u32] = Uint32(0) | ||
| j += flat_stride |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check how scatter_output is allocated and whether contiguity is enforced.
rg -nP --type=py -C6 'scatter_output' -g '!**/_moe_dynamic/**'
rg -nP --type=py -C3 'is_contiguous|\.contiguous\(\)' -g 'flashinfer/fused_moe/cute_dsl/**' | head -30Repository: flashinfer-ai/flashinfer
Length of output: 162
🏁 Script executed:
#!/bin/bash
set -u
echo "Locate gated.py and related _moe_dynamic files"
fd -a 'gated\.py$|_moe_dynamic|blackwell_sm12x' . | sed 's#^\./##' | head -200
echo
echo "Search scatter_output broadly"
rg -n --glob '!**/.git/**' 'scatter_output' . | head -200
echo
echo "Search allocation / contiguity checks broadly"
rg -n --glob '!**/.git/**' 'is_contiguous|contiguous\(\)|torch\.empty|torch\.zeros|new_empty|reshape\(|stride\(' flashinfer/fused_moe -S | head -200 || true
echo
echo "If gated.py exists, inspect outline and relevant lines"
FILE="$(fd 'gated\.py$' . | head -1 || true)"
if [ -n "${FILE:-}" ]; then
wc -l "$FILE"
ast-grep outline "$FILE" --view compact | head -200 || true
sed -n '1,80p;2640,2745p;2960,3030p' "$FILE"
fiRepository: flashinfer-ai/flashinfer
Length of output: 42156
🏁 Script executed:
#!/bin/bash
set -u
echo "Inspect frontend callers of dynamic gated MoE and output allocation"
sed -n '980,1210p' flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py
echo
sed -n '1420,1535p' flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py
echo
sed -n '1780,1905p' flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py
echo
sed -n '2186,2280p' flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py
echo
sed -n '2540,2760p' flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py
echo
echo "Inspect test allocations around scatter_output"
sed -n '470,510p' tests/moe/test_b12x_fused_moe.py
sed -n '600,620p' tests/moe/test_b12x_fused_moe.py
echo
echo "Inspect docs/API around contiguous workspace/output"
sed -n '1050,1140p' flashinfer/fused_moe/core.py
sed -n '950,990p' flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py
echo "Behavioral probe: PyTorch strides for contiguous and non-contiguous tensors"
python3 - <<'PY'
import torch
for shape in [(3, 2), (3, 8)]:
c = torch.empty(shape, dtype=torch.bfloat16)
nc = c[:, ::2]
print(shape, "contiguous_st", c.storage_offset(), c.stride(), c.is_contiguous())
print(shape, "noncontig_st", nc.storage_offset(), nc.stride(), nc.is_contiguous())
PYRepository: flashinfer-ai/flashinfer
Length of output: 34864
🏁 Script executed:
#!/bin/bash
set -u
echo "Find scatter_output validation and allocation in moe_dispatch public paths"
rg -n --glob 'flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py' -C5 'scatter_output|data_ptr\(|check_' flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py || true
echo
echo "Inspect allocation helpers and public wrapper around sm120 MoE"
rg -n --glob '**/moe_dispatch.py|**/b12x_moe.py' -C8 'Sm120DynamicMoEWorkspace|sm120_moe|scatter_output|is_contiguous|contiguous\(\)|torch\.empty' flashinfer/fused_moe/cute_dsl/blackwell_sm12x flashinfer/fused_moe -S || trueRepository: flashinfer-ai/flashinfer
Length of output: 11020
Enforce and document scatter_output contiguity / alignment.
The dynamic launcher passes scatter_output.data_ptr() while the kernel makes a [num_tokens, k] CuTe tensor with default row stride (self._k, 1). Existing public callers may pass non-contiguous buffers such as .view(-1), so either validate/contiguous the tensor before launch or add a documented contract requiring a contiguous, 16-byte-aligned output tensor.
🤖 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_dynamic/gated.py` around
lines 3004 - 3021, Ensure the dynamic launcher’s scatter_output input is
contiguous and 16-byte aligned before passing scatter_output.data_ptr() to the
kernel, either by validating and rejecting invalid tensors or materializing a
contiguous tensor and preserving the expected output behavior. Document this
requirement at the relevant public API, and keep the kernel’s default row-major
[num_tokens, k] layout consistent with the enforced contract.
| route_gs = cute.make_rmem_tensor((16,), cutlass.Float32) | ||
| cache_slot = Int32(0) | ||
| while cache_slot < num_topk: | ||
| route_slot = route_slot_base + cache_slot | ||
| expert_id = _ld_shared_i32( | ||
| route_expert_ids_addr + route_slot * Int32(4) | ||
| ) | ||
| gs_value = input_global_scale[expert_id].to(cutlass.Float32) | ||
| if ( | ||
| self.input_scales_are_reciprocal | ||
| and gs_value != cutlass.Float32(0.0) | ||
| ): | ||
| if self.fast_math: | ||
| gs_value = rcp_approx_ftz(gs_value) | ||
| else: | ||
| gs_value = cutlass.Float32(1.0) / gs_value | ||
| route_gs[cache_slot] = gs_value | ||
| cache_slot += Int32(1) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Look for a host-side top-k validation on the SM120 dynamic path.
rg -nP --type=py -C5 'topk|num_topk' -g 'flashinfer/fused_moe/cute_dsl/**' | rg -nP -C3 'raise|assert|<=|>' | head -60Repository: flashinfer-ai/flashinfer
Length of output: 162
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Locate file and snippets"
wc -l flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py
sed -n '2960,3540p' flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py
sed -n '3540,3630p' flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py
echo
echo "Search route_gs definitions/usages and any topk validation"
rg -n 'route_gs|num_topk=|total_pairs|topk_ids|top-k|top_k|topk|assert|max|min|raise' flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic flashinfer/fused_moe -g '*.py' | head -200Repository: flashinfer-ai/flashinfer
Length of output: 50380
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Focused route_gs read/write range"
sed -n '3523,3645p' flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py
echo
echo "topK argument/docs around dynamic fused_moe APIs"
rg -n --type=py -C3 'top_k|topk_ids|topk_weights|num_topk' flashinfer/flashinfer fused_moe -g '*.py' | head -240Repository: flashinfer-ai/flashinfer
Length of output: 7380
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Dynamic gated API path"
rg -n --type=py -C3 'def .*fused.*moe|def .*gated|top_k|topk_ids|route_phys_rows|share_input_across_experts|num_topk' flashinfer/fused_moe -g '*.py' | head -260
echo
echo "Host/input validation searches"
rg -n --type=py -C5 'raise ValueError|assert .*topk|topk_ids\.shape|top_k|num_topk|ValueError' flashinfer/fused_moe -g '*.py' | head -300Repository: flashinfer-ai/flashinfer
Length of output: 43256
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Blackwell gated Python frontend"
fd -a 'gated.*\.py' flashinfer/fused_moe | sed 's#^\./##' | tr '\n' ' '
echo
for f in flashinfer/fused_moe/gated.py flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py; do
if [ -f "$f" ]; then
echo "=== $f ==="
wc -l "$f"
rg -n --type=py -C5 'def .*gated|top_k|topk_ids|route_gs|num_topk|launch_params|storage|check_support|ValueError|assert' "$f" | head -260
else
echo "missing $f"
fi
done
echo
echo "Search route_gs references in repository files only"
rg -n 'route_gs|route_gs_value|route_scale|route_phys_rows|route_expert_ids' . -g '*.py' -g '*.cu' -g '*.cuh' -g '*.jinja' | head -220Repository: flashinfer-ai/flashinfer
Length of output: 12127
Enforce the 16-entry top-k cache contract.
route_gs is allocated with 16 Float32 entries, but num_topk = total_pairs // num_tokens is runtime data. Routes >16 loop past this tensor when loading scales and again when reading them before quantization. Add a host-side check that rejects top_k > 16; a kernel-side clamp is not a valid fix because it changes routing behavior.
🤖 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_dynamic/gated.py` around
lines 3506 - 3523, Add a host-side validation wherever the runtime
num_topk/top_k is established, rejecting values greater than 16 before launching
the kernel. Preserve routing behavior by failing explicitly rather than
clamping, and ensure the validation covers both the route_gs population and
subsequent scale-read paths in the relevant MoE setup flow.
| self.a_dtype = packed_a.element_type | ||
| self.b_dtype = b_w13.element_type | ||
| self.sf_dtype = sfa_ptr.dtype | ||
| self.a_layout = utils.LayoutEnum.from_tensor(packed_a) | ||
| self.b_layout = utils.LayoutEnum.from_tensor(b_w13) | ||
| # Dynamic never materializes the intermediate C tensor. Preserve the | ||
| # original row-major epilogue layout without carrying a dead memref. | ||
| self.c_layout = utils.LayoutEnum.ROW_MAJOR | ||
|
|
||
| hidden_size = a_input.shape[1] | ||
| self._setup_attributes(hidden_size=hidden_size) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Find where the kernel instance is created and whether it is cached or shared.
rg -nP --type=py -C8 'MoEGatedDynamicKernel|MoEDynamicKernel' -g '!**/_moe_dynamic/**'
rg -nP --type=py -C4 'functools\.cache|lru_cache' -g 'flashinfer/fused_moe/cute_dsl/**' | head -40Repository: flashinfer-ai/flashinfer
Length of output: 162
🏁 Script executed:
#!/bin/bash
set -u
echo "== repo files top-level =="
git ls-files | sed -n '1,80p'
echo "== locate relevant files =="
fd -a 'gated.py|.*dynamic.*|.*moe.*kernel.*|.*cute.*dsl.*' . | sed 's#^\./##' | sed -n '1,120p'
echo "== search symbols with precise strings =="
rg -n --type=py -C 10 'class .*MoEGatedDynamicKernel|class .*MoEDynamicKernel|MoEGatedDynamicKernel|MoEDynamicKernel|CuteDslMoEWrapper|gated_dynamic|moe_dynamic' . | sed -n '1,240p'
echo "== inspect target file outline/size =="
if [ -f flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py ]; then
wc -l flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py
ast-grep outline flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py --match 'MoEGatedDynamicKernel' --view expanded || true
sed -n '3800,3940p' flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py
else
echo "target file missing"
fi
echo "== inspect _setup_attributes if present =="
rg -n --type=py -C 8 '_setup_attributes|self\.tiled_mma|self\._hidden_size' flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py | sed -n '1,260p'Repository: flashinfer-ai/flashinfer
Length of output: 46964
🏁 Script executed:
#!/bin/bash
set -u
echo "== git status/stat =="
git status --short
git diff --stat || true
echo "== find all py files with exact identifiers outside target dir =="
rg -n --type=py 'MoEGatedDynamicKernel|MoEDynamicKernel|CuteDslMoEWrapper|cute_dsl' . | sed -n '1,240p'
echo "== list tracked paths under flashinfer fused_moe =="
git ls-files 'flashinfer/fused_moe/*' | sed -n '1,200p'Repository: flashinfer-ai/flashinfer
Length of output: 27066
🏁 Script executed:
#!/bin/bash
set -u
echo "== inspect fused_moe.py relevant classes/functions =="
fd -a 'fused_moe.py' flashinfer/fused_moe/cute_dsl | xargs -r wc -l
ast-grep outline flashinfer/fused_moe/cute_dsl/fused_moe.py --view expanded | sed -n '1,220p'
sed -n '1,260p' flashinfer/fused_moe/cute_dsl/fused_moe.py
sed -n '260,660p' flashinfer/fused_moe/cute_dsl/fused_moe.py
sed -n '660,1100p' flashinfer/fused_moe/cute_dsl/fused_moe.py
echo "== inspect run_sm12x helper and wrappers =="
wc -l flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_host.py
ast-grep outline flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_host.py --view expanded | sed -n '1,220p'
rg -n --type=py -C 12 'run_sm12x|def .*sm12x|MoEDynamicKernel\\(' flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_host.py | sed -n '1,260p'
sed -n '1,260p' flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_host.py
echo "== inspect tests around instance reuse/shared kernel =="
rg -n --type=py -C 10 'CuteDslMoEWrapper|cute_dsl_fused_moe_nvfp4|launch_sm120_dynamic_moe|MoE|cache|reuse|stream' tests benchmark* | sed -n '1,260p'
echo "== inspect cache implementation =="
fd -a 'cute_dsl_core.py|tuner.py' flashinfer flashinfer/fused_moe/cute_dsl | xargs -r wc -l
ast-grep outline flashinfer/jit/cute_dsl_core.py --view expanded | sed -n '1,180p'
sed -n '180,340p' flashinfer/jit/cute_dsl_core.pyRepository: flashinfer-ai/flashinfer
Length of output: 50380
🏁 Script executed:
#!/bin/bash
set -u
echo "== cute_dsl fused_moe tuner outline relevant =="
ast-grep outline flashinfer/fused_moe/cute_dsl/tuner.py --view expanded | sed -n '1,220p'
wc -l flashinfer/fused_moe/cute_dsl/tuner.py
sed -n '1,220p' flashinfer/fused_moe/cute_dsl/tuner.py
sed -n '220,380p' flashinfer/fused_moe/cute_dsl/tuner.py
sed -n '380,700p' flashinfer/fused_moe/cute_dsl/tuner.py
echo "== precise MoEDynamicKernel dispatch call =="
rg -n --type=py -C 18 'kernel: Any = MoEDynamicKernel\(' flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py
sed -n '1560,1680p' flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py
echo "== precise cute_dsl_fused_moe_nvfp4 implementation call =="
rg -n --type=py -C 20 'def blockscaled_contiguous_gather_grouped_gemm_act_fusion_nvfp4' flashinfer/fused_moe/cute_dsl/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py
sed -n '1,260p' flashinfer/fused_moe/cute_dsl/blockscaled_contiguous_gather_grouped_gemm_act_fusion.pyRepository: flashinfer-ai/flashinfer
Length of output: 43162
Derive per-call kernel state locally instead of mutating self.
__call__ sets a_dtype, b_dtype, sf_dtype, a_layout, b_layout, and c_layout, then self._setup_attributes assigns the tiled MMA/SME layouts. The MoE dispatch caches MoEDynamicKernel by config, so calls with different input dtypes/scaling share the object; concurrent reuse can interleave these writes and expose mixed state to the compiled kernel. Move these values into a local configuration object or use a thread-safe cache per input configuration.
🤖 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_dynamic/gated.py` around
lines 3889 - 3899, Update MoEDynamicKernel.__call__ to derive a per-invocation
configuration containing a_dtype, b_dtype, sf_dtype, a_layout, b_layout,
c_layout, and the layouts produced by _setup_attributes instead of mutating
shared self state. Ensure compiled-kernel generation and dispatch consume only
that local configuration, so cached kernels remain safe across differing inputs
and concurrent calls.
Source: Learnings
| ): | ||
| """Kernel entry point.""" | ||
| tidx, _, _ = cute.arch.thread_idx() | ||
| bidx, _, bidz = cute.arch.block_idx() |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Prefix the unused bidx with an underscore.
Ruff reports RUF059 for this unpacking. Only bidz is used.
🔧 Proposed fix
- bidx, _, bidz = cute.arch.block_idx()
+ _bidx, _, bidz = cute.arch.block_idx()📝 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.
| bidx, _, bidz = cute.arch.block_idx() | |
| _bidx, _, bidz = cute.arch.block_idx() |
🧰 Tools
🪛 Ruff (0.16.0)
[warning] 4095-4095: Unpacked variable bidx is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
🤖 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_dynamic/gated.py` at line
4095, Update the block index unpacking in the surrounding function to bind the
unused first value to an underscore instead of bidx, while preserving bidz for
subsequent use.
Source: Linters/SAST tools
| cute.recast_tensor(sA, cutlass.Uint8) | ||
| cute.recast_tensor(sB, cutlass.Uint8) | ||
| cute.recast_tensor(sB_phase2_extra, cutlass.Uint8) | ||
| cute.recast_tensor(sB_fc1_all, cutlass.Uint8) | ||
| cute.recast_tensor(sB_fc1, cutlass.Uint8) | ||
| cute.recast_tensor(sB_up_fc1, cutlass.Uint8) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Inspect recast_tensor in the installed CUTLASS DSL and check repository usage patterns.
python - <<'PY'
import inspect
try:
import cutlass.cute as cute
print(inspect.signature(cute.recast_tensor))
print((inspect.getdoc(cute.recast_tensor) or "")[:800])
except Exception as exc:
print("unavailable:", exc)
PY
# Compare with other call sites that keep the result.
rg -nP --type=py -C1 'recast_tensor\(' -g 'flashinfer/**' | head -40Repository: flashinfer-ai/flashinfer
Length of output: 201
🌐 Web query:
nvidia-cutlass-dsl cute.recast_tensor return value semantics
💡 Result:
In NVIDIA CUTLASS, specifically within the CuTe DSL, the cute.recast_tensor function returns a new tensor object that points to the same underlying memory as the original input tensor, but reinterprets that memory as the specified target data type [1][2]. The key semantic behaviors of cute.recast_tensor are as follows: 1. Memory Interpretation: The function changes the element interpretation of the tensor [1][2]. It effectively reinterprets the memory buffer using the new data type, adjusting both the iterator pointer and the tensor's layout to ensure the data is accessed consistently with the new type [1][2]. 2. Return Value: It returns a new instance of the original tensor type (e.g., Tensor), constructed with the re-calculated iterator and layout [1][2]. 3. Relationship: The new tensor and the original tensor share the same physical memory, but they provide different logical views of that memory due to the change in data type [1][2]. 4. C++ Context: In the C++ header implementation (cute/tensor_impl.hpp), the recast operation (implemented as the recast function) performs a similar role, recomputing the layout to match the new type [3][4]. It includes warnings that this is a dangerous operation, as it does not inherently check for dynamic integer divisibility or memory alignment requirements [3][4]. Note that in the Python DSL, if certain conditions (like the presence of a stride-1 dimension) are not met, the operation may behave as a no-op or return the original tensor structure depending on the specific layout configuration [5].
Citations:
- 1: https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/cutlass/cute/tensor.py
- 2: https://github.com/NVIDIA/cutlass/blob/f74fea9c/python/CuTeDSL/cutlass/cute/tensor.py
- 3: https://github.com/NVIDIA/cutlass/blob/main/include/cute/tensor_impl.hpp
- 4: https://github.com/NVIDIA/cutlass/blob/f74fea9c/include/cute/tensor_impl.hpp
- 5: [QST] [CuTeDSL]
cast_tensorcorner case when no stride = 1 NVIDIA/cutlass#3026
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== target context =="
sed -n '4230,4325p' flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py | cat -n
echo
echo "== all recast_tensor call sites =="
rg -nP --type=py -C2 'recast_tensor\(' -g 'flashinfer/**' || true
echo
echo "== dependency/version references =="
rg -n "nvidia-cutlass-dsl|cutlass|cute" -g 'pyproject.toml' -g 'requirements*.txt' -g 'setup.py' -g 'setup.cfg' . || trueRepository: flashinfer-ai/flashinfer
Length of output: 5552
🌐 Web query:
site:github.com NVIDIA cutlass recast_tensor implementation cute Python cute recast_tensor
💡 Result:
In NVIDIA CUTLASS, recast_tensor is a utility function used to reinterpret the memory of an existing tensor with a different data type [1][2]. It effectively changes the element interpretation of a tensor without copying or modifying the underlying data, acting similarly to a reinterpret_cast in C++ [3][4]. Implementation and Functionality: The recast_tensor operation is implemented in both the CUTLASS C++ template library (as recast) [3] and the Python CuTe DSL (as recast_tensor) [1][2]. Python DSL Implementation: The Python function recast_tensor(src, dtype, ...) takes a source Tensor and a target dtype [1][2]. It performs the following steps: 1. Calculates the element width (in bits) for both the source and target data types [1][2]. 2. Updates the iterator pointer type using recast_ptr [1][2]. 3. Adjusts the tensor layout using recast_layout to ensure consistency with the new element size [1][2]. 4. Returns a new Tensor object that shares the same underlying memory but interprets it according to the new dtype [1][2]. Technical Considerations and Risks: - Memory Consistency: recast_tensor is considered an inherently dangerous operation [3][4]. It does not perform automatic checks for alignment or dynamic integer divisibility [3][4]. If the preconditions regarding the tensor's layout and memory alignment are not met, the operation may produce incorrect results or lead to illegal memory access (IMA) [5][6][4]. - Layout Requirements: For recast_tensor to function correctly, the tensor layout typically must have a stride of 1 in the dimensions being recast, or meet specific divisibility requirements [5][6]. In cases where these conditions are not satisfied (e.g., no stride is 1), the operation may result in a no-op or incorrect address computation [5][6]. - Use Cases: It is frequently used when frameworks do not natively support narrow precision (e.g., when converting between 32-bit and 16-bit representations) or when manipulating bits at a lower level [1][7][8]. It is also employed internally by higher-level mechanisms like TiledCopy [4].
Citations:
- 1: https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/cutlass/cute/tensor.py
- 2: https://github.com/NVIDIA/cutlass/blob/f74fea9c/python/CuTeDSL/cutlass/cute/tensor.py
- 3: https://github.com/NVIDIA/cutlass/blob/main/include/cute/tensor_impl.hpp
- 4: [BUG] CuTe: Inconsistent results when using a dynamically vs statically defined layouts and recast. NVIDIA/cutlass#2132
- 5: [QST] [CuTeDSL]
cast_tensorcorner case when no stride = 1 NVIDIA/cutlass#3026 - 6: [Cute] Fix shuffle sync and enable pack gqa for varlen sm100 Dao-AILab/flash-attention#2097
- 7: https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/cutlass/cute/runtime.py
- 8: https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn%2Fcute%2Futils.py
🌐 Web query:
"def recast_tensor" "cutlass.cute" cutlass/cute/tensor.py
💡 Result:
In the CUTLASS CuTe Python DSL, recast_tensor is a function used to change the data type of a tensor [1][2]. It allows you to reinterpret the elements of a tensor as a different type, which is useful for type conversions [2]. Usage: new_tensor = cute.recast_tensor(tensor, dtype) Parameters: tensor: The source tensor to be recasted [1]. dtype: The new element data type (e.g., cute.Float16, cute.Float32) [1][2]. Key details: - It returns a new tensor view with the specified data type [1][2]. - The operation relies on the underlying tensor layout; specifically, it has been noted in community discussions that the operation may behave as a no-op if the tensor's layout does not include a stride of 1, as it relies on the ability to access contiguous elements for the recasting to be meaningful [3]. - It is part of the broader CuTe DSL API for memory operations [4][2]. Note: Some earlier internal documentation or code comments have referred to recast_tensor as being part of an evolving API story for narrow precision handling, but it is now a standard part of the public CuTe DSL API for type-casting operations [5][2].
Citations:
- 1: https://docs.nvidia.com/cutlass/4.3.5/media/docs/pythonDSL/cute_dsl_api/cute.html
- 2: https://mintlify.wiki/NVIDIA/cutlass/api/dsl/memory-operations
- 3: [QST] [CuTeDSL]
cast_tensorcorner case when no stride = 1 NVIDIA/cutlass#3026 - 4: https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/cute_dsl_api/cute.html
- 5: https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/cutlass/cute/runtime.py
Assign the recast_tensor return values to avoid stale dtype views.
cute.recast_tensor(...) returns a new tensor with the target dtype reinterpretation. Store each result and pass the assigned tensor to later GPU operations so those lanes do not keep reading the original dtype view.
🤖 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_dynamic/gated.py` around
lines 4270 - 4275, Update the tensor recasting block to assign each
cute.recast_tensor result back to its corresponding tensor variable, including
sA, sB, sB_phase2_extra, sB_fc1_all, sB_fc1, and sB_up_fc1, so subsequent GPU
operations use the Uint8 views rather than stale dtype views.
| q1_a_stage_idx = Int32(3) | ||
| defer_a = Int32(0) | ||
| if slice_idx == Int32(1): | ||
| q1_a_stage_idx = Int32(4) | ||
| elif slice_idx == Int32(2): | ||
| q1_a_stage_idx = Int32(0) | ||
| defer_a = Int32(1) | ||
| elif slice_idx == Int32(3): | ||
| q1_a_stage_idx = Int32(1) | ||
|
|
||
| q1_sfa_stage_idx = Int32(3) | ||
| defer_sfa = Int32(0) | ||
| deferred_sfa_slot = Int32(0) | ||
| if slice_idx == Int32(1): | ||
| q1_sfa_stage_idx = Int32(0) | ||
| defer_sfa = Int32(1) | ||
| elif slice_idx == Int32(2): | ||
| q1_sfa_stage_idx = Int32(0) | ||
| defer_sfa = Int32(1) | ||
| deferred_sfa_slot = Int32(1) | ||
| elif slice_idx == Int32(3): | ||
| q1_sfa_stage_idx = Int32(0) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Determine the intermediate sizes and w13 shapes exercised on the SM120 dynamic path.
rg -nP --type=py -C6 'intermediate_size|I_tp' -g 'flashinfer/fused_moe/cute_dsl/**' | head -60
# Find tests that cover this kernel and the shapes they use.
rg -nP --type=py -C6 'dynamic' -g 'tests/moe/**' | head -60Repository: flashinfer-ai/flashinfer
Length of output: 162
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate the target file and surrounding publication/read logic without reading the whole file.
target='flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py'
wc -l "$target"
echo '--- publication variables and methods outline/search ---'
rg -n "publish_variable_deferred_tasks|task_slice_count|slice_chunk|gate_tile_cnt|_TASK_SLICE_CHUNK|slice_count|q1_a_stage_idx|q1_sfa_stage_idx|storageGated|w13|intermediate_size|num_tokens|bidx|fc1_storage_alias|sequential_branch_compact" "$target"
echo '--- sections around publisher ---'
sed -n '1050,1145p' "$target" | nl -ba -v1050
echo '--- sections around dynamic kernel host bounds ---'
sed -n '3720,3790p' "$target" | nl -ba -v3720
echo '--- sections around stage maps ---'
sed -n '4785,4925p' "$target" | nl -ba -v4785
echo '--- target ranges 4265-4305 and 4811-4833 ---'
sed -n '4265,4305p' "$target" | nl -ba -v4265
sed -n '4811,4833p' "$target" | nl -ba -v4811Repository: flashinfer-ai/flashinfer
Length of output: 13004
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the exact code around the targeted publication logic and FC2 read stage maps.
target='flashinfer/fuse
d_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py'
if [ -f "$target" ]; then
sed -n '1090,1125p' "$target" | nl -ba -v1090
sed -n '4795,4920p' "$target" | nl -ba -v4795
fiRepository: flashinfer-ai/flashinfer
Length of output: 162
🏁 Script executed:
#!/bin/bash
set -euo pipefail
target='flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py'
echo '--- lines 1,105-125 ---'
python3 - <<'PY'
target='flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py'
for start,end in [(1,130),(105,125)]:
print(f'--- {target}:{start}-{end} ---')
with open(target) as f:
for i,line in enumerate(f,1):
if start <= i <= end:
print(f'{i}: {line}', end='')
PY
echo '--- lines 1090-1117 ---'
python3 - <<'PY'
target='flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py'
with open(target) as f:
for i,line in enumerate(f,1):
if 1090 <= i <= 1117:
print(f'{i}: {line}', end='')
PY
echo '--- lines 2986-3006 ---'
python3 - <<'PY'
target='flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py'
with open(target) as f:
for i,line in enumerate(f,1):
if 2986 <= i <= 3006:
print(f'{i}: {line}', end='')
PY
echo '--- lines 3716-3796 ---'
python3 - <<'PY'
target='flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py'
with open(target) as f:
for i,line in enumerate(f,1):
if 3716 <= i <= 3796:
print(f'{i}: {line}', end='')
PY
echo '--- lines 3956-3986 ---'
python3 - <<'PY'
target='flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py'
with open(target) as f:
for i,line in enumerate(f,1):
if 3956 <= i <= 3986:
print(f'{i}: {line}', end='')
PY
echo '--- lines 4785-4838 ---'
python3 - <<'PY'
target='flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py'
with open(target) as f:
for i,line in enumerate(f,1):
if 4785 <= i <= 4838:
print(f'{i}: {line}', end='')
PY
echo '--- lines 4890-4918 ---'
python3 - <<'PY'
target='flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py'
with open(target) as f:
for i,line in enumerate(f,1):
if 4890 <= i <= 4918:
print(f'{i}: {line}', end='')
PY
echo '--- deterministic model of publication / slice count bounds ---'
python3 - <<'PY'
TASK_SLICE_CHUNK = 4
def publish_variable_deferred_tasks(gate_tile_cnt, split_tile_count=0):
num_groups = (gate_tile_cnt + (1 + 1 - 1)) // 1
tasks = []
for g in range(num_groups):
if g < split_tile_count:
slice_chunk = 2
num_groups_split = (gate_tile_cnt + (split_tile_count + 1 - 1)) // split_tile_count
slice_begin = g * slice_chunk
slice_count = gate_tile_cnt - slice_begin
if slice_count > slice_chunk:
slice_count = slice_chunk
tasks.append(slice_count)
else:
slice_chunk = gate_tile_cnt
num_groups_split = (gate_tile_cnt + (1 + 1 - 1)) // 1
slice_begin = g * slice_chunk
slice_count = gate_tile_cnt - slice_begin
if slice_count > slice_chunk:
slice_count = slice_chunk
tasks.append(slice_count)
return tasks
def publish_other(gate_tile_cnt, num_tokens):
if num_tokens <= 2048:
task_slice_chunk = 2
else:
task_slice_chunk = TASK_SLICE_CHUNK
tasks = []
num_groups = max(1, gate_tile_cnt) // task_slice_chunk
if 0 == 0:
num_groups += 1
for g in range(num_groups):
slice_begin = g * task_slice_chunk
slice_count = gate_tile_cnt - slice_begin
if slice_count > task_slice_chunk:
slice_count = task_slice_chunk
tasks.append(slice_count)
return tasks
def max_task_slice_count(gate_tile_cnt, num_tokens):
if num_tokens > 4096:
return TASK_SLICE_CHUNK
elif num_tokens > 2048:
split_tile_count = (gate_tile_cnt + 1) // 2
return max(publish_variable_deferred_tasks(gate_tile_cnt, split_tile_count))
elif num_tokens > 256:
return max(publish_variable_deferred_tasks(gate_tile_cnt, 0))
else:
return publish_other(gate_tile_cnt, num_tokens)[0]
for shape_n in [1, 64, 320, 640, 641, 1280, 2048, 4096, 8192, 16384]:
tiles = shape_n // 64 // 2
print(f'shape_n={shape_n} gate_tile_cnt={tiles} max_count_4096={max_task_slice_count(tiles, 4096)} max_count_2049={max_task_slice_count(tiles, 2049)}')
print(' tasks_4096=', publish_variable_deferred_tasks(tiles, ((tiles + 1) // 2) if num_tokens > 2048 else (0)))
PYRepository: flashinfer-ai/flashinfer
Length of output: 17469
🏁 Script executed:
#!/bin/bash
set -euo pipefail
target='flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py'
python3 - <<'PY'
gated='flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py'
with open(gated) as f:
txt=f.read().splitlines()
for start,end in [(1,130),(105,125),(1090,1117),(2986,3006),(3716,3796),(3956,3986),(4785,4838),(4890,4918)]:
print(f'--- {gated}:{start}-{end} ---')
for i in range(start,end+1):
print(f'{i}: {txt[i-1]}')
PY
python3 - <<'PY'
TASK_SLICE_CHUNK = 4
def publish_variable_deferred_tasks(gate_tile_cnt, split_tile_count=0):
num_groups = (gate_tile_cnt + (1 + 1 - 1)) // 1
tasks = []
for g in range(num_groups):
if g < split_tile_count:
slice_chunk = 2
num_groups_split = (gate_tile_cnt + (split_tile_count + 1 - 1)) // split_tile_count if split_tile_count > 0 else gate_tile_cnt
slice_begin = g * slice_chunk
slice_count = gate_tile_cnt - slice_begin
if slice_count > slice_chunk:
slice_count = slice_chunk
tasks.append(slice_count)
else:
slice_chunk = gate_tile_cnt
num_groups_split = (gate_tile_cnt + (1 + 1 - 1)) // 1
slice_begin = g * slice_chunk
slice_count = gate_tile_cnt - slice_begin
if slice_count > slice_chunk:
slice_count = slice_chunk
tasks.append(slice_count)
return tasks
for shape_n in [1, 64, 320, 640, 641, 1280, 2048, 4096, 8192, 16384]:
tiles = shape_n // 64 // 2
print(f'shape_n={shape_n} gate_tile_cnt={tiles} task_slice_count={publish_variable_deferred_tasks(tiles)}')
PYRepository: flashinfer-ai/flashinfer
Length of output: 19814
🏁 Script executed:
#!/bin/bash
set -euo pipefail
target='flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py'
python3 - <<'PY'
from pathlib import Path
text = Path('flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py').read_text()
lines = text.splitlines()
print('--- lines 1115-1138 ---')
for i in range(1115, 1138+1):
print(f'{i}: {lines[i-1]}')
# Read-only, static text extraction of the publication loop without executing repository code.
# This is a simplified model: it extracts only assignments/constants from the method source
# and prints whether `slice_chunk` is the remainder expression or the constant `Int32(2)`.
src_start = text.index(' def publish_variable_deferred_tasks(')
src_end = text.index(' def publish_uniform_deferred_tasks(', src_start)
body = text[src_start:src_end]
print('--- method contains slice_chunk assignments ---')
for line in body.splitlines():
if 'slice_chunk' in line:
print(line)
def publish_variable_deferred_tasks(gate_tile_cnt, split_tile_count):
split_groups = (gate_tile_cnt + 1) // 2
extra_per_split = split_groups - 1
split_tiles_before = split_tile_count
if split_tiles_before > split_tile_count:
split_tiles_before = split_tile_count
start = split_tiles_before + split_tiles_before * extra_per_split
tasks = []
num_groups = 1
slice_chunk = gate_tile_cnt
g = 0
while g < num_groups:
if g < split_tile_count:
slice_chunk = 2
slice_begin = g * slice_chunk
slice_count = gate_tile_cnt - slice_begin
if slice_count > slice_chunk:
slice_count = slice_chunk
tasks.append(slice_count)
g += 1
return tasks
for N in [320, 640, 641, 1280, 2048, 4096, 8192]:
gate_tile_cnt = N // 64 // 2
for num_tokens in [4096, 2049, 257]:
if num_tokens <= 256:
continue
split_tile_count = 0
if num_tokens <= 4096:
if gate_tile_cnt % 2 != 0:
extra_per_split = 0
else:
extra_per_split = (gate_tile_cnt + 1) // 2 - 1
if extra_per_split > 0:
target_task_count = 128 if num_tokens <= 2048 else 0
if num_tokens > 2048:
target_task_count = (125 + 31) // 32
split_tile_count = max(0, (target_task_count + extra_per_split - 1) // extra_per_split)
if split_tile_count > 1:
split_tile_count = 1
print(f'N={N} gate_tile_cnt={gate_tile_cnt} num_tokens={num_tokens} split_tile_count={split_tile_count} tasks={publish_variable_deferred_tasks(gate_tile_cnt, split_tile_count)[-num_tokens%4+1:][:10]}')
PYRepository: flashinfer-ai/flashinfer
Length of output: 1208
Cap slice_chunk in the variable deferred publisher.
publish_variable_deferred_tasks sets slice_chunk = gate_tile_cnt for the final group, while the FC2 read paths only define stage maps for slice_idx 0–3. If slice_count also wraps at gate_tile_cnt, later slices can reuse slice 0’s sA/sSFA stages. Limit this path to the same slice chunk cap used by the uniform publisher, or bound the stages by number-of-slices arithmetic.
🤖 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_dynamic/gated.py` around
lines 4811 - 4832, Update publish_variable_deferred_tasks so its final-group
slice_chunk is capped to the same maximum used by the uniform publisher, keeping
slice_idx within the defined 0–3 FC2 stage mappings and preventing later slices
from reusing slice 0 stages.
| """ | ||
| MoEDynamicKernel — queue-driven routed NVFP4 MoE kernel for SM120/SM121. | ||
|
|
||
| Ported from the b12x kernel library to FlashInfer. | ||
|
|
||
| This is the first dynamic fused control-plane kernel derived from the current | ||
| static implementation. It keeps the proven FC1 / activation / quant / FC2 / | ||
| scatter compute body, but replaces the resident-grid route/pack -> compute | ||
| barrier with a global ready-task queue. | ||
|
|
||
| Supports two activation modes selected at construction time: | ||
| SiLU (gated, activation="silu"): | ||
| FC1: A x gate^T, A x up^T (paired FP4 block-scaled GEMMs) | ||
| Act: SiLU(gate) * up (fused SwiGLU activation) | ||
| ReLU2 (non-gated, activation="relu2"): | ||
| FC1: A x W1^T (single FP4 block-scaled GEMM) | ||
| Act: max(0, x)^2 (squared ReLU activation) | ||
|
|
||
| Execution model | ||
| Phase 0: cooperative init / clear scratch state | ||
| Phase 1: all CTAs start as producers | ||
| - claim routed (token, topk_slot) pairs from pair_head | ||
| - append expert rows | ||
| - write token_map + token_weights | ||
| - quantize each routed token row into expert-major packed A + scales | ||
| - publish one compute task per ready (expert, m_tile, slice_group) | ||
| as soon as a tile is fully written | ||
| Phase 2: CTAs that finish producing become consumers immediately | ||
| - CTA leader pops one ready task into shared ctrl state | ||
| - MMA warps run FC1 -> SiLU -> quant -> FC2 -> scatter for that task | ||
| - DMA warp streams the corresponding FC1 / FC2 weights | ||
|
|
||
| This is intentionally conservative: | ||
| - still one CTA per SM | ||
| - still the static per-slice microkernel, now executed sequentially for a | ||
| small grouped slice task | ||
| - still one initial resident-grid barrier after init | ||
|
|
||
| What changes relative to the static path | ||
| - no global route/pack -> compute barrier | ||
| - no static scheduler in the compute steady state | ||
| - route/pack is warp-private instead of CTA-broadcast | ||
| - compute work is driven by a global append-only ready-task queue | ||
|
|
||
| This file is a first implementation pass, not a compiled or profiled artifact. | ||
| It is meant as a concrete CuTeDSL starting point for the next iteration. | ||
| """ |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Update the module docstring to list all supported activations.
__init__ accepts silu, relu2, gelu_tanh, and swigluoai_uninterleave. The docstring documents only SiLU and ReLU2. The closing lines also state that the file is "not a compiled or profiled artifact", which no longer matches the validated implementation.
📝 Proposed docstring update
-Supports two activation modes selected at construction time:
- SiLU (gated, activation="silu"):
+Supports gated and non-gated activation modes selected at construction time:
+ Gated (activation="silu", "gelu_tanh", "swigluoai_uninterleave"):
FC1: A x gate^T, A x up^T (paired FP4 block-scaled GEMMs)
- Act: SiLU(gate) * up (fused SwiGLU activation)
- ReLU2 (non-gated, activation="relu2"):
+ Act: act(gate) * up_term (see moe_activation.gated_activation_f32)
+ Non-gated (activation="relu2"):
FC1: A x W1^T (single FP4 block-scaled GEMM)
Act: max(0, x)^2 (squared ReLU activation)Also remove or revise the final paragraph at Lines 45-46.
As per coding guidelines: "Keep documentation synchronized with code changes, including infrastructure, conventions, deprecated approaches, error-handling patterns, and relevant skill tutorials."
🤖 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_dynamic/generic.py` around
lines 1 - 47, Update the module docstring for MoEDynamicKernel to document all
activation modes accepted by __init__: silu, relu2, gelu_tanh, and
swigluoai_uninterleave, including their relevant behavior. Revise or remove the
closing statement that calls the implementation uncompiled or unprofiled so the
documentation reflects the validated implementation.
Source: Coding guidelines
|
/bot run tests/moe |
There was a problem hiding this comment.
Actionable comments posted: 5
♻️ Duplicate comments (2)
flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py (2)
3415-3432: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winThe 16-entry
route_gscache still has nonum_topkbound.
route_gsholds 16Float32entries, andnum_topk = total_pairs // num_tokens(line 2946) is runtime data. The population loop at line 3417 and the comparison loop at line 3453 both iterate tonum_topk, so a top-k above 16 reads and writes past the register tensor. Neither__call__nor this function rejects that case.Add a host-side rejection for
top_k > 16on this path. Do not clamp in the kernel, because clamping changes routing.Also applies to: 3451-3456
🤖 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_dynamic/gated.py` around lines 3415 - 3432, Add a host-side validation on the affected dynamic gated MoE path, before launching or executing the kernel, that rejects any top_k/num_topk value greater than 16 with an appropriate error. Ensure both route_gs population and comparison remain bounded by valid inputs, and do not clamp the value in the kernel because routing must not change.
3701-3718: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy liftPer-call state is still written to
self.
__call__assignsa_dtype,b_dtype,sf_dtype,a_layout,b_layout, andc_layouton the instance, and_setup_attributesthen writes the tiled MMA and shared-memory layouts.moe_dispatchcaches kernel instances by configuration, so two callers with different input dtypes or scaling share one object. Move this state into a per-invocation configuration object.🤖 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_dynamic/gated.py` around lines 3701 - 3718, The kernel invocation currently stores call-specific dtypes, layouts, and `_setup_attributes` results on the cached instance, allowing callers to overwrite shared state. Introduce a per-invocation configuration object in `__call__`, move `a_dtype`, `b_dtype`, `sf_dtype`, `a_layout`, `b_layout`, `c_layout`, and the tiled MMA/shared-memory attributes initialized by `_setup_attributes` into it, and update downstream dynamic-kernel logic to read from that object instead of `self`.
🧹 Nitpick comments (3)
flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py (2)
3929-3934: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove the dead
sequential_branch_compactandfc1_storage_aliasconfiguration.Neither attribute is ever assigned on
MoEGatedDynamicKernel, so bothgetattrcalls always returnFalse. That makes the alias branches at lines 4012-4015 and 4089-4093 and theup_pipeline.producer_tailcall at lines 4849-4850 unreachable. Delete the flags and the unreachable branches, or set the attributes explicitly in__init__so the intent is visible.🤖 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_dynamic/gated.py` around lines 3929 - 3934, Remove the unused sequential_branch_compact and fc1_storage_alias configuration from MoEGatedDynamicKernel, including their getattr declarations, dependent alias branches, and unreachable up_pipeline.producer_tail call; do not preserve dead configuration paths unless the attributes are explicitly initialized and intentionally supported in __init__.
442-492: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueRemove
load_global_bf16x16_to_f32x16.The global BF16x16 loader is not called anywhere in the repository, while the shared variant is used for Q0 staging. Keep the module surface minimal.
🤖 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_dynamic/gated.py` around lines 442 - 492, Remove the unused load_global_bf16x16_to_f32x16 function, including its inline assembly and result-conversion logic, while leaving the shared Q0 staging loader and other module functionality unchanged.flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dynamic_kernel.py (1)
113-113: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSort
__all__and export_MAX_SHARED_INPUT_TOPK.Ruff reports RUF022 for this line.
moe_dispatch.pyalso imports_MAX_SHARED_INPUT_TOPKfrom this module at Line 35, so include it in the public export list for consistency.♻️ Proposed change
-__all__ = ["MoEDynamicKernel", "_TASK_SLICE_CHUNK"] +__all__ = ["MoEDynamicKernel", "_MAX_SHARED_INPUT_TOPK", "_TASK_SLICE_CHUNK"]🤖 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_dynamic_kernel.py` at line 113, Update the module’s __all__ declaration to include _MAX_SHARED_INPUT_TOPK and sort all exported names alphabetically, preserving the existing MoEDynamicKernel and _TASK_SLICE_CHUNK exports.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/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py`:
- Line 3695: Update the public entry point containing the scatter_output
parameter to document and validate that scatter_output is a contiguous,
16-byte-aligned [num_tokens, K] tensor before launching the scatter path; reject
invalid layouts before the global address calculation and 128-bit reduction
operations.
- Around line 3336-3346: Update the packed-A offset calculation in the dynamic
gated store path to use Int64 for both phys_row/output_bytes_per_row and
sf_idx/scale-byte contributions, and apply the same change to the corresponding
scale_storage offset arithmetic. Ensure the Int64 offset is passed through
get_ptr_as_int64 so large routed rows and hidden sizes cannot overflow 32-bit
arithmetic.
- Around line 3783-3789: Update the optimized-kernel selection in
_can_use_gated_optimized_kernel() to reject intermediate_size values from 1
through 128, matching _pad_intermediate_to_tile’s padding behavior. Preserve the
existing fallback for intermediate_size <= 0 and ensure larger dimensions
continue through MoEGatedDynamicKernel.
- Around line 2778-2786: Update the Stage2 gate/up pipeline handling around the
fc1_k_tile_cnt loop so a pending gate wait is always consumed before the next
Task 0 slice can reuse Stage2 storage. Either reset the pipeline index for each
fc1_half and restrict the TMA-stall wait to the active half, or add a final wait
that drains gate_wait_pending while preserving the existing deferred in-loop
wait.
In `@tests/moe/test_b12x_fused_moe.py`:
- Around line 108-129: Add the existing cute_dsl_available pytest marker to
test_gated_dynamic_optimized_capability_bounds, matching the two following
tests, so imports of moe_dynamic_kernel and its CuteDSL dependencies are skipped
when CuteDSL is unavailable.
---
Duplicate comments:
In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py`:
- Around line 3415-3432: Add a host-side validation on the affected dynamic
gated MoE path, before launching or executing the kernel, that rejects any
top_k/num_topk value greater than 16 with an appropriate error. Ensure both
route_gs population and comparison remain bounded by valid inputs, and do not
clamp the value in the kernel because routing must not change.
- Around line 3701-3718: The kernel invocation currently stores call-specific
dtypes, layouts, and `_setup_attributes` results on the cached instance,
allowing callers to overwrite shared state. Introduce a per-invocation
configuration object in `__call__`, move `a_dtype`, `b_dtype`, `sf_dtype`,
`a_layout`, `b_layout`, `c_layout`, and the tiled MMA/shared-memory attributes
initialized by `_setup_attributes` into it, and update downstream dynamic-kernel
logic to read from that object instead of `self`.
---
Nitpick comments:
In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.py`:
- Around line 3929-3934: Remove the unused sequential_branch_compact and
fc1_storage_alias configuration from MoEGatedDynamicKernel, including their
getattr declarations, dependent alias branches, and unreachable
up_pipeline.producer_tail call; do not preserve dead configuration paths unless
the attributes are explicitly initialized and intentionally supported in
__init__.
- Around line 442-492: Remove the unused load_global_bf16x16_to_f32x16 function,
including its inline assembly and result-conversion logic, while leaving the
shared Q0 staging loader and other module functionality unchanged.
In `@flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dynamic_kernel.py`:
- Line 113: Update the module’s __all__ declaration to include
_MAX_SHARED_INPUT_TOPK and sort all exported names alphabetically, preserving
the existing MoEDynamicKernel and _TASK_SLICE_CHUNK exports.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: f67335ff-21b9-4872-a7c6-4d3d05f1940e
📒 Files selected for processing (5)
flashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/gated.pyflashinfer/fused_moe/cute_dsl/blackwell_sm12x/_moe_dynamic/generic.pyflashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.pyflashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dynamic_kernel.pytests/moe/test_b12x_fused_moe.py
| @pytest.mark.parametrize( | ||
| "overrides,expected", | ||
| [ | ||
| ({}, True), | ||
| ({"activation": "relu2"}, False), | ||
| ({"sf_vec_size": 8}, False), | ||
| ({"mma_tiler_mn": (64, 128)}, False), | ||
| ({"hidden_size": 16384}, True), | ||
| ({"hidden_size": 16385}, False), | ||
| ({"intermediate_size": 512}, True), | ||
| ({"intermediate_size": 640}, False), | ||
| ({"num_topk": 16}, True), | ||
| ({"num_topk": 17}, False), | ||
| ({"num_topk": 32, "share_input_across_experts": True}, True), | ||
| ({"num_topk": 33, "share_input_across_experts": True}, False), | ||
| ], | ||
| ) | ||
| def test_gated_dynamic_optimized_capability_bounds(overrides, expected): | ||
| """Unsafe Q0/Q1 and route-cache shapes must use the generic fallback.""" | ||
| from flashinfer.fused_moe.cute_dsl.blackwell_sm12x.moe_dynamic_kernel import ( | ||
| _can_use_gated_optimized_kernel, | ||
| ) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add the @cute_dsl_available marker to this test.
The import of moe_dynamic_kernel pulls in ._moe_dynamic.generic and ._moe_dynamic.gated, which both import cutlass and cutlass.cute. Without CuteDSL installed, this test errors instead of skipping. The two following tests already carry the marker.
💚 Proposed fix
+@cute_dsl_available
`@pytest.mark.parametrize`(
"overrides,expected",
[
({}, True),📝 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.
| @pytest.mark.parametrize( | |
| "overrides,expected", | |
| [ | |
| ({}, True), | |
| ({"activation": "relu2"}, False), | |
| ({"sf_vec_size": 8}, False), | |
| ({"mma_tiler_mn": (64, 128)}, False), | |
| ({"hidden_size": 16384}, True), | |
| ({"hidden_size": 16385}, False), | |
| ({"intermediate_size": 512}, True), | |
| ({"intermediate_size": 640}, False), | |
| ({"num_topk": 16}, True), | |
| ({"num_topk": 17}, False), | |
| ({"num_topk": 32, "share_input_across_experts": True}, True), | |
| ({"num_topk": 33, "share_input_across_experts": True}, False), | |
| ], | |
| ) | |
| def test_gated_dynamic_optimized_capability_bounds(overrides, expected): | |
| """Unsafe Q0/Q1 and route-cache shapes must use the generic fallback.""" | |
| from flashinfer.fused_moe.cute_dsl.blackwell_sm12x.moe_dynamic_kernel import ( | |
| _can_use_gated_optimized_kernel, | |
| ) | |
| `@cute_dsl_available` | |
| `@pytest.mark.parametrize`( | |
| "overrides,expected", | |
| [ | |
| ({}, True), | |
| ({"activation": "relu2"}, False), | |
| ({"sf_vec_size": 8}, False), | |
| ({"mma_tiler_mn": (64, 128)}, False), | |
| ({"hidden_size": 16384}, True), | |
| ({"hidden_size": 16385}, False), | |
| ({"intermediate_size": 512}, True), | |
| ({"intermediate_size": 640}, False), | |
| ({"num_topk": 16}, True), | |
| ({"num_topk": 17}, False), | |
| ({"num_topk": 32, "share_input_across_experts": True}, True), | |
| ({"num_topk": 33, "share_input_across_experts": True}, False), | |
| ], | |
| ) | |
| def test_gated_dynamic_optimized_capability_bounds(overrides, expected): | |
| """Unsafe Q0/Q1 and route-cache shapes must use the generic fallback.""" | |
| from flashinfer.fused_moe.cute_dsl.blackwell_sm12x.moe_dynamic_kernel import ( | |
| _can_use_gated_optimized_kernel, | |
| ) |
🤖 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_b12x_fused_moe.py` around lines 108 - 129, Add the existing
cute_dsl_available pytest marker to
test_gated_dynamic_optimized_capability_bounds, matching the two following
tests, so imports of moe_dynamic_kernel and its CuteDSL dependencies are skipped
when CuteDSL is unavailable.
jiahanc
left a comment
jiahanc
left a comment
There was a problem hiding this comment.
lgtm, thanks for the contribution
|
[FAILED] Pipeline #61137611 — 14/18 executed test jobs passed Compared with nightly #60831563. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 4/6 passed
Failure detailsTimeouts, infrastructure, or incomplete jobs
|
Optimize the SM12x static MoE path while keeping NVFP4 and MXFP4 in one `MoEStaticKernel` implementation. ## What changed - Fold the retained NVFP4 schedule into the existing `MoEStaticKernel`; no separate retained-kernel class or source file remains. - Split oversized routed experts into 32-row virtual tasks, so skewed experts remain within the tile's computed M extent. - Retain two FC1 N-slices per scheduled work group and use the compact O(1)-claim scheduler. - Retune the static tile/MAC ladder and extend the default NVFP4 static cutover from 640 to 1024 routed rows (`M <= 128` for top-k 8). - Keep MXFP4 on the same unified class with its original 640-row cutover. - Update workspace sizing, compile/launch ABI, source tracking, and dispatch tests for the unified kernel. The branch is rebased on current `origin/main`. Upstream #4603 now provides the shared B12x workspace support that was previously a separate commit in this PR, so this update deliberately drops that redundant commit. The PR is now one commit touching four files. ## Why The former static schedule could assign more rows from a skewed expert than a tile64/tile128 launch computed, and it repeated scheduling/staging work across the two FC1 N-slices. Virtual 32-row tasks bound each physical work item while the retained schedule reuses pipeline state across both slices. ## Correctness Validated source-exact against the upstream measurement base. The branch is now rebased on the current upstream head; the intervening upstream commits do not modify any of this PR's four changed files. - Measurement upstream base: `083012d6819cf97128e559616b12acb666f2fffe` - Current upstream base: `3bbfeba6218b6de32d1e894243c010c8d3aacb21` - PR head: `e184523e35f59144132d750e085243d409a16cf4` - Local SM120 GPU: `GPU-1c189c11-e797-a795-cefd-495b190afebc` - Shape: Qwen3.5-35B TP1, `E=256`, `H=2048`, `I=512`, `topk=8` - Routes: three exact-marginal Zipf-0.75 samples - M: 32, 64, 96, 128, 256, 512 - Candidate repeats: three per `(M, route)` Result: **18/18 cases passed**. | Metric | Result | |---|---:| | Maximum relative L2 | 0.00791765 | | Minimum cosine similarity | 0.99996866 | | Maximum zero rows | 0 | | Static candidate repeats | Bitwise equal | The baseline dispatches M=96/128 to dynamic while this PR intentionally dispatches them to static; those cross-backend cases are included in the 18/18. MXFP4 targeted GPU tests also passed: - static functional accuracy - intermediate-size padding accuracy - wrapper CUDA Graph accuracy ## Performance Protocol: local SM120, A-B-B-A order, fresh cache per arm, exact-marginal 100-route replay, CUDA Graph event timing, 192 MiB L2 flush, and warmup/iterations/repeats = 5/50/7. No kernel-selection override was used. | M | Upstream main (us) | This PR (us) | Latency reduction | Dispatch | |---:|---:|---:|---:|---| | 1 | 28.408 | 28.548 | -0.49% | direct_micro → direct_micro | | 2 | 43.077 | 43.103 | -0.06% | direct_micro → direct_micro | | 4 | 79.381 | 78.937 | 0.56% | static → static | | 8 | 111.938 | 95.736 | 14.47% | static → static | | 16 | 176.482 | 150.483 | 14.73% | static → static | | 24 | 245.054 | 204.748 | 16.45% | static → static | | 32 | 289.707 | 228.786 | 21.03% | static → static | | 48 | 339.848 | 281.051 | 17.30% | static → static | | 64 | 383.466 | 325.493 | 15.12% | static → static | | 96 | 394.043 | 377.587 | 4.18% | dynamic → static | | 128 | 450.756 | 415.284 | 7.87% | dynamic → static | | 256 | 464.900 | 465.468 | -0.12% | dynamic → dynamic | | 512 | 463.755 | 465.730 | -0.43% | dynamic → dynamic | | 1024 | 481.790 | 483.247 | -0.30% | dynamic → dynamic | | 1536 | 538.665 | 539.207 | -0.10% | dynamic → dynamic | | 2048 | 549.935 | 550.471 | -0.10% | dynamic → dynamic | | 3072 | 599.754 | 597.256 | 0.42% | dynamic → dynamic | | 4096 | 640.822 | 639.772 | 0.16% | dynamic → dynamic | | 5120 | 694.750 | 693.507 | 0.18% | dynamic → dynamic | | 6144 | 790.377 | 788.815 | 0.20% | dynamic → dynamic | | 7168 | 935.373 | 937.726 | -0.25% | dynamic → dynamic | | 8192 | 968.798 | 967.644 | 0.12% | dynamic → dynamic | - Static-band geometric-mean latency reduction: **12.63%** - M=32–128 geometric-mean latency reduction: **13.32%** - Full 22-point geometric-mean latency reduction: **5.34%** - Maximum upstream A-arm drift: **0.37%** - Maximum PR B-arm drift: **0.21%** Dynamic-only points are non-regression controls; differences there are within the predeclared 1% maintenance threshold. ## Tests ```text pytest tests/moe/test_b12x_fused_moe.py -k 'cutover or share_cached_workspace' 2 passed, 192 deselected pytest \ tests/moe/test_b12x_fused_moe.py::TestB12xFunctional::test_mxfp4_static_functional_accuracy \ tests/moe/test_b12x_fused_moe.py::TestB12xFunctional::test_mxfp4_intermediate_padding_accuracy \ tests/moe/test_b12x_fused_moe.py::TestB12xWrapper::test_mxfp4_wrapper_cuda_graph_accuracy 3 passed pre-commit run --files <four changed files> all hooks passed ``` ## Relationship to #4329 #4329 optimized the gated dynamic NVFP4 path. This PR targets the complementary small-token static path and leaves that merged dynamic implementation unchanged. ## Reviewer notes Please pay particular attention to virtual-task allocation/publication ordering, workspace sizing, the shared NVFP4/MXFP4 ABI, and the quant-mode-specific 1024 vs 640 routed-row cutover. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Reuse externally provided workspaces and output buffers during CUDA graph execution. * Improved NVFP4 performance and static execution support for larger workloads. * Maintained optimized MXFP4 execution selection across supported workload sizes. * **Bug Fixes** * Improved workspace sizing, routing, and buffer handling for expanded workloads. * Added validation for incompatible shared buffers, output shapes, data types, devices, capacities, and execution modes. * Expanded regression coverage for buffer reuse, backend selection, and static/dynamic execution paths. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: EricChen02 <EricChen02@users.noreply.github.com>
4 participants
Route gated activations through the optimized branch-paired dynamic kernel while preserving the generic fallback for non-gated activations. Support SiLU, GELU-tanh, and SwiGLU-OAI without an environment toggle.
📌 Description
This PR integrates an optimized branch-paired dynamic NVFP4 MoE kernel for gated activations on Blackwell SM12x into the regular FlashInfer CuTeDSL dispatch path.
What changed
_moe_dynamic/generic.py: the existing generic fallback._moe_dynamic/gated.py: the optimized gated implementation.silu,gelu_tanh, andswigluoai_uninterleaveto the optimized gated kernel.relu2.MoEDynamicKernelAPI unchanged.Why
A gated MoE FC1 computes two projections:
The previous generic dynamic implementation handled the two branches without exploiting their shared scheduling structure.
The new implementation pairs the gate and up branch work in the dynamic N64 kernel. This shortens intermediate lifetimes and reduces staging overhead while preserving the existing routing, NVFP4 quantization, FC2, and scatter behavior.
The top-level activation dispatch selects the gated implementation. The internal
is_gateddistinction is still required because the same kernel object participates in both stages:Correctness
Validated on an SM120 Blackwell GPU against the BF16 reference with:
silugelu_tanhswigluoai_uninterleaveNo NaN or Inf values were observed.
Static validation also passed:
Performance
The benchmark compares this PR against the native FlashInfer CuTeDSL implementation at the PR's current upstream base. No kernel overlay or environment-variable kernel override was used.
Compared revisions:
d7e390c17844f493db23320cb7952375f03bc6c4594ca394c9328db582a689591392e13adcab092bBenchmark environment:
Both revisions were measured on the same GPU with the same benchmark command and protocol. Each revision used an isolated JIT/cache directory. Lower latency is better.
Geometric-mean latency reduction across the sweep:
The static/micro M<96 subset changed by only -0.19% to +0.22% across the five configurations, consistent with measurement noise on the unchanged path. The larger full-sweep gains therefore come from the dynamic path targeted by this PR.
🔍 Related Issues
N/A
🚀 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 my preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Please pay particular attention to:
is_gatedbehavior.The non-gated fallback and public
MoEDynamicKernelinterface are intentionally unchanged.Summary by CodeRabbit
New Features
Refactor