submission 605758

  #!POPCORN leaderboard amd-mixed-mla
  #!POPCORN gpu MI355X
-  
- """V727: v722 with safer Q scale (0.20 instead of 0.14/0.15) for precision."""
-  
+ """V1092: ALL ps=2 ALL fp8 with Q_scale=5.0/max (larger to avoid clipping)."""
  import os
-  
  os.environ.setdefault("HIP_FORCE_DEV_KERNARG", "1")
  os.environ.setdefault("AMD_DIRECT_DISPATCH", "1")
-  
  import torch
  from task import input_t, output_t
-  
  import aiter
- from aiter import dtypes as aiter_dtypes
- from aiter import mla as aiter_mla
-  
+ from aiter import dtypes as aiter_dtypes, mla as aiter_mla
  try:
-     from aiter.jit.module_quant import static_per_tensor_quant as _static_per_tensor_quant
+     from aiter.jit.module_quant import static_per_tensor_quant as _quant
  except Exception:
-     from aiter.ops.quant import static_per_tensor_quant as _static_per_tensor_quant
-  
+     from aiter.ops.quant import static_per_tensor_quant as _quant
  try:
-     from aiter.jit.module_mla_asm import mla_decode_stage1_asm_fwd as _mla_stage1
+     from aiter.jit.module_mla_metadata import get_mla_metadata_v1 as _mv1
  except Exception:
