submission 736348

  #!POPCORN leaderboard amd-mixed-mla
  #!POPCORN gpu MI355X
- """V029a: v027m + 256x8K ps=1024.
- Fresh namespace copy of the deep 256x8K page ladder on top of the now-safe v027m stack."""
+ """V029n: Push ALL 8K shapes to maximum page sizes.
+ - 4x8K: ps=2048 (8192/2048=4 pages, 4*4=16 total)
+ - 32x8K: ps=1024 (8192/1024=8 pages, 32*8=256 total)
+ - 64x8K: ps=512 (8192/512=16 pages, 64*16=1024 total)
+ - 256x8K: ps=2048 (proven in v029b leaderboard)
+ 1K shapes: best proven configs from v027m."""
  import os
  os.environ.setdefault("HIP_FORCE_DEV_KERNARG", "1")
  os.environ.setdefault("AMD_DIRECT_DISPATCH", "1")
  os.environ.setdefault("HIPBLASLT_ALLOW_FLUSH_DENORM", "1")
  os.environ.setdefault("GPU_MAX_HW_QUEUES", "2")
- import torch, triton, triton.language as tl
+ import torch
  from task import input_t, output_t
  import aiter
- from aiter import dtypes as aiter_dtypes, mla as aiter_mla
+ from aiter import dtypes as _ad, mla as _am
  try:
      from aiter.jit.module_quant import static_per_tensor_quant as _quant
  except Exception:
  ⋯ 10 unchanged lines
      from aiter.jit.module_mla_metadata import get_mla_metadata_v1 as _mv1
  except Exception:
      _mv1 = aiter.get_mla_metadata_v1
- _mi = aiter.get_mla_metadata_info_v1
- _fwd = aiter.mla.mla_decode_fwd
- NH=16;QKD=576;VD=512;SM=float(1.0/(QKD**0.5));FP8=aiter_dtypes.fp8;BF16=torch.bfloat16
+ _mi=aiter.get_mla_metadata_info_v1;_fwd=aiter.mla.mla_decode_fwd
+ NH=16;QKD=576;VD=512;SM=float(1.0/(QKD**0.5));FP8=_ad.fp8;BF16=torch.bfloat16
  QSV=float(2.0/torch.finfo(FP8).max);_st={};_nc={};_pc={};_s2c={}
  def _gs(d):
      t=_st.get(0)
  ⋯ 13 unchanged lines
      key=(bs,kvl,ps,intra,split_override,kg_override);c=_pc.get(key)
      if c: return c
      tot=bs*kvl;ppb=kvl//ps
-     ns = split_override if split_override is not None else aiter_mla.get_meta_param(None,bs,tot,NH,1,FP8)[0]
+     ns=split_override if split_override else _am.get_meta_param(None,bs,tot,NH,1,FP8)[0]
      ki=torch.arange(bs*ppb,dtype=torch.int32,device=dev);kl=torch.full((bs,),ps,dtype=torch.int32,device=dev)
      kip=torch.arange(bs+1,dtype=torch.int32,device=dev)*ppb;out=torch.empty((bs,NH,VD),dtype=BF16,device=dev)
-     kg=kg_override if kg_override is not None else max(ps,16)
+     kg=kg_override if kg_override else max(ps,16)
      info=_mi(bs,1,NH,FP8,FP8,is_sparse=False,fast_mode=True,num_kv_splits=ns,intra_batch_mode=intra)
      w=[torch.empty(s,dtype=t,device=dev) for s,t in info]
      _mv1(qo,kip,kl,16,1,False,w[0],w[2],w[1],w[3],w[4],w[5],page_size=ps,kv_granularity=kg,max_seqlen_qo=1,uni_seqlen_qo=1,fast_mode=True,max_split_per_batch=ns,intra_batch_mode=intra,dtype_q=FP8,dtype_kv=FP8)
  ⋯ 1 unchanged lines
      pt=int(w[5].numel());po=torch.empty((pt,1,NH,VD),dtype=torch.float32,device=dev);pl=torch.empty((pt,1,NH,1),dtype=torch.float32,device=dev)
      c=(ki,kl,kip,out,w[0],w[1],w[2],w[3],w[4],w[5],ns,qf,po,pl);_pc[key]=c;return c
  def _g_np(dev,bs,kvl,ps):
-     key=(bs,kvl,ps,"npm");c=_s2c.get(key)
+     key=(bs,kvl,ps,"npn2");c=_s2c.get(key)
      if c: return c
      ppb=kvl//ps
      ki=torch.arange(bs*ppb,dtype=torch.int32,device=dev)
  ⋯ 5 unchanged lines
      sd=torch.empty((bs,1,NH,VD),dtype=torch.float32,device=dev)
      sl=torch.empty((bs,1,NH,1),dtype=torch.float32,device=dev)
      c=(ki,kl,kip,out,qf,nsi,sd,sl);_s2c[key]=c;return c
