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
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={}
# (ps, nks) per config — ps=4 for kv=8192, V79 ps for kv=1024
_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),
}
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