submission 631834

  #!POPCORN leaderboard amd-mixed-mla
  #!POPCORN gpu MI355X
- """V1191: ps=4 ONLY for 64×8K and 256×8K (proven in test). ps=2 for everything else.
- Test covers 64×8K and 256×8K at ps=4, so these are safe. 4×8K/32×8K stay at ps=2."""
+ """V010i: v010d plus intra_batch_mode=False for 256x1k."""
  import os
  os.environ.setdefault("HIP_FORCE_DEV_KERNARG", "1")
  os.environ.setdefault("AMD_DIRECT_DISPATCH", "1")
  ⋯ 11 unchanged lines
      _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
- QSV=float(2.0/torch.finfo(FP8).max);_st={};_pc={};_nc={}
+ QSV=float(2.0/torch.finfo(FP8).max);_st={};_pc={}
  def _gs(d):
      t=_st.get(0)
      if t is None: t=torch.tensor([QSV],dtype=torch.float32,device=d);_st[0]=t
      return t
- def _gpc(dev,qo,bs,kvl,ps):
-     key=(bs,kvl,ps);c=_pc.get(key)
+ def _gpc(dev,qo,bs,kvl,ps,split_override=None,intra=True):
+     key=(bs,kvl,ps,split_override,intra);c=_pc.get(key)
      if c: return c
-     tot=bs*kvl;ppb=kvl//ps;np_=tot//ps;ns,_=aiter_mla.get_meta_param(None,bs,tot,NH,1,FP8)
+     tot=bs*kvl;ppb=kvl//ps;np_=tot//ps
+     ns = split_override if split_override is not None else aiter_mla.get_meta_param(None,bs,tot,NH,1,FP8)[0]
      ki=torch.arange(np_,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=max(ps,16)
-     info=_mi(bs,1,NH,FP8,FP8,is_sparse=False,fast_mode=True,num_kv_splits=ns,intra_batch_mode=True)
+     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=True,dtype_q=FP8,dtype_kv=FP8)
-     qf=torch.empty((bs,NH,QKD),dtype=FP8,device=dev);c=(ki,kl,kip,out,w[0],w[1],w[2],w[3],w[4],w[5],ns,qf);_pc[key]=c;return c
- def _gnc(dev,bs,kvl,ps):
-     key=(bs,kvl,ps);c=_nc.get(key)
-     if c: return c
-     tot=bs*kvl;ppb=kvl//ps;np_=tot//ps
-     ki=torch.arange(np_,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)
-     qf=torch.empty((bs,NH,QKD),dtype=FP8,device=dev);c=(ki,kl,kip,out,qf);_nc[key]=c;return c
- _NP_SHAPES = {(4,1024), (32,1024), (256,1024)}
- # Use ps=4 only for shapes proven in test (64×8K and 256×8K)
- _PS4_SHAPES = {(64,8192), (256,8192)}
+     _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)
+     qf=torch.empty((bs,NH,QKD),dtype=FP8,device=dev);c=(ki,kl,kip,out,w[0],w[1],w[2],w[3],w[4],w[5],ns,qf,intra);_pc[key]=c;return c
+ _SPLIT_CFG = {
+     (4,8192): 6,    # was 4, try 6 (more K-parallelism)
+     (32,8192): 6,   # proven
+     (64,8192): 4,   # was 2, try 4 (more K-parallelism)
+ }
  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"]
-     ps = 4 if (bs,kvl) in _PS4_SHAPES else 2
+     ps = 4 if kvl == 8192 else 2
      kb=kv_fp8.view(-1,ps,1,QKD);qs=_gs(dev)
-     if (bs,kvl) in _NP_SHAPES:
-         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]
-     c=_gpc(dev,qo_indptr,bs,kvl,ps);_quant(c[11],q,qs)
-     _fwd(c[11],kb,c[3],qo_indptr,c[2],c[0],c[1],1,ps,1,SM,num_kv_splits=c[10],work_meta_data=c[4],work_indptr=c[5],work_info_set=c[6],reduce_indptr=c[7],reduce_final_map=c[8],reduce_partial_map=c[9],q_scale=qs,kv_scale=kv_scale,intra_batch_mode=True)
+     split_ov = _SPLIT_CFG.get((bs,kvl))
+     intra = False if (bs,kvl) in {(32,8192), (256,1024)} else True
+     c=_gpc(dev,qo_indptr,bs,kvl,ps,split_ov,intra);_quant(c[11],q,qs)
+     _fwd(c[11],kb,c[3],qo_indptr,c[2],c[0],c[1],1,ps,1,SM,num_kv_splits=c[10],work_meta_data=c[4],work_indptr=c[5],work_info_set=c[6],reduce_indptr=c[7],reduce_final_map=c[8],reduce_partial_map=c[9],q_scale=qs,kv_scale=kv_scale,intra_batch_mode=c[12])
      return c[3]