Skip to content
KernelIndex
Search⌘K
Sign in

submission 529737

.jonnss · python · License unknown
Kernel source · 61 lines ↓

Use it

Vendorable · source mirrored · license unknownView source →

No package. Vendor the mirrored source: 61 lines, June 9 Researcher Reciprocity License v1.0.

Submission_v1.py
curl "https://kernelindex.com/api/v1/implementations/kernelbot-amd-mxfp4-mm-529737?include=source"
interfacepython
Compatibility
measured onAMD Instinct MI355X
declared hardwareAMD Instinct MI355X
architecturesgfx950
dtypesbf16, mxfp4

Benchmark evidence

1 measurement across 1 GPU, fastest first.

Operation / workload
Hardware
Latency
Rank
Observed
AMD MXFP4 GEMMsuite of 6 cases
AMD Instinct MI355X
14.9µs
#549 of 1143
2026-03-11
Run →

Reported · How evidence levels are derived →

Source and license

sourceavailable
revision digestsha256:41d6b56761d0db36fa6c2c21fb30e55c4199598422db04b9b08dac2cf660ba9d
license declaredunknown
license concludedunknown
authors.jonnss
imported2026-08-15

Techniques

Extracted from the mirrored source by pattern, never inferred. Each row cites its line.

fp4FP4 quant + FP4 GEMM reference: bf16 A, MXFP4 B -> MXFP4 per-1x32 quant A -> gemm_a4w4 -> bf16 C.

Kernel source

Submission_v1.py61 lines
#!POPCORN leaderboard amd-mxfp4-mm
"""
FP4 quant + FP4 GEMM reference: bf16 A, MXFP4 B -> MXFP4 per-1x32 quant A -> gemm_a4w4 -> bf16 C.
Quant logic follows aiter op_tests/test_gemm_a4w4.py (get_triton_quant(QuantType.per_1x32)).
"""
import os

import torch

from task import input_t, output_t

VARIANT = os.environ.get("MXFP4_MM_VARIANT", "baseline")
SPLIT_ROW_M = 32
SPLIT_ROW_K = 512
SPLIT_CHUNK_M = 16


def _run_quant_gemm(aiter, quant_func, dtypes, A: torch.Tensor, B_shuffle: torch.Tensor, B_scale_sh: torch.Tensor) -> torch.Tensor:
    A_q, A_scale_sh = quant_func(A, shuffle=True)
    return aiter.gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=dtypes.bf16,
        bpreshuffle=True,
    )


def _run_split_row_two_pass(
    aiter,
    quant_func,
    dtypes,
    A: torch.Tensor,
    B_shuffle: torch.Tensor,
    B_scale_sh: torch.Tensor,
) -> torch.Tensor:
    top = _run_quant_gemm(aiter, quant_func, dtypes, A[:SPLIT_CHUNK_M], B_shuffle, B_scale_sh)
    bottom = _run_quant_gemm(aiter, quant_func, dtypes, A[SPLIT_CHUNK_M:], B_shuffle, B_scale_sh)
    return torch.cat((top, bottom), dim=0)


def custom_kernel(data: input_t) -> output_t:
    """
    Reference: MXFP4 per-1x32 quant on A; B_shuffle, B_scale_sh from generate_input.
    gemm_a4w4 with bpreshuffle=True.
    """
    import aiter
    from aiter import QuantType, dtypes

    A, _B, _B_q, B_shuffle, B_scale_sh = data
    A = A.contiguous()
    m, k = A.shape

    quant_func = aiter.get_triton_quant(QuantType.per_1x32)
    if VARIANT == "shape_aware":
        if m == SPLIT_ROW_M and k == SPLIT_ROW_K:
            return _run_split_row_two_pass(aiter, quant_func, dtypes, A, B_shuffle, B_scale_sh)

    return _run_quant_gemm(aiter, quant_func, dtypes, A, B_shuffle, B_scale_sh)
scrolls · 61 lines total

Source code from GPU Mode and the KernelBot dataset · June 9 Researcher Reciprocity License v1.0

Changes from previous submission

Against this author's previous submission submission 526261.

  ⋯ 9 unchanged lines
  from task import input_t, output_t
   
  VARIANT = os.environ.get("MXFP4_MM_VARIANT", "baseline")
- SMALL_M_PAD_THRESHOLD = 32
- PADDED_M = 64
+ SPLIT_ROW_M = 32
+ SPLIT_ROW_K = 512
+ SPLIT_CHUNK_M = 16
   
   
- def _maybe_pad_a(A: torch.Tensor) -> tuple[torch.Tensor, int]:
-     orig_m = A.shape[0]
-     if orig_m > SMALL_M_PAD_THRESHOLD:
-         return A, orig_m
+ def _run_quant_gemm(aiter, quant_func, dtypes, A: torch.Tensor, B_shuffle: torch.Tensor, B_scale_sh: torch.Tensor) -> torch.Tensor:
+     A_q, A_scale_sh = quant_func(A, shuffle=True)
+     return aiter.gemm_a4w4(
+         A_q,
+         B_shuffle,
+         A_scale_sh,
+         B_scale_sh,
+         dtype=dtypes.bf16,
+         bpreshuffle=True,
+     )
   
-     padded = torch.zeros((PADDED_M, A.shape[1]), dtype=A.dtype, device=A.device)
-     padded[:orig_m].copy_(A)
-     return padded, orig_m
   
+ def _run_split_row_two_pass(
+     aiter,
+     quant_func,
+     dtypes,
+     A: torch.Tensor,
+     B_shuffle: torch.Tensor,
+     B_scale_sh: torch.Tensor,
+ ) -> torch.Tensor:
+     top = _run_quant_gemm(aiter, quant_func, dtypes, A[:SPLIT_CHUNK_M], B_shuffle, B_scale_sh)
+     bottom = _run_quant_gemm(aiter, quant_func, dtypes, A[SPLIT_CHUNK_M:], B_shuffle, B_scale_sh)
+     return torch.cat((top, bottom), dim=0)
   
+  
  def custom_kernel(data: input_t) -> output_t:
      """
      Reference: MXFP4 per-1x32 quant on A; B_shuffle, B_scale_sh from generate_input.
  ⋯ 4 unchanged lines
   
      A, _B, _B_q, B_shuffle, B_scale_sh = data
      A = A.contiguous()
+     m, k = A.shape
   
-     if VARIANT == "small_m_pad":
-         A_work, orig_m = _maybe_pad_a(A)
-     else:
-         A_work, orig_m = A, A.shape[0]
-  
      quant_func = aiter.get_triton_quant(QuantType.per_1x32)
-     A_q, A_scale_sh = quant_func(A_work, shuffle=True)
-     out_gemm = aiter.gemm_a4w4(
-         A_q,
-         B_shuffle,
-         A_scale_sh,
-         B_scale_sh,
-         dtype=dtypes.bf16,
-         bpreshuffle=True,
-     )
-     return out_gemm[:orig_m]
+     if VARIANT == "shape_aware":
+         if m == SPLIT_ROW_M and k == SPLIT_ROW_K:
+             return _run_split_row_two_pass(aiter, quant_func, dtypes, A, B_shuffle, B_scale_sh)
+  
+     return _run_quant_gemm(aiter, quant_func, dtypes, A, B_shuffle, B_scale_sh)
scrolls · 73 diff lines total

Best evidence level for this revision: reported

JSON