submission 737583 · KernelIndex

Use it

Vendorable · source mirrored · license unknownView source →

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

submission.py

curl "https://kernelindex.com/api/v1/implementations/kernelbot-amd-mixed-mla-737583?include=source"

interfacepython

Compatibility

measured onAMD Instinct MI355X

declared hardwareAMD Instinct MI355X

architecturesgfx950

dtypesbf16, int32

Benchmark evidence

1 measurement across 1 GPU, fastest first.

Operation / workload

Hardware

Latency

Rank

Observed

AMD mixed-precision MLA decodesuite of 8 cases

AMD Instinct MI355X

33.7µs

#48 of 766

2026-04-05

Run →

Reported · How evidence levels are derived →

Source and license

sourceavailable

revision digestsha256:62834fde051f68feb7644fb50b2feb810f87ce3143302376debf68a59efaf2e1

license declaredunknown

license concludedunknown

authorsmars-compute

imported2026-08-15

Kernel source

submission.py50 lines

import os
os.environ["HIP_FORCE_DEV_KERNARG"]="1";os.environ["HSA_ENABLE_SDMA"]="0"
import torch
from task import input_t, output_t
import aiter
from aiter import dtypes as aiter_dtypes
from aiter import get_mla_metadata_info_v1, get_mla_metadata_v1
NH=16;NKH=1;QKD=576;VD=512;SM=1.0/(QKD**0.5);FP8=aiter_dtypes.fp8
_stage1=aiter.mla_decode_stage1_asm_fwd
_reduce=aiter.mla_reduce_v1
_qs=None;_cache={}
# kv=1024: V188 proven configs; kv=8192: tuned ps=8 configs
_CFG={
    (4,1024):(1,16),(4,8192):(8,10),
    (32,1024):(2,2),(32,8192):(8,10),
    (64,1024):(2,4),(64,8192):(8,2),
    (256,1024):(2,2),(256,8192):(8,1),
}
def _build(bs,kvs,qd,kd,ps,nks):
    npp=kvs//ps
    qo=torch.arange(0,bs+1,dtype=torch.int32,device="cuda")
    kvi=torch.arange(0,bs+1,dtype=torch.int32,device="cuda")*npp
    klp=torch.full((bs,),ps,dtype=torch.int32,device="cuda")
    kvidx=torch.arange(bs*npp,dtype=torch.int32,device="cuda")
    info=get_mla_metadata_info_v1(bs,1,NH,qd,kd,is_sparse=False,fast_mode=False,num_kv_splits=nks,intra_batch_mode=True)
    w=[torch.empty(s,dtype=t,device="cuda") for s,t in info]
    wm,wi,wis,ri,rfm,rpm=w
    get_mla_metadata_v1(qo,kvi,klp,NH,NKH,True,wm,wis,wi,ri,rfm,rpm,
        page_size=ps,kv_granularity=max(ps,16),max_seqlen_qo=1,uni_seqlen_qo=1,
        fast_mode=False,max_split_per_batch=nks,intra_batch_mode=True,dtype_q=qd,dtype_kv=kd)
    lg=torch.empty((rpm.size(0),1,NH,VD),dtype=torch.float32,device="cuda")
    al=torch.empty((rpm.size(0),1,NH,1),dtype=torch.float32,device="cuda")
    o=torch.empty((bs,NH,VD),dtype=torch.bfloat16,device="cuda")
    return (wm,wi,wis,ri,rfm,rpm,kvidx,klp,qo,kvi,lg,al,o,ps)
def custom_kernel(data:input_t)->output_t:
    global _qs
    q,kv_data,_,_2,cfg=data
    bs=cfg["batch_size"];kvs=cfg["kv_seq_len"]
    kf,ks=kv_data["fp8"];qf=q.to(FP8)
    if _qs is None:_qs=torch.ones(1,dtype=torch.float32,device="cuda")
    ps,nks=_CFG.get((bs,kvs),(2,32))
    key=(bs,kvs)
    if key not in _cache:_cache[key]=_build(bs,kvs,qf.dtype,kf.dtype,ps,nks)
    wm,wi,wis,ri,rfm,rpm,kvidx,klp,qo,kvi,lg,al,o,ps_=_cache[key]
    np_=kf.shape[0]//ps_;kv4d=kf.view(np_,ps_,NKH,QKD)
    _stage1(qf.view(-1,NH,QKD),kv4d,qo,kvi,kvidx,klp,None,wm,wi,wis,1,ps_,NKH,SM,lg,al,o,_qs,ks)
    _reduce(lg,al,ri,rfm,rpm,1,o,None)
    return o

scrolls · 50 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 721303.

- """V188 — V185 splits + ps=4 for kv=8192 + safer (64,1024) nks=4.
- kv=8192 ps=4 gives ~20-47% speedup. kv=1024 uses V79's proven ps."""
  import os
  os.environ["HIP_FORCE_DEV_KERNARG"]="1";os.environ["HSA_ENABLE_SDMA"]="0"
  import torch
  ⋯ 5 unchanged lines
  _stage1=aiter.mla_decode_stage1_asm_fwd
  _reduce=aiter.mla_reduce_v1
  _qs=None;_cache={}
- # (ps, nks) per config — ps=4 for kv=8192, V79 ps for kv=1024
+ # kv=1024: V188 proven configs; kv=8192: tuned ps=8 configs
  _CFG={
-     (4,1024):(1,16),(4,8192):(4,16),
-     (32,1024):(2,2),(32,8192):(4,8),
-     (64,1024):(2,4),(64,8192):(4,8),   # nks=4 for safety (V79 value)
-     (256,1024):(2,2),(256,8192):(4,4),
+     (4,1024):(1,16),(4,8192):(8,10),
+     (32,1024):(2,2),(32,8192):(8,10),
+     (64,1024):(2,4),(64,8192):(8,2),
+     (256,1024):(2,2),(256,8192):(8,1),
  }
  def _build(bs,kvs,qd,kd,ps,nks):
      npp=kvs//ps
  ⋯ 25 unchanged lines
      _stage1(qf.view(-1,NH,QKD),kv4d,qo,kvi,kvidx,klp,None,wm,wi,wis,1,ps_,NKH,SM,lg,al,o,_qs,ks)
      _reduce(lg,al,ri,rfm,rpm,1,o,None)
      return o
+

scrolls · 28 diff lines total

Best evidence level for this revision: reported

JSON