+ # Per-shape optimal page sizes
+ _8K_PS = {4: 2048, 32: 1024, 64: 512, 256: 2048}
  def custom_kernel(data: input_t) -> output_t:
      q,kv_data,qo_indptr,kv_indptr,config=data
      bs=int(config["batch_size"]);kvl=int(config["kv_seq_len"]);dev=q.device
-     kv_fp8,kv_scale=kv_data["fp8"];qs=_gs(dev)
-     if kvl == 1024 and bs <= 32:
-         ps=2;kb=kv_fp8.view(-1,ps,1,QKD);c=_gnc(dev,bs,kvl,ps);_quant(c[4],q,qs)
-         _fwd(c[4],kb,c[3],qo_indptr,c[2],c[0],c[1],1,ps,1,SM,q_scale=qs,kv_scale=kv_scale,intra_batch_mode=True)
-         return c[3]
-     if bs == 64 and kvl == 1024:
-         ps=2;kb=kv_fp8.view(-1,ps,1,QKD)
-         c=_gpc(dev,qo_indptr,bs,kvl,ps,intra=True,split_override=1,kg_override=4)
-         _quant(c[11],q,qs)
-         _s1(c[11],kb,qo_indptr,c[2],c[0],c[1],None,c[4],c[5],c[6],1,ps,1,SM,c[12],c[13],c[3],qs,kv_scale)
-         if c[10]>1: _rd(c[12],c[13],c[7],c[8],c[9],1,c[3],None)
-         return c[3]
-     if bs == 256 and kvl == 1024:
-         ps=2;kb=kv_fp8.view(-1,ps,1,QKD)
-         c=_gpc(dev,qo_indptr,bs,kvl,ps,intra=True,split_override=1,kg_override=2)
-         _quant(c[11],q,qs)
-         _s1(c[11],kb,qo_indptr,c[2],c[0],c[1],None,c[4],c[5],c[6],1,ps,1,SM,c[12],c[13],c[3],qs,kv_scale)
-         if c[10]>1: _rd(c[12],c[13],c[7],c[8],c[9],1,c[3],None)
-         return c[3]
-     # ALL 8K: keep 128 except push 256x8K to 1024
-     ps=1024 if bs == 256 else 128;kb=kv_fp8.view(-1,ps,1,QKD);c=_g_np(dev,bs,kvl,ps);_quant(c[4],q,qs)
-     c[3].zero_(); c[7].fill_(-float("inf"))
-     _s1(c[4],kb,qo_indptr,c[2],c[0],c[1],c[5],None,None,None,1,ps,1,SM,c[6],c[7],c[3],qs,kv_scale)
+     kv,kvs=kv_data["fp8"];qs=_gs(dev)
+     if kvl==1024:
+         if bs<=32:
+             c=_gnc(dev,bs,kvl,2);_quant(c[4],q,qs)
+             _fwd(c[4],kv.view(-1,2,1,QKD),c[3],qo_indptr,c[2],c[0],c[1],1,2,1,SM,q_scale=qs,kv_scale=kvs,intra_batch_mode=True)
+             return c[3]
+         elif bs==64:
+             c=_gpc(dev,qo_indptr,64,1024,2,True,1,4);_quant(c[11],q,qs)
+             _s1(c[11],kv.view(-1,2,1,QKD),qo_indptr,c[2],c[0],c[1],None,c[4],c[5],c[6],1,2,1,SM,c[12],c[13],c[3],qs,kvs)
+             if c[10]>1:_rd(c[12],c[13],c[7],c[8],c[9],1,c[3],None)
+             return c[3]
+         else:
+             c=_gpc(dev,qo_indptr,256,1024,2,True,1,2);_quant(c[11],q,qs)
+             _s1(c[11],kv.view(-1,2,1,QKD),qo_indptr,c[2],c[0],c[1],None,c[4],c[5],c[6],1,2,1,SM,c[12],c[13],c[3],qs,kvs)
+             if c[10]>1:_rd(c[12],c[13],c[7],c[8],c[9],1,c[3],None)
+             return c[3]
+     # 8K: NP ns=1 with per-shape optimal large ps
+     ps=_8K_PS[bs]
+     c=_g_np(dev,bs,kvl,ps);_quant(c[4],q,qs)
+     c[3].zero_();c[7].fill_(float("-inf"))
+     _s1(c[4],kv.view(-1,ps,1,QKD),qo_indptr,c[2],c[0],c[1],c[5],None,None,None,1,ps,1,SM,c[6],c[7],c[3],qs,kvs)
      return c[3]