-     _mla_stage1 = aiter.mla_decode_stage1_asm_fwd
-  
- try:
-     from aiter.jit.module_mla_reduce import mla_reduce_v1 as _mla_reduce
- except Exception:
-     _mla_reduce = aiter.mla_reduce_v1
-  
- try:
-     from aiter.jit.module_mla_metadata import get_mla_metadata_v1 as _get_mla_metadata_v1
- except Exception:
-     _get_mla_metadata_v1 = aiter.get_mla_metadata_v1
-  
- _get_mla_metadata_info_v1 = aiter.get_mla_metadata_info_v1
- _mla_decode_fwd = aiter.mla.mla_decode_fwd
-  
- NUM_HEADS = 16
- NUM_KV_HEADS = 1
- QK_HEAD_DIM = 576
- V_HEAD_DIM = 512
- PAGE_SIZE = 1
- SM_SCALE = float(1.0 / (QK_HEAD_DIM ** 0.5))
-  
- FP8_DTYPE = aiter_dtypes.fp8
- BF16_DTYPE = torch.bfloat16
- _FP8_FINFO = torch.finfo(FP8_DTYPE)
- _Q020_SCALE = float(0.20 / _FP8_FINFO.max)
- _Q014_SCALE = _Q020_SCALE  # Use safer 0.20 scale for all shapes
- _Q015_SCALE = _Q020_SCALE
-  
- _scale_tensors = {}
- _direct_cache = {}
- _decode_cache = {}
- _nonpersist_cache = {}
-  
- _SPLITBOOST_8K = {
-     (32, 8192): 8,
-     (64, 8192): 4,
-     (256, 8192): 2,
- }
-  
-  
- def _get_scale(dev, scale):
-     key = (dev.index or 0, scale)
-     t = _scale_tensors.get(key)
-     if t is None:
-         t = torch.tensor([scale], dtype=torch.float32, device=dev)
-         _scale_tensors[key] = t
+     _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
+ # LARGER Q scale to avoid clipping (Q values can be up to ~4.5)
+ QSV=float(5.0/torch.finfo(FP8).max)
+ PS=2
+ _st={};_c={}
+ 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 _resolve_num_splits(batch_size, kv_seq_len, dtype_kv, split_override):
-     if split_override is not None:
-         return int(split_override)
-     num_splits, _ = aiter_mla.get_meta_param(None, batch_size, batch_size * kv_seq_len, NUM_HEADS, 1, dtype_kv)
-     return int(num_splits)
-  
-  
- def _quantize_q(out_fp8, q, q_scale):
-     _static_per_tensor_quant(out_fp8, q, q_scale)
-  
-  
- def _get_direct_cache(dev, qo_indptr, kv_indptr, bs, kvlen, dtype_q, dtype_kv, kv_gran, intra, causal, need_fp8):
-     key = (dev.index or 0, bs, kvlen, dtype_q, dtype_kv, kv_gran, intra, causal, need_fp8)
-     cached = _direct_cache.get(key)
-     if cached is not None:
-         return cached
-  
-     total_kv = bs * kvlen
-     split_dtype = FP8_DTYPE if dtype_kv == FP8_DTYPE else BF16_DTYPE
-     num_splits, _ = aiter_mla.get_meta_param(None, bs, total_kv, NUM_HEADS, 1, split_dtype)
-  
-     kv_indices = torch.arange(total_kv, dtype=torch.int32, device=dev)
-     kv_lpl = torch.full((bs,), kvlen, dtype=torch.int32, device=dev)
-     out = torch.empty((bs, NUM_HEADS, V_HEAD_DIM), dtype=BF16_DTYPE, device=dev)
-  
-     info = _get_mla_metadata_info_v1(bs, 1, NUM_HEADS, dtype_q, dtype_kv,
-         is_sparse=False, fast_mode=True, num_kv_splits=num_splits, intra_batch_mode=intra)
-     wmd, wi, wis, ri, rfm, rpm = [torch.empty(s, dtype=t, device=dev) for s, t in info]
-  
-     _get_mla_metadata_v1(qo_indptr, kv_indptr, kv_lpl, 16, 1, causal,
-         wmd, wis, wi, ri, rfm, rpm,
-         page_size=1, kv_granularity=kv_gran, max_seqlen_qo=1, uni_seqlen_qo=1,
-         fast_mode=True, max_split_per_batch=num_splits, intra_batch_mode=intra,
-         dtype_q=dtype_q, dtype_kv=dtype_kv)
-  
-     pt = int(rpm.numel())
-     po = torch.empty((pt, 1, NUM_HEADS, V_HEAD_DIM), dtype=torch.float32, device=dev)
-     pl = torch.empty((pt, 1, NUM_HEADS, 1), dtype=torch.float32, device=dev)
-     q_fp8 = torch.empty((bs, NUM_HEADS, QK_HEAD_DIM), dtype=FP8_DTYPE, device=dev) if need_fp8 else None
-  
-     cached = (kv_indices, kv_lpl, out, wmd, wi, wis, ri, rfm, rpm, po, pl, q_fp8)
-     _direct_cache[key] = cached
-     return cached
-  
-  
- def _get_decode_cache(dev, qo_indptr, kv_indptr, bs, kvlen, kv_gran, intra, split_override):
-     key = (dev.index or 0, bs, kvlen, kv_gran, intra, split_override)
-     cached = _decode_cache.get(key)
-     if cached is not None:
-         return cached
-  
-     num_splits = _resolve_num_splits(bs, kvlen, FP8_DTYPE, split_override)
-     total_kv = bs * kvlen
-  
-     kv_indices = torch.arange(total_kv, dtype=torch.int32, device=dev)
-     kv_lpl = torch.full((bs,), kvlen, dtype=torch.int32, device=dev)
-     out = torch.empty((bs, NUM_HEADS, V_HEAD_DIM), dtype=BF16_DTYPE, device=dev)
-  
-     info = _get_mla_metadata_info_v1(bs, 1, NUM_HEADS, FP8_DTYPE, FP8_DTYPE,
-         is_sparse=False, fast_mode=True, num_kv_splits=num_splits, intra_batch_mode=intra)
-     wmd, wi, wis, ri, rfm, rpm = [torch.empty(s, dtype=t, device=dev) for s, t in info]
-  
-     _get_mla_metadata_v1(qo_indptr, kv_indptr, kv_lpl, 16, 1, False,
-         wmd, wis, wi, ri, rfm, rpm,
-         page_size=1, kv_granularity=kv_gran, max_seqlen_qo=1, uni_seqlen_qo=1,
-         fast_mode=True, max_split_per_batch=num_splits, intra_batch_mode=intra,
-         dtype_q=FP8_DTYPE, dtype_kv=FP8_DTYPE)
-  
-     q_fp8 = torch.empty((bs, NUM_HEADS, QK_HEAD_DIM), dtype=FP8_DTYPE, device=dev)
-     cached = (kv_indices, kv_lpl, out, wmd, wi, wis, ri, rfm, rpm, num_splits, q_fp8)
-     _decode_cache[key] = cached
-     return cached
-  
-  
- def _get_nonpersistent_cache(dev, bs, kvlen, split_override):
-     key = (dev.index or 0, bs, kvlen, split_override)
-     cached = _nonpersist_cache.get(key)
-     if cached is not None:
-         return cached
-  
-     total_kv = bs * kvlen
-     num_splits, num_splits_indptr = aiter_mla.get_meta_param(split_override, bs, total_kv, NUM_HEADS, 1, FP8_DTYPE)
-     kv_indices = torch.arange(total_kv, dtype=torch.int32, device=dev)
-     kv_lpl = torch.full((bs,), kvlen, dtype=torch.int32, device=dev)
-     out = torch.empty((bs, NUM_HEADS, V_HEAD_DIM), dtype=BF16_DTYPE, device=dev)
-     logits = (
-         out.view(bs, 1, NUM_HEADS, V_HEAD_DIM)
-         if num_splits == 1
-         else torch.empty((bs, num_splits, NUM_HEADS, V_HEAD_DIM), dtype=torch.float32, device=dev)
-     )
-     attn_lse = torch.empty((bs, num_splits, NUM_HEADS, 1), dtype=torch.float32, device=dev)
-     q_fp8 = torch.empty((bs, NUM_HEADS, QK_HEAD_DIM), dtype=FP8_DTYPE, device=dev)
-  
-     cached = (num_splits, num_splits_indptr, kv_indices, kv_lpl, out, logits, attn_lse, q_fp8)
-     _nonpersist_cache[key] = cached
-     return cached
-  
-  
+ def _gc(dev,qo,bs,kvl):
+     key=(bs,kvl)
+     c=_c.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)
+     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)
+     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)
+     _c[key]=c; return c
  def custom_kernel(data: input_t) -> output_t:
