submission 676583 · KernelIndex

Use it

Vendorable · source mirrored · license unknownView source →

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

submission.py

curl "https://kernelindex.com/api/v1/implementations/kernelbot-amd-moe-mxfp4-676583?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

155.9µs

#251 of 782

2026-03-31

Run →

Reported · How evidence levels are derived →

Source and license

sourceavailable

revision digestsha256:215ddf7bfce936eb58e752e2222e2ed38136a2777310958466e08b6c5a91b87e

license declaredunknown

license concludedunknown

authorslgc0338

imported2026-08-15

Kernel source

submission.py45 lines

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

"""
Best per-shape strategy:
- Small batch (M≤128): KSPLIT=2 + BYPASS for E=257 → CKTile a16w4 (skip quant)
- Large batch (M>128): default CK kernels (no KSPLIT, no BYPASS)
"""
import os
import torch
from task import input_t, output_t
from aiter import ActivationType, QuantType
from aiter.fused_moe import fused_moe

_ACT = ActivationType.Silu
_QT = QuantType.per_1x32

@torch.inference_mode()
def custom_kernel(data: input_t) -> output_t:
    (hs, _,_,_,_, w1s,w2s,w1ss,w2ss, tw,ti, cfg) = data
    M = hs.shape[0]
    E = cfg["n_routed_experts"] + cfg["n_shared_experts"]
    dhp = cfg["d_hidden_pad"]
    dh = cfg["d_hidden"]
    dep = cfg["d_expert_pad"]
    de = cfg["d_expert"]

    if M <= 128:
        # Small batch: CKTile a16w4 path (skip quant, block_m=16)
        os.environ['AITER_KSPLIT'] = '2'
        if E > 64:
            os.environ['AITER_BYPASS_TUNE_CONFIG'] = '1'
        else:
            os.environ.pop('AITER_BYPASS_TUNE_CONFIG', None)
    else:
        # Large batch: default CK kernels
        os.environ.pop('AITER_KSPLIT', None)
        os.environ.pop('AITER_BYPASS_TUNE_CONFIG', None)

    return fused_moe(hs, w1s, w2s, tw, ti,
        expert_mask=None, activation=_ACT, quant_type=_QT,
        doweight_stage1=False, w1_scale=w1ss, w2_scale=w2ss,
        a1_scale=None, a2_scale=None,
        hidden_pad=dhp-dh, intermediate_pad=dep-de)

scrolls · 45 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 671802.

  #!POPCORN leaderboard amd-moe-mxfp4
  #!POPCORN gpu MI355X
   
+ """
+ Best per-shape strategy:
+ - Small batch (M≤128): KSPLIT=2 + BYPASS for E=257 → CKTile a16w4 (skip quant)
+ - Large batch (M>128): default CK kernels (no KSPLIT, no BYPASS)
+ """
  import os
  import torch
  from task import input_t, output_t
+ from aiter import ActivationType, QuantType
+ from aiter.fused_moe import fused_moe
   
- # Per-shape env var tuning
- def _set_env(M, E, dep):
-     """Set AITER env vars for optimal per-shape performance."""
-     # E=33 bs=16: SplitK helps (-36% from 96→62 μs in earlier tests)
-     if E <= 33 and M <= 16:
-         os.environ['AITER_KSPLIT'] = '2'
-     else:
-         os.environ.pop('AITER_KSPLIT', None)
+ _ACT = ActivationType.Silu
+ _QT = QuantType.per_1x32
   
  @torch.inference_mode()
  def custom_kernel(data: input_t) -> output_t:
      (hs, _,_,_,_, w1s,w2s,w1ss,w2ss, tw,ti, cfg) = data
-  
-     from aiter import ActivationType, QuantType
-     from aiter.fused_moe import fused_moe
-  
      M = hs.shape[0]
      E = cfg["n_routed_experts"] + cfg["n_shared_experts"]
      dhp = cfg["d_hidden_pad"]
  ⋯ 1 unchanged lines
      dep = cfg["d_expert_pad"]
      de = cfg["d_expert"]
   
-     _set_env(M, E, dep)
+     if M <= 128:
+         # Small batch: CKTile a16w4 path (skip quant, block_m=16)
+         os.environ['AITER_KSPLIT'] = '2'
+         if E > 64:
+             os.environ['AITER_BYPASS_TUNE_CONFIG'] = '1'
+         else:
+             os.environ.pop('AITER_BYPASS_TUNE_CONFIG', None)
+     else:
+         # Large batch: default CK kernels
+         os.environ.pop('AITER_KSPLIT', None)
+         os.environ.pop('AITER_BYPASS_TUNE_CONFIG', None)
   
-     return fused_moe(hs, w1s, w2s, tw, ti, expert_mask=None,
-         activation=ActivationType.Silu,
-         quant_type=QuantType.per_1x32,
-         doweight_stage1=False,
-         w1_scale=w1ss, w2_scale=w2ss,
+     return fused_moe(hs, w1s, w2s, tw, ti,
+         expert_mask=None, activation=_ACT, quant_type=_QT,
+         doweight_stage1=False, w1_scale=w1ss, w2_scale=w2ss,
          a1_scale=None, a2_scale=None,
-         hidden_pad=dhp-dh,
-         intermediate_pad=dep-de)
+         hidden_pad=dhp-dh, intermediate_pad=dep-de)

scrolls · 64 diff lines total

Best evidence level for this revision: reported

JSON