[WIP] [Feature] Add Turbo MXFP8 Grouped GEMM (gfx950) for MoE - #330
[WIP] [Feature] Add Turbo MXFP8 Grouped GEMM (gfx950) for MoE#330kyle-256 wants to merge 1 commit into
Conversation
There was a problem hiding this comment.
Pull request overview
This PR adds a new Turbo MXFP8 (MX_BLOCKWISE) grouped GEMM path targeting gfx950 (MI350/MI355) to accelerate MoE FFN workloads, including fused per-group padding/quantization and dedicated wgrad (variable-K) support.
Changes:
- Added gfx950 Turbo grouped GEMM kernels (fwd/dgrad) plus a variable-K Turbo kernel for wgrad, along with new PyTorch extension entry points.
- Added fused MXFP8 quantization ops for grouped inputs (including per-group M padding) to keep the flow async and avoid host syncs.
- Added MXFP8 Turbo wrapper + parameter-sweep test coverage and a benchmark option.
Reviewed changes
Copilot reviewed 19 out of 19 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/pytorch/ops/test_grouped_gemm_fp8.py | Adds gfx950-only MX_BLOCKWISE Turbo backend test coverage and ensures E8M0 scale dtype configuration. |
| primus_turbo/pytorch/ops/grouped_gemm_fp8.py | Introduces GroupedGemmFP8MXFunc wrapper for MX_BLOCKWISE (handles trans_b=False and per-group padding) and calls new Turbo ops. |
| primus_turbo/pytorch/kernels/quantization/quantization_impl.py | Adds Python wrapper for fused grouped MXFP8 dual quantization with per-group M padding. |
| primus_turbo/pytorch/kernels/grouped_gemm/grouped_gemm_fp8_impl.py | Registers Turbo backend entries for MX_BLOCKWISE grouped GEMM and variable-K wgrad dispatch. |
| csrc/pytorch/quantization/quantization.cpp | Implements quantize_mxfp8_dual_perg and quantize_mxfp8_dual_grouped (grouped + padded) C++ entry points. |
| csrc/pytorch/quantization/quantization_meta.cpp | Adds Meta kernels for the new MXFP8 quantization ops to support shape inference. |
| csrc/pytorch/grouped_gemm/turbo_grouped_gemm.cpp | Adds PyTorch bindings for Turbo MXFP8 grouped GEMM and Turbo variable-K wgrad entry points. |
| csrc/pytorch/extensions.h | Declares new Turbo grouped GEMM and MXFP8 quantization APIs. |
| csrc/pytorch/bindings_pytorch.cpp | Registers new ops (quantize_mxfp8_dual_grouped, quantize_mxfp8_dual_perg, Turbo grouped GEMMs) for CUDA and Meta. |
| csrc/kernels/quantization/quantization_mxfp8.cu | Extends MXFP8 quantization kernels for per-group/per-layout variants and adds per-group padded-layout computation. |
| csrc/kernels/grouped_gemm/turbo/turbo_grouped_gemm_mxfp8_kernel.h | Adds Turbo MXFP8 grouped GEMM persistent kernel (fwd/dgrad) implementation for gfx950. |
| csrc/kernels/grouped_gemm/turbo/turbo_grouped_gemm_mxfp8_wgrad_kernel.h | Adds Turbo MXFP8 variable-K grouped GEMM persistent kernel for wgrad on gfx950. |
| csrc/kernels/grouped_gemm/turbo_grouped_gemm.cu | Adds grouped-GEMM-specific preshuffle + workspace handling and launches Turbo grouped GEMM kernels. |
| csrc/include/primus_turbo/quantization.h | Adds constants and declares new MXFP8 grouped quantization + padded-layout helpers. |
| csrc/include/primus_turbo/grouped_gemm.h | Adds public parameter structs and APIs for Turbo MXFP8 grouped GEMM and variable-K wgrad. |
| csrc/include/primus_turbo/device/register.cuh | Strengthens inline-asm clobbering to improve race-freeness/ordering assumptions. |
| csrc/include/primus_turbo/device/mfma.cuh | Adds memory clobbers to MFMA inline asm to match intended synchronization semantics. |
| csrc/include/primus_turbo/device/memory.cuh | Adds VGPR clobbers after LDS reads to align with race-freeness assumptions. |
| benchmark/ops/bench_grouped_gemm_turbo.py | Adds an mxfp8 benchmark configuration and CLI option for comparison runs. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| // ``quantize_mxfp8_dual_perg``: per-group dual quant for the grouped GEMM | ||
| // B-side. Input is 3D ``(G, M, N)``; per-group M/N are uniform. The | ||
| // kernel writes: | ||
| // rowwise_output: ``(G, M_pad, N_pad)`` -- naturally fits ``(G, N_of_B, | ||
| // K_pad)`` for the fwd GEMM. | ||
| // colwise_output: ``(G, N, M_pad)`` -- naturally fits ``(G, K_of_B, | ||
| // N_pad)`` for the dgrad GEMM | ||
| // (no separate regroup). |
| if constexpr (Bytes == 16) { | ||
| clobber_vgpr_one<VDST + 2>(); | ||
| clobber_vgpr_one<VDST + 3>(); | ||
| } |
There was a problem hiding this comment.
Why do we need to add clobber_vgpr_one here for ds_read? The related registers should already be clobbered before the ds_read.
There was a problem hiding this comment.
Tried removing them — 100% race in grouped MXFP8 dgrad (max diff inf) and 0.7% race in dense MXFP8 mxgemm at 4096³, exactly matching the historical race profile we hit before.
The pre-ds_read clobber_vgpr_one<> reserves the dst VGPRs from being reused, but doesn't fence the compiler from sinking the consuming MFMA above the ds_read. The post-asm clobber on the dst VGPR is what creates that ordering edge.
There was a problem hiding this comment.
-
How was the 0.7% race obtained?
-
In the Turbo MXFP8 GEMM, there should be no race in any phase between MFMA and DS. MFMA computes the current block, while DS loads the data for the next block. Any cross-block race has already been constrained using __builtin_amdgcn_s_barrier.
There was a problem hiding this comment.
The role of clobber_vgpr_one is to allocate/register VGPRs. It cannot prevent the compiler from using them.
There was a problem hiding this comment.
Additionally, my personal understanding is that using clobber_vgpr_one on ds_read is not a semantically correct way to enforce ordering, because it only operates at the level of register allocation and does not involve modeling of memory dependencies or execution ordering. Any apparent “correctness effect” is likely just an indirect result of changes in compiler scheduling behavior, rather than a reliable ordering guarantee. Race conditions should instead be addressed using standard synchronization and memory-dependency mechanisms such as s_waitcnt and barriers.
There was a problem hiding this comment.
It seems to be some bug that led to this result. No race in main branch. I will take time to find out.
There was a problem hiding this comment.
The grouped GEMM race is triggered by compiler-emitted SGPR spill: the spill reload uses vmcnt, which gfx950 does not FIFO-order against buffer_load_lds, so a stale m0 can leak into the LDS write address. The latest commits both reduce spill (CK refactor + smem packing + -sink-insts-to-avoid-spills; private_segment drops 44 → 20 bytes) and move the prologue LDS commit onto lgkmcnt (2-step buffer_load → ds_write), decoupling it from the spill's vmcnt traffic — race goes to zero.
Tested on both gemm kernel and grouped gemm kernels, no race condition happened again.
There was a problem hiding this comment.
Why do we need the two-step path? It copies the data through VGPR first and then into SMEM, right? If so, it seems to break the whole design philosophy here:
- it breaks the counter separation model.
vmcnt is supposed to track only buffer_load_lds (GMEM→LDS)
lgkmcnt is supposed to track only ds_read_pinned (LDS→VGPR)
The two-step path introduces ds_write into lgkmcnt, which muddies the semantics of the two counters.
- the two-step path is effectively a reverse optimization.
The original purpose of buffer_load_lds was zero-VGPR buffering plus fully asynchronous transfer. The two-step path regresses into GMEM→VGPR→LDS, which either consumes extra VGPRs or sacrifices asynchronicity.
- SGPR resources on the hardware should already be extremely abundant.
There was a problem hiding this comment.
The 2-step prologue path is removed, and VGPR / SGPR spill is fully eliminated — all kernel specs now show sgpr_spill = 0 / vgpr_spill = 0 / private_segment = 0, and scratch_load / scratch_store are completely gone from disasm. The vmcnt race root cause is fixed at the source.
But clobber_vgpr_one / clobber_agpr_one still need to stay:
- The spill race and the clobber race are two different things. Clobber prevents compiler reordering —
ds_read_pinnedand MFMA both use"n"(VDST)constant operands that encode the register name into the asm string, so the compiler sees no def-use chain. Per GCC: twoasm volatilestatements are not ordered with respect to each other. Without the clobbers, the compiler is free to reorder MFMA before the consumingds_read→ reads stale VGPR. - Verified empirically: with spill at 0, removing both clobbers still races — grouped fwd 5.5%, wgrad 0.05%.
| else | ||
| asm volatile("v_mfma_scale_f32_16x16x128_f8f6f4 v[%0:%1], v[%2:%3], v[%4:%5], v[%0:%1], v[%6], v[%7] op_sel_hi:[0,0,0] cbsz:1 blgp:1" | ||
| : : "n"(PIN_ACC), "n"(PIN_ACC + 3), "n"(PIN_A), "n"(PIN_A + 7), "n"(PIN_B), "n"(PIN_B + 7), "n"(PIN_SA), "n"(PIN_SB)); | ||
| : : "n"(PIN_ACC), "n"(PIN_ACC + 3), "n"(PIN_A), "n"(PIN_A + 7), "n"(PIN_B), "n"(PIN_B + 7), "n"(PIN_SA), "n"(PIN_SB) : "memory"); |
There was a problem hiding this comment.
v_mfma_scale_f32_16x16x128_f8f6f4 only modifies registers, so I don't think we need the "memory" here.
There was a problem hiding this comment.
removed all "memory"
| static_assert(AGPR >= 0 && AGPR <= 255, "AGPR must be in [0, 255]"); | ||
| // clang-format off | ||
| #define CLOBBER_AREG_CASE(N) case N: asm volatile("" ::: "a" #N); break; | ||
| #define CLOBBER_AREG_CASE(N) case N: asm volatile("" ::: "a" #N, "memory"); break; |
| template <int AC> __device__ __forceinline__ void zero_agpr() { | ||
| asm volatile("v_accvgpr_write_b32 a[%0], 0" : : "n"(AC)); | ||
| asm volatile("v_accvgpr_write_b32 a[%0], 0" : : "n"(AC) : "memory"); | ||
| clobber_agpr_one<AC>(); |
| if constexpr (N >= 13) asm volatile("v_accvgpr_read_b32 %0, a[%1]" : "=v"(raw[12]) : "n"(AC + 12) : "memory"); | ||
| if constexpr (N >= 14) asm volatile("v_accvgpr_read_b32 %0, a[%1]" : "=v"(raw[13]) : "n"(AC + 13) : "memory"); | ||
| if constexpr (N >= 15) asm volatile("v_accvgpr_read_b32 %0, a[%1]" : "=v"(raw[14]) : "n"(AC + 14) : "memory"); | ||
| if constexpr (N >= 16) asm volatile("v_accvgpr_read_b32 %0, a[%1]" : "=v"(raw[15]) : "n"(AC + 15) : "memory"); |
| typename GemmTile::AScaleSmemSubtile (*a_s_smem_tile)[4], | ||
| typename GemmTile::BScaleSmemSubtile (*b_s_smem_tile)[4], const AType *a_ptr, | ||
| const BType *b_ptr, const uint32_t *a_s_ptr, const uint32_t *b_s_ptr, CType *c_ptr, | ||
| const int64_t *group_offs_ptr, const int64_t *c_group_offs_ptr, const int32_t group_id, |
There was a problem hiding this comment.
Naming is asymmetric: group_offs_ptr is actually input-side (padded) offsets while c_group_offs_ptr is output-side (compact). Consider renaming to group_offs_in_ptr / group_offs_out_ptr (or a_group_offs_ptr / c_group_offs_ptr) so the asymmetry is visible at the callsite.
There was a problem hiding this comment.
Just realized that in real workloads the inputs probably always need to be padded anyway. In that case, it may be cleaner to stop overloading the presence of c_group_offs_ptr as a mode switch, and instead always pass explicit a_group_offs_ptr, b_group_offs_ptr, and c_group_offs_ptr. The interface becomes much more straightforward.
There was a problem hiding this comment.
Done. Launcher now always passes a non-null c_group_offs_ptr (caller's compact-output offsets when supplied; a_group_offs_ptr as fallback for the in-place layout). Kernel takes a single straight-line extract path with no mode switch on the pointer.
On b_group_offs_ptr: B is laid out as [G, N, K] and indexed by group_id directly — no per-group row offsets for B, so no third pointer to add.
| // per-group region to 128-aligned, so this matches | ||
| // group_offs_padded[g+1] - group_offs_padded[g] without an | ||
| // extra scalar load. | ||
| M_g_in = (M_g + 127) & ~127; |
| static_assert(VGPR >= 0 && VGPR <= 255, "VGPR must be in [0, 255]"); | ||
| // clang-format off | ||
| #define CLOBBER_VREG_CASE(N) case N: asm volatile("" ::: "v" #N); break; | ||
| #define CLOBBER_VREG_CASE(N) case N: asm volatile("" ::: "v" #N, "memory"); break; |
| granularity: ScalingGranularity, | ||
| num_cu: int | None, | ||
| **kwargs, | ||
| ) -> bool: |
There was a problem hiding this comment.
Done — added: n_fwd % 16 == 0 and k_fwd % 16 == 0 (wgrad MFMA tile is 16×16, output axes must clear)
78ccce9 to
847872c
Compare
…ler spill race Address PR #330 review #12 properly: launcher always passes a non-null c_group_offs_ptr (caller's compact-output offsets, or a_group_offs as fallback for the in-place layout); kernel takes a single straight-line extract path with no mode switch. Side effect: removing the if-else collapses the SSA phi-merge that had been sharing c_m_global / M_g_in register slots with m_global / M_g. Peak SGPR pressure across the LDS-direct buffer_load_lds prologue then crosses the per-wave limit (~100), and the compiler spills via scratch_store/load + readfirstlane. The reload's vmcnt-tracked completion can lose to a younger buffer_load on the shared counter (no FIFO guarantee), so a stale m0 leaks in and the buffer writes to the wrong LDS slot. Empirically reproduces at the gpt_oss-20B Expert shape (B=4 M=2048 N=5760 K=2880) at ~5% fwd in 10K reruns; tested 9 variants at the C++ level (assertions, fences, __builtin_assume, sched_barrier, deferral, formulation rewrites) — none recover baseline race rate. Workaround: 4 explicit s_waitcnt vmcnt(0) drains in the prologue between the first 4 buffer_load_lds calls, fencing the spill reload against the m0 setup. Empirically only the first four are needed (SGPR pressure peak); the rest run on already-flushed pressure. 0/20000 fwd reruns at gpt_oss-20B Expert + 0/10000 dense MXFP8. +4% fwd latency vs. baseline (338us vs 325us at the same shape). Filing the underlying compiler issue (vmcnt scoping for spill consumes against buffer_load_lds) separately.
|
@xiaobochen-amd Thanks for the commit! I have updated my code by the comments. Please review again, thanks. |
0df887d to
07e9940
Compare
| "alignment); also required: each per-group M_g % 128 == 0."); | ||
|
|
| // hint we use it; otherwise fall back to the pessimistic upper bound | ||
| // derivable from input metadata alone (no D2H sync of group_lens). | ||
| // Worst case: a single group holds all real rows, then per-group | ||
| // padding to MXFP8_PADDING_ALIGN_SIZE adds (group_num - 1)*ALIGN | ||
| // padding rows on top. | ||
| int32_t grid_x; | ||
| if (grid_x_hint > 0) { | ||
| grid_x = (int32_t) grid_x_hint; | ||
| } else { | ||
| constexpr int64_t ALIGN = primus_turbo::detail::MXFP8_PADDING_ALIGN_SIZE; | ||
| const int64_t total_m_upper = | ||
| ((int64_t) total_m_in + group_num * ALIGN + ALIGN - 1) / ALIGN * ALIGN; | ||
| grid_x = (int32_t) ((total_m_upper + 255) / 256); |
| // rowwise_output: ``(G, M_pad, N_pad)`` -- naturally fits ``(G, N_of_B, | ||
| // K_pad)`` for the fwd GEMM. |
| // | ||
| // Per-group dB[g] = LHS[g] @ RHS[g]^T, with LHS (N, total_M), RHS | ||
| // (K, total_M), dB (group_num, N, K); reduction is over total_M. | ||
| // Constraints: n % 16 == 0, k % 16 == 0, M_g % 32 == 0. |
| # MFMA tile = 16x16; output axes (N_fwd, K_fwd) must clear that. | ||
| if a.dim() == 2 and b.dim() == 2: | ||
| n_fwd = a.shape[0] | ||
| k_fwd = b.shape[1] | ||
| supported &= n_fwd % 16 == 0 and k_fwd % 16 == 0 | ||
| supported &= a.shape[1] == b.shape[0] |
|
|
||
| class GroupedGEMMFP8VariableKKernelDispatcher(BaseGroupedGEMMVariableKKernelDispatcher): | ||
| _backends = { | ||
| BackendType.TURBO: BackendEntry(GroupedGEMMFP8VariableKTurboBackend), |
| @pytest.mark.parametrize("balance", BALANCE_VALUES) | ||
| def test_grouped_gemm_fp8_mx_blockwise(B, M, NK, ori_dtype, format, trans_b, balance): |
| ScalingRecipe colwise_recipe, hipStream_t stream); | ||
|
|
||
| // Per-group dual quant for the MXFP8 grouped GEMM B-side: writes the | ||
| // rowwise output as ``(G, M_pad, N_pad)`` and the colwise output as |
07e9940 to
8244d7a
Compare
| // ``quantize_mxfp8_dual_perg``: per-group dual quant for the grouped GEMM | ||
| // B-side. Input is 3D ``(G, M, N)``; per-group M/N are uniform. The | ||
| // kernel writes: | ||
| // rowwise_output: ``(G, M_pad, N_pad)`` -- naturally fits ``(G, N_of_B, | ||
| // K_pad)`` for the fwd GEMM. | ||
| // colwise_output: ``(G, N, M_pad)`` -- naturally fits ``(G, K_of_B, | ||
| // N_pad)`` for the dgrad GEMM | ||
| // (no separate regroup). |
| detail::ScalingRecipe colwise_recipe, hipStream_t stream); | ||
|
|
||
| // Per-group dual quant with uniform per-group (M, N): input (G, M, N) -> | ||
| // rowwise (G, M_pad, N_pad), colwise (G, N, M_pad). Used by the grouped |
abaeaab to
e0263a0
Compare
e0263a0 to
17df201
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 21 out of 21 changed files in this pull request and generated 3 comments.
Comments suppressed due to low confidence (2)
tests/pytorch/ops/test_grouped_gemm_fp8.py:554
- Same issue as the deterministic MX_BLOCKWISE test above: this uses an exact (9, 5) compute capability check, which will skip on newer gfx95x devices even if they support MXFP8. Use check_mxfp8_support() or >= (9, 5) instead of equality.
if get_device_compute_capability() != (9, 5):
pytest.skip("MXFP8 grouped GEMM requires gfx950 (MI350/MI355).")
primus_turbo/pytorch/kernels/grouped_gemm/grouped_gemm_fp8_impl.py:628
- Same concern as the fwd/dgrad TURBO backend: can_handle() for the variable-K TURBO backend doesn’t verify gfx950/MXFP8 support. Adding a device capability gate would prevent the dispatcher from selecting a kernel that asserts at runtime on unsupported GPUs when invoked with pre-quantized tensors.
supported = True
supported &= a.dim() == 2 and b.dim() == 2
supported &= (a.dtype, b.dtype, out_dtype) in GroupedGEMMFP8VariableKTurboBackend.SUPPORTED_DTYPES
supported &= granularity in GroupedGEMMFP8VariableKTurboBackend.SUPPORTED_GRANULARITIES
supported &= not trans_a and not trans_b and not trans_c
# MFMA tile = 16x16; output axes (N_fwd, K_fwd) must clear that.
if a.dim() == 2 and b.dim() == 2:
n_fwd = a.shape[0]
k_fwd = b.shape[1]
supported &= n_fwd % 16 == 0 and k_fwd % 16 == 0
supported &= a.shape[1] == b.shape[0]
return supported
| @pytest.mark.parametrize("balance", [True, False]) | ||
| @pytest.mark.deterministic | ||
| def test_grouped_gemm_fp8_mx_blockwise_deterministic(B, M, NK, ori_dtype, format, trans_b, balance): | ||
| if get_device_compute_capability() != (9, 5): |
| // ``quantize_mxfp8_dual_perg``: per-group dual quant for the grouped GEMM | ||
| // B-side. Input is 3D ``(G, M, N)``; per-group M/N are uniform. The | ||
| // kernel writes: | ||
| // rowwise_output: ``(G, M_pad, N_pad)`` -- naturally fits ``(G, N_of_B, | ||
| // K_pad)`` for the fwd GEMM. | ||
| // colwise_output: ``(G, N, M_pad)`` -- naturally fits ``(G, K_of_B, | ||
| // N_pad)`` for the dgrad GEMM | ||
| // (no separate regroup). |
| supported = True | ||
| supported &= a.dim() == 2 and b.dim() == 3 | ||
| supported &= (a.dtype, b.dtype, out_dtype) in GroupedGEMMFP8TurboBackend.SUPPORTED_DTYPES | ||
| supported &= granularity in GroupedGEMMFP8TurboBackend.SUPPORTED_GRANULARITIES | ||
| supported &= not trans_a and trans_b | ||
| total_m = a.shape[0] | ||
| n = b.shape[-2] if trans_b else b.shape[-1] | ||
| k = a.shape[1] | ||
| supported &= total_m % 16 == 0 and n % 16 == 0 and k % 128 == 0 and k >= 384 | ||
| return supported |
| return grad_out if grad_out.is_contiguous() else grad_out.contiguous() | ||
|
|
||
|
|
||
| def _turbo_grouped_gemm_mxfp8( |
There was a problem hiding this comment.
The code here does not align with our backend design
| @pytest.mark.parametrize("balance", [True, False]) | ||
| @pytest.mark.deterministic | ||
| def test_grouped_gemm_fp8_mx_blockwise_deterministic(B, M, NK, ori_dtype, format, trans_b, balance): | ||
| if get_device_compute_capability() != (9, 5): |
There was a problem hiding this comment.
check_mxfp8_support
17df201 to
c5c03ad
Compare
| assert out_dtype in [torch.float16, torch.bfloat16] | ||
|
|
||
| # Kernel only supports NT, so materialise b in (G, N, K). | ||
| b_internal = b if trans_b else b.transpose(-1, -2).contiguous() |
| # Fused grouped quant for A: produces row + col MXFP8 outputs | ||
| # in the padded layout AND computes the padded per-group layout | ||
| # on GPU (no D2H sync). | ||
| a_fp8_row, a_scale_inv_row, a_fp8_col, a_scale_inv_col, _, group_offs_padded = ( | ||
| quantize_mxfp8_dual_grouped_impl( | ||
| a, | ||
| a_dtype, | ||
| group_lens, | ||
| group_offs, | ||
| ) |
| // ``quantize_mxfp8_dual_perg``: per-group dual quant for the grouped GEMM | ||
| // B-side. Input is 3D ``(G, M, N)``; per-group M/N are uniform. The | ||
| // kernel writes: | ||
| // rowwise_output: ``(G, M_pad, N_pad)`` -- naturally fits ``(G, N_of_B, | ||
| // K_pad)`` for the fwd GEMM. | ||
| // colwise_output: ``(G, N, M_pad)`` -- naturally fits ``(G, K_of_B, | ||
| // N_pad)`` for the dgrad GEMM | ||
| // (no separate regroup). |
792e1b5 to
d0ec992
Compare
| # fallback computes from input shape alone. Without this the loop | ||
| # iterates ~B× more tile_ids than necessary (all early-skip), which | ||
| # for B≥16 collapses fwd/dgrad TFLOPS to ~40-60% of peak. Single | ||
| # D2H sync per call; negligible vs kernel time. | ||
| grid_x_hint = (int(group_lens.max().item()) + 255) // 256 |
| // rowwise_output: ``(G, M_pad, N_pad)`` -- naturally fits ``(G, N_of_B, | ||
| // K_pad)`` for the fwd GEMM. | ||
| // colwise_output: ``(G, N, M_pad)`` -- naturally fits ``(G, K_of_B, | ||
| // N_pad)`` for the dgrad GEMM | ||
| // (no separate regroup). |
| num_cu: int | None, | ||
| default_backend: int, | ||
| maybe_pre_sync: bool = False, | ||
| c_group_offs: torch.Tensor | None = None, |
There was a problem hiding this comment.
In mxfp8 gemm, we padded a and b in quantization kernel. So we assume a and b are padded in grouped gemm kernel. Please remove these unused arguments.
| // Fold layout computation into this op (no separate D2H round-trip). | ||
| auto group_lens_padded = at::empty_like(group_lens); | ||
| auto group_offs_padded = at::empty({G + 1}, group_offs.options()); | ||
| primus_turbo::compute_padded_layout_gpu( |
There was a problem hiding this comment.
The MXFP8_PADDING_ALIGN_SIZE means n dim padding align size. For group offs, we may want to pad m dim if group_offs is unbalanced but I think the align size set to 128 is to huge. Maybe we should align with the m dim padding align size of MXFP8 GEMM (
) ?| if a.dim() == 2 and b.dim() == 2: | ||
| n_fwd = a.shape[0] | ||
| k_fwd = b.shape[0] | ||
| supported &= n_fwd % 16 == 0 and k_fwd % 16 == 0 | ||
| supported &= a.shape[1] == b.shape[1] | ||
| return supported |
| // Bitmask used to round M_g up to the input-layout padding alignment | ||
| // for SRD bound calculation: ``M_g_in = (M_g + mask) & ~mask``. | ||
| // 127 → padded input (per-group regions aligned to 128 rows) | ||
| // 0 → unpadded input (M_g_in == M_g) | ||
| // Branchless on the kernel side; lets the outer kernel resolve | ||
| // c_group_offs_ptr unconditionally, removing a per-tile null-check | ||
| // and the duplicate i64 split-load chain that came with it. | ||
| int32_t c_padding_align_mask = 0; |
| def _run_once(): | ||
| # Prevent buffer-alias race across pytest cases: the caching allocator | ||
| # can hand a fresh tensor the same physical GPU memory still being | ||
| # written by a pending op from a previous test case. Force full sync | ||
| # + cache release so each _run_once gets clean memory. | ||
| torch.cuda.synchronize() | ||
| torch.cuda.empty_cache() |
| int64_t colwise_scale_N_pad = cdiv(colwise_scale_N, 8) * 8; | ||
|
|
||
| at::Tensor colwise_scale; | ||
| int colwise_scale_stride = 1; |
| // | ||
| // Per-group dB[g] = LHS[g] @ RHS[g]^T, with LHS (N, total_M), RHS | ||
| // (K, total_M), dB (group_num, N, K); reduction is over total_M. | ||
| // Constraints: n % 16 == 0, k % 16 == 0, M_g % 32 == 0. | ||
|
|
df25df8 to
0b4f24d
Compare
xiaobochen-amd
changed the title
Description
Add a turbo MXFP8 grouped GEMM backend (gfx950 / MI350) for MoE FFN
workloads. Forward, dgrad and wgrad all run on hand-tuned 256x256x128
persistent kernels with E8M0 scales preshuffled into 16x4 col-major
blocks. The wrapper handles
trans_b=Falseand unaligned per-groupM_g transparently (zero-pad along M), so the new path passes the full
tensorwise parametrization (HYBRID excluded for now).
End-to-end overhead vs
tensorwise + Tritonon gpt-oss-20B shapes:median +12.8%, with M=4096 cases matching or beating the tensorwise
baseline.
Type of change
Changes
csrc/kernels/grouped_gemm/):fwd / dgrad share a 256x256x128 persistent kernel with optional
store-side unpad via
c_group_offs; wgrad uses a variable-K (M_greduction) kernel. uint32 BufferSRD num_records is clamped to avoid
the >4 GB silent-truncation that masks valid LDGs as OOB.
to the single-GEMM kernel but stages each 16-row tile through LDS in
256-byte column slabs, restoring fully coalesced HBM traffic on the
fwd / dgrad / wgrad hot paths.
quantize_mxfp8_dual_perg): writesrowwise
(G, M, N_pad)and colwise(G, N, M_pad)directly,eliminating the torch transpose+contiguous regroup that previously
bottlenecked the dgrad B-side.
quantize_mxfp8_dual_grouped):computes the padded layout on GPU (no D2H sync) and materialises the
padded layout straight into the output tensors, allocating outputs
at the host-known upper bound so the call is fully async.
clobber_*gpr/zero_agpr/read_agprinline asm +clobber_vgpr_oneafterds_read_b*, matching the single-GEMM kernel's race profile.GroupedGemmFP8MXFunc): handlestrans_b=False(NT materialisation + grad transpose-back) and unaligned per-group
M_g (128-aligned per-group zero-pad on A / grad_out) transparently.
test_grouped_gemm_fp8_mx_blockwisemirrors the fulltensorwise param sweep (
B,M,NK,ori_dtype,format,trans_b,balance); HYBRID format is intentionally excluded.bench_grouped_gemm_turbo.pyadds anmxfp8granularityfor direct comparison with
tensorwise + Triton.Checklist
MX Grouped GEMM Bench
GPU: MI355 (gfx950) · dtype: bfloat16 · MX_BLOCKWISE (E4M3 + E8M0, block=32) ·
NT layout · CUDA-event timed (warmup 20 / iters 100). All numbers are
TFLOPS (higher is better).
Fwd / Dgrad / Wgrad): pureturbo_grouped_gemm_fp8/turbo_grouped_gemm_variable_k_fp8, inputs pre-quantized outside thetimed region. Includes the in-kernel scale preshuffle.
FLOPs =
2·B·M·N·Kper kernel.Fwd+Q / Bwd+Q): full Python op times.Fwd+Q=grouped_gemm_fp8(...)(A grouped quant + B per-group quant +GEMM, FLOPs =
2·B·M·N·K).Bwd+Q=out.sum().backward()only (grad_out grouped quant + dgrad +wgrad, FLOPs =
4·B·M·N·K).bal= balanced (M_g = M);unb= random per-group split, total =B·M.DeepSeek-V3 (n_routed_experts=256, hidden=7168, ffn=2048)
Kernel-only · Balanced
Kernel-only · Unbalanced
With Quant · Balanced
With Quant · Unbalanced
gpt_oss_20B (n_routed_experts=32, hidden=2880, ffn=2880)
Kernel-only · Balanced
Kernel-only · Unbalanced
With Quant · Balanced
With Quant · Unbalanced