submission 588910

  #!POPCORN leaderboard amd-mixed-mla
  #!POPCORN gpu MI355X
   
- """v707 - bf16 for small shapes, 1-split for (64,1K)+(256,1K), persistent for 8K."""
+ """V727: v722 with safer Q scale (0.20 instead of 0.14/0.15) for precision."""
   
  import os
   
  ⋯ 6 unchanged lines
  import aiter
  from aiter import dtypes as aiter_dtypes
  from aiter import mla as aiter_mla
- from aiter import get_mla_metadata_info_v1, get_mla_metadata_v1
   
  try:
-     from aiter.jit.module_quant import static_per_tensor_quant
+     from aiter.jit.module_quant import static_per_tensor_quant as _static_per_tensor_quant
  except Exception:
-     try:
-         from aiter.ops.quant import static_per_tensor_quant
-     except Exception:
-         static_per_tensor_quant = None
+     from aiter.ops.quant import static_per_tensor_quant as _static_per_tensor_quant
   
+ try:
+     from aiter.jit.module_mla_asm import mla_decode_stage1_asm_fwd as _mla_stage1
+ 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
   
- _cache = {}
+ _scale_tensors = {}
+ _direct_cache = {}
+ _decode_cache = {}
+ _nonpersist_cache = {}
   
+ _SPLITBOOST_8K = {
+     (32, 8192): 8,
+     (64, 8192): 4,
+     (256, 8192): 2,
+ }
   
- def _quantize_q_fp8(q, q_fp8_buf):
-     amax = q.abs().amax().clamp(min=1e-12)
-     scale = (amax / _FP8_FINFO.max).reshape(1).to(torch.float32)
-     if static_per_tensor_quant is not None:
-         static_per_tensor_quant(q_fp8_buf, q, scale)
-     else:
-         q_fp8_buf.copy_(
-             (q / scale).clamp(min=_FP8_FINFO.min, max=_FP8_FINFO.max).to(FP8_DTYPE)
-         )
-     return scale
   
+ 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
+     return t
   
- def _build_nonpersist_1split(dev, bs, kvlen, kv_indptr):
-     """Non-persistent mode with 1 split: output written directly, no reduce needed."""
-     total_kv = bs * kvlen
-     _, num_splits_indptr = aiter_mla.get_meta_param(1, 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=torch.bfloat16, device=dev)
-     # With 1 split, logits shares memory with out
-     logits = out.view(bs, 1, NUM_HEADS, V_HEAD_DIM)
-     attn_lse = torch.empty((bs, 1, NUM_HEADS, 1), dtype=torch.float32, device=dev)
-     q_fp8 = torch.empty((bs, NUM_HEADS, QK_HEAD_DIM), dtype=FP8_DTYPE, device=dev)
-     return (num_splits_indptr, kv_indices, kv_lpl, out, logits, attn_lse, q_fp8)
   
+ 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 _build_persistent(dev, qo_indptr, kv_indptr, bs, kvlen, kv_gran, intra):
-     """Persistent mode with metadata — uses mla_reduce_v1."""
+  
+ 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
-     num_splits, _ = aiter_mla.get_meta_param(None, bs, total_kv, NUM_HEADS, 1, FP8_DTYPE)
+     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=torch.bfloat16, 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,
-     )
-     bufs = [torch.empty(s, dtype=t, device=dev) for s, t in info]
-     wmd, wi, wis, ri, rfm, rpm = bufs
-     get_mla_metadata_v1(
-         qo_indptr, kv_indptr, kv_lpl, 16, 1, False,
+     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=FP8_DTYPE, dtype_kv=FP8_DTYPE,
-     )
+         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)
-     return (kv_indices, kv_lpl, out, wmd, wi, wis, ri, rfm, rpm, po, pl, q_fp8)
+     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 _build_bf16_persistent(dev, qo_indptr, kv_indptr, bs, kvlen, kv_gran, intra):
-     """Persistent mode with bf16 Q + bf16 KV — no quantization needed."""
+  
+ 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
-     num_splits, _ = aiter_mla.get_meta_param(None, bs, total_kv, NUM_HEADS, 1, torch.bfloat16)
+  
      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=torch.bfloat16, 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, torch.bfloat16, torch.bfloat16,