-     q, kv_data, qo_indptr, kv_indptr, config = data
-     batch_size = int(config["batch_size"])
-     kv_seq_len = int(config["kv_seq_len"])
-     sm_scale = float(config["sm_scale"])
-  
-     if batch_size == 4 and kv_seq_len == 1024:
-         c = _get_direct_cache(q.device, qo_indptr, kv_indptr, 4, 1024,
-             BF16_DTYPE, BF16_DTYPE, 16, True, False, False)
-         kv_buf = kv_data["bf16"].view(-1, 1, 1, QK_HEAD_DIM)
-         _mla_stage1(
-             q, kv_buf, qo_indptr, kv_indptr, c[0], c[1], None,
-             c[3], c[4], c[5], 1, 1, 1, sm_scale, c[9], c[10], c[2], None, None)
-         _mla_reduce(c[9], c[10], c[6], c[7], c[8], 1, c[2], None)
-         return c[2]
-  
-     kv_fp8, kv_scale = kv_data["fp8"]
-     kv_buf = kv_fp8.view(-1, 1, 1, QK_HEAD_DIM)
-     dev = q.device
-  
-     if batch_size == 4 and kv_seq_len == 8192:
-         c = _get_direct_cache(dev, qo_indptr, kv_indptr, 4, 8192,
-             FP8_DTYPE, FP8_DTYPE, 64, True, False, True)
-         q_scale = _get_scale(dev, _Q014_SCALE)
-         _quantize_q(c[11], q, q_scale)
-         _mla_stage1(
-             c[11], kv_buf, qo_indptr, kv_indptr, c[0], c[1], None,
-             c[3], c[4], c[5], 1, 1, 1, sm_scale, c[9], c[10], c[2], q_scale, kv_scale)
-         _mla_reduce(c[9], c[10], c[6], c[7], c[8], 1, c[2], None)
-         return c[2]
-  
-     if kv_seq_len == 1024:
-         # ALL 1K shapes: persistent fp8 (safe — no non-persistent kernel)
-         q_scale_val = _Q015_SCALE if batch_size == 256 else _Q014_SCALE
-         c = _get_direct_cache(dev, qo_indptr, kv_indptr, batch_size, 1024,
-             FP8_DTYPE, FP8_DTYPE, 8, True, False, True)
-         q_scale = _get_scale(dev, q_scale_val)
-         _quantize_q(c[11], q, q_scale)
-         _mla_stage1(
-             c[11], kv_buf, qo_indptr, kv_indptr, c[0], c[1], None,
-             c[3], c[4], c[5], 1, 1, 1, sm_scale, c[9], c[10], c[2], q_scale, kv_scale)
-         _mla_reduce(c[9], c[10], c[6], c[7], c[8], 1, c[2], None)
-         return c[2]
-  
-     split_override = _SPLITBOOST_8K.get((batch_size, kv_seq_len))
-     c = _get_decode_cache(dev, qo_indptr, kv_indptr, batch_size, 8192, 16, False, split_override)
-     q_scale = _get_scale(dev, _Q014_SCALE)
-     _quantize_q(c[10], q, q_scale)
-     _mla_decode_fwd(
-         c[10], kv_buf, c[2], qo_indptr, kv_indptr, c[0], c[1],
-         1, 1, 1, sm_scale,
-         num_kv_splits=c[9],
-         work_meta_data=c[3], work_indptr=c[4], work_info_set=c[5],
-         reduce_indptr=c[6], reduce_final_map=c[7], reduce_partial_map=c[8],
-         q_scale=q_scale, kv_scale=kv_scale,
-         intra_batch_mode=False)
-     return c[2]
+     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"]
+     kb=kv_fp8.view(-1,PS,1,QKD)
+     c=_gc(dev,qo_indptr,bs,kvl)
+     qs=_gs(dev)
+     _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)
+     return c[3]