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!

There was a bug with non-contiguous inputs, that should be fixed now.

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.

I've added some comments, hopefully it clears up what is 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

wangye805

in the PR description, add the ticket or confluence doc link for performance

wangye805 · 2026-07-17T02:53:04Z

+        n_arr[i]   = n_i;
+    }
+
+    // Activation data (B), activation scales (SB), and output (D) must be contiguous.


For main stream MOE models, permuted tokens will still be guaranteed to be continuous but usually the scales is hard to guarantee to be still continuous, especially after padding.

So what's the motivation of this continuous assumption? Without it, the hipkitten kernel will drop a lot? Does our CK grouped gemm need this assumption?

This was simply a simplification, but you are right that if data is not 256 aligned and scales are not 16 byte aligned then we are unable to use this method. Since we don't do padding for rocm though, generally I would expect us to be okay for mainstream MoE models like you mention.

A note on performance though, if we guarantee that all of output, activations, scales, and weights are contiguous we see about a 20-25% increase in performance for these grouped GEMM kernels. Currently coming from Megatron, only weights are non-contiguous after quantization. I have some ideas on how we can make that possible, or even better, that I am planning on working on in a followup ticket.

wangye805 · 2026-07-17T03:26:14Z

  diff algo_tune.csv algo_tune_check.csv


+Grouped GEMM Backends on ROCm


Technically multi-stream based hipblaslt/hipkitten is another backend. Maybe we should say explicitly "single-launch" grouped gemm in the section title

I think technically speaking multistream is not really a grouped gemm, but a set of gemms... -- I see where you are coming from though and have added that clarification!

wangye805 · 2026-07-17T03:28:28Z

+
+The following environment variables control backend selection:
+
+* ``NVTE_USE_CUTLASS_GROUPED_GEMM=1`` -- Enable the grouped GEMM path. On gfx950, defaults to HipKittens with CK fallback.


Is this env by default 1?

No, this is unset by default. I have added a note at the top of this subsection.

wangye805 · 2026-07-17T03:30:16Z

        """Return 64 MiB for gfx50x, 32 MiB for all other architectures."""
        if get_device_compute_capability() == (9, 5):
-            return 67_108_864
+            return 134_217_728


Emm, this workspace size was originally for hipblaslt

Sorry, I was doing some experiments and forgot to revert this. Should be back to 64 MB now.

wangye805 · 2026-07-17T03:32:39Z

+
+    auto *ws = convertNVTETensorCheck(workspace[0]);
+
+    return kittens_grouped_mxfp8_wgrad(


Just curious, does hipkitten wgrad need continuous b and b scale requirement?

No, as this kernel was designed fully per-expert due to the varying K dimension.

wangye805 · 2026-07-17T03:35:39Z

+}
+
+static size_t align256(size_t x) {
+    return (x + 255) & ~(size_t)255;


nit: since hipkittens kernels are a standalone lib.so, we can use its own env prefix, like HK_. Or still use NVTE_ but pass it from libte.so into libhipkitten.so

Sounds good, I have added hk prefix. I am hoping to avoid introducing a dependency on libte to keep cmake/builds clean

wangye805 · 2026-07-17T03:40:31Z

+
+template <bool ACCUMULATE, typename OutGL>
+__global__ __launch_bounds__(NUM_THREADS, 2)
+void mxfp8_wgrad_nt_kernel(


Emm, is this wgrad_nt_kernel based on/modified from regular nt_kernel? If so, can you have some comments on how the new one is derived.

If this is a complete new one, maybe have some tiling + pipelining comments or a pointer to our internal doc?

Done. This is similar to our regular nt kernel but the varying K dimension means that templating was getting messy so I separated it out. Also, a 1D grid performed better for the nature of expert wgrad where k is relatively small to m and n.

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

alextmagro added 2 commits July 17, 2026 02:13

NT kernel and bugfixes

fd55913

Merge remote-tracking branch 'origin/dev' into hipkittens_groupedgemm

7da3eb0

wangye805 requested changes Jul 17, 2026

View reviewed changes

alextmagro added 2 commits July 19, 2026 13:34

fused scale packing kernel

54b8faf

Added comments and var renaming

3fbd06e

alextmagro requested review from aris134 and wangye805 July 20, 2026 19:36

	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;
	}
	}

		diff algo_tune.csv algo_tune_check.csv


		Grouped GEMM Backends on ROCm


		The following environment variables control backend selection:

		* ``NVTE_USE_CUTLASS_GROUPED_GEMM=1`` -- Enable the grouped GEMM path. On gfx950, defaults to HipKittens with CK fallback.


		auto *ws = convertNVTETensorCheck(workspace[0]);

		return kittens_grouped_mxfp8_wgrad(

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!

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!

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!

wangye805 left a comment

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

aris134 Jul 8, 2026 •

edited

Loading

aris134 Jul 13, 2026 •

edited

Loading

aris134 Jul 13, 2026 •

edited

Loading