-         is_sparse=False, fast_mode=True, num_kv_splits=num_splits, intra_batch_mode=intra,
-     )
-     bufs = [torch.empty(s, dtype=t, device=dev) for s, t in info]
-     wmd, wi, wis, ri, rfm, rpm = bufs
-     get_mla_metadata_v1(
-         qo_indptr, kv_indptr, kv_lpl, 16, 1, False,
+     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=torch.bfloat16, dtype_kv=torch.bfloat16,
+         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)
      )
-     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)
-     return (kv_indices, kv_lpl, out, wmd, wi, wis, ri, rfm, rpm, po, pl)
+     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 custom_kernel(data: input_t) -> output_t:
      q, kv_data, qo_indptr, kv_indptr, config = data
-     bs = int(config["batch_size"])
-     kvlen = int(config["kv_seq_len"])
+     batch_size = int(config["batch_size"])
+     kv_seq_len = int(config["kv_seq_len"])
      sm_scale = float(config["sm_scale"])
-     dev = q.device
   
+     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
   
-     key = (dev.index, bs, kvlen)
-  
-     # --- bf16 path for small batches: skip Q quantization entirely ---
-     if bs <= 32 and kvlen == 1024:
-         bkey = ("bf16", *key)
-         if bkey not in _cache:
-             _cache[bkey] = _build_bf16_persistent(dev, qo_indptr, kv_indptr, bs, kvlen, 16, True)
-         c = _cache[bkey]
-         kv_bf16 = kv_data["bf16"].view(-1, 1, 1, QK_HEAD_DIM)
-         aiter.mla_decode_stage1_asm_fwd(
-             q, kv_bf16, 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,
-         )
-         aiter.mla_reduce_v1(c[9], c[10], c[6], c[7], c[8], 1, c[2], None)
+     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 bs == 4 and kvlen == 8192:
-         bkey = ("bf16", *key)
-         if bkey not in _cache:
-             _cache[bkey] = _build_bf16_persistent(dev, qo_indptr, kv_indptr, bs, kvlen, 64, True)
-         c = _cache[bkey]
-         kv_bf16 = kv_data["bf16"].view(-1, 1, 1, QK_HEAD_DIM)
-         aiter.mla_decode_stage1_asm_fwd(
-             q, kv_bf16, 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,
-         )
-         aiter.mla_reduce_v1(c[9], c[10], c[6], c[7], c[8], 1, c[2], None)
+     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]
   
-     # --- 1-split non-persistent for select shapes: skip reduce ---
-     if (bs, kvlen) in ((64, 1024), (256, 1024)):
-         nkey = ("np1", *key)
-         if nkey not in _cache:
-             _cache[nkey] = _build_nonpersist_1split(dev, bs, kvlen, kv_indptr)
-         c = _cache[nkey]
-         q_scale = _quantize_q_fp8(q, c[6])
-         aiter.mla_decode_stage1_asm_fwd(
-             c[6], kv_buf, qo_indptr, kv_indptr, c[1], c[2], c[0],
-             None, None, None, 1, 1, 1, sm_scale, c[4], c[5], c[3], q_scale, kv_scale,
-         )
-         return c[3]
-  
-     # --- Persistent mode for remaining shapes ---
-     pkey = ("persist", *key)
-     if pkey not in _cache:
-         kv_gran = 64 if kvlen == 8192 else 8
-         intra = True
-         _cache[pkey] = _build_persistent(dev, qo_indptr, kv_indptr, bs, kvlen, kv_gran, intra)
-     c = _cache[pkey]
-     q_scale = _quantize_q_fp8(q, c[11])
-     aiter.mla_decode_stage1_asm_fwd(
-         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,
-     )
-     aiter.mla_reduce_v1(c[9], c[10], c[6], c[7], c[8], 1, c[2], None)
+     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]