submission 634359 · KernelIndex

Use it

Vendorable · source mirrored · license unknownView source →

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

submission.py

curl "https://kernelindex.com/api/v1/implementations/kernelbot-amd-moe-mxfp4-634359?include=source"

interfacepython

Compatibility

measured onAMD Instinct MI355X

declared hardwareAMD Instinct MI355X

architecturesgfx950

dtypesbf16, fp32, fp8_e8m0, int32, mxfp4

Benchmark evidence

1 measurement across 1 GPU, fastest first.

Operation / workload

Hardware

Latency

Rank

Observed

AMD MXFP4 MoEsuite of 7 cases

AMD Instinct MI355X

154.9µs

#238 of 782

2026-03-25

Run →

Reported · How evidence levels are derived →

Source and license

sourceavailable

revision digestsha256:74f49a22d38bf58bcd68b957d5945d792175fae210d84617665c77584f85ef59

license declaredunknown

license concludedunknown

authorsHamza

imported2026-08-15

Kernel source

submission.py65 lines

#!POPCORN leaderboard amd-moe-mxfp4
#!POPCORN gpu MI355X

import os
import torch
from task import input_t, output_t

from aiter import ActivationType, QuantType
from aiter.fused_moe import fused_moe

_PAD_CACHE = {}


def custom_kernel(data: input_t) -> output_t:
    (
        hidden_states,
        gate_up_weight,
        down_weight,
        gate_up_weight_scale,
        down_weight_scale,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        gate_up_weight_scale_shuffled,
        down_weight_scale_shuffled,
        topk_weights,
        topk_ids,
        config,
    ) = data

    M = topk_ids.shape[0]

    cfg_id = id(config)
    cached = _PAD_CACHE.get(cfg_id)
    if cached is not None:
        hidden_pad, intermediate_pad = cached
    else:
        hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
        intermediate_pad = config["d_expert_pad"] - config["d_expert"]
        _PAD_CACHE[cfg_id] = (hidden_pad, intermediate_pad)

    # Per-shape backend selection via env vars (LRU-cached on first call per shape).
    # M<=128: cktile (ksplit=2) skips fp4 quantization entirely.
    # M>128: default ck2stages with CSV-tuned kernels.
    if M <= 128:
        os.environ["AITER_KSPLIT"] = "2"
        os.environ["AITER_BYPASS_TUNE_CONFIG"] = "1"
        # For M=128 with fewer experts (<=64), block_m=32 halves block count
        # vs cktile default 16, improving CU utilization.
        num_experts = gate_up_weight_shuffled.shape[0]
        block_size_M = 32 if (M > 16 and num_experts <= 64) else None
    else:
        os.environ["AITER_KSPLIT"] = "0"
        os.environ["AITER_BYPASS_TUNE_CONFIG"] = "0"
        block_size_M = None

    return fused_moe(
        hidden_states, gate_up_weight_shuffled, down_weight_shuffled,
        topk_weights, topk_ids,
        activation=ActivationType.Silu, quant_type=QuantType.per_1x32,
        w1_scale=gate_up_weight_scale_shuffled,
        w2_scale=down_weight_scale_shuffled,
        hidden_pad=hidden_pad, intermediate_pad=intermediate_pad,
        block_size_M=block_size_M,
    )

scrolls · 65 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 632161.

  ⋯ 28 unchanged lines
   
      M = topk_ids.shape[0]
   
-     # Per-shape kernel backend selection via env vars.
-     # get_2stage_cfgs() and get_ksplit() are LRU-cached per shape,
-     # so env vars only matter on first (warmup) call.
-     if M <= 128:
-         # Trigger cktile backend with split_k=2 for decode shapes.
-         # BYPASS_TUNE_CONFIG=1 skips CSV kernel selection, forcing
-         # the default heuristic path which respects AITER_KSPLIT.
-         os.environ["AITER_KSPLIT"] = "2"
-         os.environ["AITER_BYPASS_TUNE_CONFIG"] = "1"
-     else:
-         # Keep CSV-tuned ck2stages for large-M 256-expert shapes,
-         # and default ck backend for large-M 32-expert shapes.
-         os.environ["AITER_KSPLIT"] = "0"
-         os.environ["AITER_BYPASS_TUNE_CONFIG"] = "0"
-  
      cfg_id = id(config)
      cached = _PAD_CACHE.get(cfg_id)
      if cached is not None:
  ⋯ 3 unchanged lines
          intermediate_pad = config["d_expert_pad"] - config["d_expert"]
          _PAD_CACHE[cfg_id] = (hidden_pad, intermediate_pad)
   
+     # Per-shape backend selection via env vars (LRU-cached on first call per shape).
+     # M<=128: cktile (ksplit=2) skips fp4 quantization entirely.
+     # M>128: default ck2stages with CSV-tuned kernels.
+     if M <= 128:
+         os.environ["AITER_KSPLIT"] = "2"
+         os.environ["AITER_BYPASS_TUNE_CONFIG"] = "1"
+         # For M=128 with fewer experts (<=64), block_m=32 halves block count
+         # vs cktile default 16, improving CU utilization.
+         num_experts = gate_up_weight_shuffled.shape[0]
+         block_size_M = 32 if (M > 16 and num_experts <= 64) else None
+     else:
+         os.environ["AITER_KSPLIT"] = "0"
+         os.environ["AITER_BYPASS_TUNE_CONFIG"] = "0"
+         block_size_M = None
+  
      return fused_moe(
          hidden_states, gate_up_weight_shuffled, down_weight_shuffled,
          topk_weights, topk_ids,
  ⋯ 1 unchanged lines
          w1_scale=gate_up_weight_scale_shuffled,
          w2_scale=down_weight_scale_shuffled,
          hidden_pad=hidden_pad, intermediate_pad=intermediate_pad,
+         block_size_M=block_size_M,
      )

scrolls · 50 diff lines total

Best evidence level for this revision: reported

JSON