Grouped MXFP8 GEMMs with HipKittens by alextmagro · Pull Request #661 · ROCm/TransformerEngine

alextmagro · 2026-07-07T16:44:30Z

Modifies HipKittens MXFP8 GEMM kernels to handled single launch grouped GEMM. Usage is enabled through NVTE_USE_CUTLASS_GROUPED_GEMM, similar to CK FP8

github-actions · 2026-07-07T16:55:08Z

Claude review

Reviewed the 5-file diff (transformer_engine/common/gemm/{cublaslt_gemm.cu, rocm_gemm.cu, kittens/mxfp8_gemm.{cpp,h}} and transformer_engine/pytorch/quantization.py) adding grouped MXFP8 GEMM through the HipKittens kernel with a single-launch, expert-concatenated layout.

Verdict: approach is sound and the CUDA path is untouched (__HIP_PLATFORM_AMD__ guards on the include, extern decl, and call in cublaslt_gemm.cu), but there are a few issues worth addressing before merging — see inline comments:

try_kittens_grouped_mxfp8_gemm declaration/call in cublaslt_gemm.cu is not guarded by USE_HIPKITTENS_GEMM, while its definition in rocm_gemm.cu is — this will fail to link when building with -DUSE_HIPKITTENS_GEMM=OFF.
The wrapper collects per-expert pointers into arrays but the kernel only uses [0] from each; the packed-contiguous layout is a silent invariant with no assertion or documentation.
The grouped kernels accept bias/bias_dtype marked [[maybe_unused]], but gemm_epilogue's bias math uses a global block_m — a latent correctness bug if a future GROUPED+BIAS instantiation is added.
NVTE_USE_CUTLASS_GROUPED_GEMM is the sole gate for the HipKittens path but it only affects M-alignment in Python; consider making that gate explicit in C++ too, and either renaming or commenting on the shared env-var name.
No test appears to cover the new path (existing grouped tests target CK); given the env-var + gfx950 gating, coverage is easy to miss.

Copyright headers: OK — all five files updated to 2026 AMD end-year, NVIDIA years unchanged. New AMD-only files (kittens/mxfp8_gemm.{cpp,h}) already carry a 2026 AMD-only header in the correct style.

matthiasdiener · 2026-07-07T16:57:41Z

Usage is enabled through NVTE_USE_CUTLASS_GROUPED_GEMM, similar to CK FP8

This probably needs an update in the Readme, hipkittens is essentially "overriding" the CK grouped Gemm (which is also used for bf16 I believe).

matthiasdiener · 2026-07-07T16:48:18Z

+};
+
+template <typename T>
+__device__ __forceinline__ int rocm_upper_bound(const T *arr, int n, T val) {


Isn't this function in rocm_device_utils.cuh?

Yes, however the kittens folder is built as a separate shared library that doesn't reference TransformerEngine, so we need a copy of this function here.

matthiasdiener · 2026-07-07T17:04:03Z

+void mxfp8_gemm_tn_kernel(const gl_fp8_rt A, const gl_fp8_rt B, const OutGL C, const AuxGLType AuxGL,
+    const gl_scale_rt scale_A_gl, const gl_scale_rt scale_B_gl,
+    [[maybe_unused]] const void *__restrict__ bias, [[maybe_unused]] int bias_dtype,
+    [[maybe_unused]] TileOffsets tile_offsets, [[maybe_unused]] int num_experts,


tile_offsets is passed by value into each launch, including the non-grouped launches via dummy.

Not sure if it is an issue, but could we maybe template this argument out, or pass a pointer instead?

Addressed this with the changes to how tileOffsets is used.

aris134 · 2026-07-08T12:37:19Z

+    if (transa && transb) return false;
+    if (N % BLOCK_COL != 0 || K % BLOCK_K != 0 || K < 256) return false;
+    if (num_experts <= 0 || num_experts > MAX_EXPERTS) return false;


Would it be useful to optionally emit warning messages on the reason for fallback, similarly to the CK setup?

I see you use the NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK flag to enable these warnings, so I have added them similarly.

aris134 · 2026-07-08T12:41:02Z

+    int total_M = 0;
+    int total_m_tiles = 0;
+    for (int g = 0; g < num_experts; g++) {
+        if (M_array[g] % BLOCK_ROW != 0) return false;


Same comment here regarding warning emission. Given that the CK grouped GEMM has padding support:

TransformerEngine/transformer_engine/common/gemm/ck_grouped_gemm/ck_grouped_gemm_fp16.cpp

Lines 14 to 41 in 71eb59d

bool ck_tile_grouped_gemm_fp16_dispatch(DType a_dtype,

DType b_dtype,

DType d_dtype,

const GroupedGemmRunContext& ctx) {

// Check M and K alignment across all groups.

// All tile configs share the same M_Tile (256) and K_Tile (64).

constexpr ck_tile::index_t M_Tile = TileCfg_256x256x64::M_Tile;

constexpr ck_tile::index_t K_Tile = TileCfg_256x256x64::K_Tile;

bool need_m_pad = false;

bool need_k_pad = false;

for (int i = 0; i < ctx.group_num; ++i) {

const transformer_engine::Tensor* A_te =

transformer_engine::convertNVTETensorCheck(ctx.A[i]);

int64_t Ad0 = 0, Ad1 = 0;

if (get_flat_2d_dims(*A_te, Ad0, Ad1)) {

const int64_t M = ctx.transA ? Ad1 : Ad0;

const int64_t K = ctx.transA ? Ad0 : Ad1;

if (M % M_Tile != 0)

need_m_pad = true;

if (K % K_Tile != 0)

need_k_pad = true;

if (need_m_pad && need_k_pad)

break;

}

}

would it make sense to have that for HK as well? Otherwise my question would be how frequently does HK grouped MXFP8 fallback for Qwen or DS runs?

In practice, I don't think we would ever fall back for Qwen and DS as expert weights are all 256 aligned. Activations are padded already by Megatron, so no issue there either.

aris134 · 2026-07-08T12:53:55Z


+constexpr int MAX_EXPERTS = 512;
+
+struct TileOffsets {


Would it make sense to carve tile_offsets out of the workspace?

Good idea. I have done that and combined it into other variables passed to the kernel itself.

aris134 · 2026-07-08T13:08:57Z

numerics tests for the HK grouped GEMM?

Numeric tests are covered by test_grouped_linear_accuracy_cutlass and other grouped tests

Can you double check test_grouped_linear_accuracy_cutlass? I ran it on gfx950 with:

NVTE_ROCM_ENABLE_MXFP8=1 NVTE_USE_HIPKITTENS_GROUPED_GEMM=1 NVTE_USE_CUTLASS_GROUPED_GEMM=0 NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK=1 pytest tests/pytorch/test_numerics.py::test_grouped_linear_accuracy_cutlass -vv -s

and am seeing assertion failures from the numerical correctness checks when comparing HK outputs against the reference implementation.

e.g.

Mismatched elements: 1559808 / 1572864 (99.2%)
Greatest absolute difference: 2.58984375 at index (1792, 0, 751) (up to 0.001 allowed)
Greatest relative difference: 9792.0 at index (0, 0, 148) (up to 0.008 allowed)
FAILED tests/pytorch/test_numerics.py::test_grouped_linear_accuracy_cutlass[hipkittens-True-True-recipe0-True-126m-1-6-torch.bfloat16] - AssertionError: Tensor-likes are not close!

aris134 · 2026-07-08T13:14:41Z

Do we have any performance numbers? I know there isn't a CK MXFP8 group GEMM for gfx950, but do we have hipBlaslt or triton baselines?

aris134 · 2026-07-09T16:29:48Z

+    G::load(As[tic][1], A, {0, 0, a_half1, 0}, sw_A, a_srd, a_base, a_lds[tic][1]);

    if (warp_m == 1) __builtin_amdgcn_s_barrier();
    asm volatile("s_waitcnt vmcnt(4)");


Consider replacing the magic vmcnt values with named constants. It's currently difficult to tell why the hard-coded values were chosen.

I don't think this is the right approach, as the number is not so much a magic number, but depends on how many G::load calls are in flight at each point. Named constants would not hold if the kernel changes.

Right, but I don't think it's strictly the number of G::load calls themselves. The required vmcnt value is really tied to the underlying VMEM operations generated by those loads, which can vary depending on how the backend lowers them (e.g., vector load width per thread).

That gets back to my readability concern. I'm not necessarily advocating for named constants, but it might be helpful to add a brief comment describing what each non-zero vmcnt is waiting on at the pipeline level and why newer loads are allowed to remain in flight.

aris134 · 2026-07-09T18:46:08Z


-        G::load(Bs[tic][1], B, {0, 0, block_col * 2 + 1, k + 2}, sw_B, b_srd, b_base, b_lds[tic][1]);
+        G::load(Bs[tic][1], B, {0, 0, b_half1, k + 2}, sw_B, b_srd, b_base, b_lds[tic][1]);
        asm volatile("s_waitcnt vmcnt(6)");


For example here, I think the vmcnt(6) here means you're retiring k+1 scales and As[toc][1] k+1, but using named constants here could confirm

or just describing in a comment as I said above

…iguous

Grouped MXFP8 GEMMs with HipKittens

34806b9

alextmagro requested review from ipanfilo, wangye805 and wenchenvincent as code owners July 7, 2026 16:44

alextmagro requested review from aris134 and matthiasdiener and removed request for ipanfilo, wangye805 and wenchenvincent July 7, 2026 16:44

alextmagro added ci-level 3 CI test level 3 claude-review labels Jul 7, 2026

github-actions Bot reviewed Jul 7, 2026

View reviewed changes

Comment thread transformer_engine/common/gemm/cublaslt_gemm.cu Outdated

github-actions Bot reviewed Jul 7, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/quantization.py Outdated

github-actions Bot reviewed Jul 7, 2026

View reviewed changes

Comment thread transformer_engine/common/gemm/kittens/mxfp8_gemm.cpp Outdated

github-actions Bot reviewed Jul 7, 2026

View reviewed changes

Comment thread transformer_engine/common/gemm/kittens/mxfp8_gemm.cpp Outdated

github-actions Bot removed the claude-review label Jul 7, 2026

matthiasdiener reviewed Jul 7, 2026

View reviewed changes

aris134 reviewed Jul 8, 2026

View reviewed changes

Comment thread transformer_engine/common/gemm/rocm_gemm.cu

aris134 reviewed Jul 8, 2026

View reviewed changes

Comment thread transformer_engine/common/gemm/kittens/mxfp8_gemm.cpp

aris134 requested changes Jul 8, 2026

View reviewed changes

aris134 reviewed Jul 9, 2026

View reviewed changes

Fixed kernel -- A and B swapped, expert pointers passed when not cont…

1ad82ae

…iguous

ROCm deleted a comment from github-actions Bot Jul 10, 2026

alextmagro added 2 commits July 10, 2026 15:39

Address copilot comments

55683c0

Updated readme, tests, and cleanup

7210715

alextmagro requested review from aris134 and matthiasdiener July 10, 2026 19:36

Revert git move for c++ test

55afb00

aris134 requested changes Jul 13, 2026

View reviewed changes

	bool ck_tile_grouped_gemm_fp16_dispatch(DType a_dtype,
	DType b_dtype,
	DType d_dtype,
	const GroupedGemmRunContext& ctx) {
	// Check M and K alignment across all groups.
	// All tile configs share the same M_Tile (256) and K_Tile (64).
	constexpr ck_tile::index_t M_Tile = TileCfg_256x256x64::M_Tile;
	constexpr ck_tile::index_t K_Tile = TileCfg_256x256x64::K_Tile;

	bool need_m_pad = false;
	bool need_k_pad = false;

	for (int i = 0; i < ctx.group_num; ++i) {
	const transformer_engine::Tensor* A_te =
	transformer_engine::convertNVTETensorCheck(ctx.A[i]);
	int64_t Ad0 = 0, Ad1 = 0;
	if (get_flat_2d_dims(*A_te, Ad0, Ad1)) {
	const int64_t M = ctx.transA ? Ad1 : Ad0;
	const int64_t K = ctx.transA ? Ad0 : Ad1;

	if (M % M_Tile != 0)
	need_m_pad = true;
	if (K % K_Tile != 0)
	need_k_pad = true;
	if (need_m_pad && need_k_pad)
	break;
	}
	}

Uh oh!

Conversation

alextmagro commented Jul 7, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Jul 7, 2026

Claude review

Uh oh!

matthiasdiener commented Jul 7, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

aris134 Jul 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

aris134 Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

aris134 commented Jul 8, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

aris134 Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

aris134 Jul 8, 2026 •

edited

Loading

aris134 Jul 13, 2026 •

edited

Loading

aris134 Jul 13, 2026 •

edited

Loading