submission 590137

  #!POPCORN leaderboard amd-mixed-mla
  #!POPCORN gpu MI355X
  """
- Ultimate MLA decode: best path per shape.
- - bmm bf16: small shapes (bs<=4 all, bs<=64 kv<=1024)
- - Triton fp8 flash-decode: bs=32/kv=8192 (beats AITER by 30%)
- - AITER fp8: large shapes (bs>=64/kv=8192, bs=256/kv=1024)
+ Ultimate dispatch: best kernel per shape based on exhaustive benchmarking.
+ - bmm bf16: small shapes
+ - Triton fp8 (BLOCK_N=32, 8 splits): bs=32/kv=8192
+ - AITER bf16Q+fp8KV: bs=256/kv=1024 (136μs, -24% vs fp8+fp8)
+ - AITER fp8+fp8: bs=64/kv=8192, bs=256/kv=8192
  """
  import torch
  import torch.nn.functional as F
  ⋯ 14 unchanged lines
  PAGE_SIZE = 1
  FP8_DTYPE = aiter_dtypes.fp8
   
- _meta_cache = {}
- _idx_cache = {}
- _out_cache = {}
- _buf = {}
- _warm = set()
   
+ # ═══════════════ Triton FP8 Flash-Decode ═══════════════
   
- # ═══════════════════════════════════════════════
- # Triton flash-decode fp8 (for medium shapes)
- # ═══════════════════════════════════════════════
-  
  @triton.jit
  def _flash_s1(
      Q, KV_FP8, kv_scale_ptr, sm_scale,
  ⋯ 8 unchanged lines
      Lnope: tl.constexpr, Lrope: tl.constexpr, Lv: tl.constexpr,
  ):
      bid = tl.program_id(0)
-     sid = tl.program_id(1)
-  
+     sid = tl.program_id(2)
      heads = tl.arange(0, BLOCK_H)
      mask_h = heads < 16
      o_nope = tl.arange(0, BLOCK_NOPE)
  ⋯ 3 unchanged lines
      mn = o_nope < Lnope
      mr = o_rope < Lrope
      mv = o_dv < Lv
-  
      ks = tl.load(kv_indptr + bid)
      ke = tl.load(kv_indptr + bid + 1)
      kl = ke - ks
-  
      ss = tl.cdiv(kl, NUM_SPLITS)
      ss = tl.cdiv(ss, BLOCK_N) * BLOCK_N
      ms = sid * ss
      me = tl.minimum(ms + ss, kl)
-  
      emax = tl.zeros([BLOCK_H], dtype=tl.float32) - float("inf")
      esum = tl.zeros([BLOCK_H], dtype=tl.float32)
      acc = tl.zeros([BLOCK_H, BLOCK_DV], dtype=tl.float32)
      sc = tl.load(kv_scale_ptr)
-  
      if me > ms:
          qb = bid * stride_qb
          qn = tl.load(Q + qb + heads[:, None] * stride_qh + o_nope[None, :],
                        mask=mask_h[:, None] & mn[None, :], other=0.0).to(tl.float16)
          qr = tl.load(Q + qb + heads[:, None] * stride_qh + o_rope_s[None, :],
                        mask=mask_h[:, None] & mr[None, :], other=0.0).to(tl.float16)
-  
          for t in range(ms, me, BLOCK_N):
              no = tl.arange(0, BLOCK_N)
              nm = (t + no) < me
              ti = (ks + t + no) * stride_kv_tok
-  
-             kn = tl.load(KV_FP8 + ti[None, :] + o_nope[:, None],
-                          mask=nm[None, :] & mn[:, None], other=0.0)
+             kn = tl.load(KV_FP8 + ti[None, :] + o_nope[:, None], mask=nm[None, :] & mn[:, None], other=0.0)
              kn16 = (kn.to(tl.float32) * sc).to(tl.float16)
-  
-             kr = tl.load(KV_FP8 + ti[None, :] + o_rope_s[:, None],
-                          mask=nm[None, :] & mr[:, None], other=0.0)
+             kr = tl.load(KV_FP8 + ti[None, :] + o_rope_s[:, None], mask=nm[None, :] & mr[:, None], other=0.0)
              kr16 = (kr.to(tl.float32) * sc).to(tl.float16)
-  
              qk = tl.dot(qn, kn16) + tl.dot(qr, kr16)
              qk = qk.to(tl.float32) * sm_scale
              qk = tl.where(mask_h[:, None] & nm[None, :], qk, float("-inf"))
-  
-             vf = tl.load(KV_FP8 + ti[:, None] + o_dv[None, :],
-                          mask=nm[:, None] & mv[None, :], other=0.0)
+             vf = tl.load(KV_FP8 + ti[:, None] + o_dv[None, :], mask=nm[:, None] & mv[None, :], other=0.0)
              v16 = (vf.to(tl.float32) * sc).to(tl.float16)
-  
              ne = tl.maximum(tl.max(qk, 1), emax)
              rs = tl.exp(emax - ne)
              p = tl.exp(qk - ne[:, None])
              acc = acc * rs[:, None] + tl.dot(p.to(tl.float16), v16).to(tl.float32)
              esum = esum * rs + tl.sum(p, 1)
              emax = ne
-  
      ob = bid * stride_ab + heads[:, None] * stride_ah + sid * stride_as + o_dv[None, :]
      tl.store(Att_Out + ob, acc / tl.maximum(esum[:, None], 1e-12), mask=mask_h[:, None] & mv[None, :])
      lb = bid * stride_lb + heads * stride_lh + sid
      tl.store(Att_Lse + lb, emax + tl.log(tl.maximum(esum, 1e-12)), mask=mask_h)
   
-  
  @triton.jit
  def _flash_s2(
      Att_Out, Att_Lse, O,
  ⋯ 14 unchanged lines
          if l > -1e30:
              pv = tl.load(Att_Out + bid * stride_ab + hid * stride_ah + s * stride_as + od, mask=md, other=0.0)
              nm = tl.maximum(l, em)
-             os = tl.exp(em - nm)
-             ns = tl.exp(l - nm)
-             ac = ac * os + ns * pv
-             es = es * os + ns
+             o_s = tl.exp(em - nm)
+             n_s = tl.exp(l - nm)
+             ac = ac * o_s + n_s * pv
+             es = es * o_s + n_s
              em = nm
      tl.store(O + bid * stride_ob + hid * stride_oh + od, (ac / tl.maximum(es, 1e-12)).to(tl.bfloat16), mask=md)
   
+ _tbuf = {}
   
  def _triton_fp8(q, kv_data, kv_indptr, config, nsplits):
      bs = config["batch_size"]
  ⋯ 1 unchanged lines
      kv_flat = kv_fp8.view(-1, QK_DIM)
      q_r = q.view(bs, NUM_Q_HEADS, QK_DIM)
      k = (bs, nsplits)
-     if k not in _buf:
+     if k not in _tbuf:
          d = q.device
-         _buf[k] = (
+         _tbuf[k] = (
              torch.empty((bs, NUM_Q_HEADS, nsplits, V_DIM), dtype=torch.float32, device=d),
              torch.empty((bs, NUM_Q_HEADS, nsplits), dtype=torch.float32, device=d),
              torch.empty((bs, NUM_Q_HEADS, V_DIM), dtype=torch.bfloat16, device=d),
          )
-     ao, al, o = _buf[k]
-     BN = 32
+     ao, al, o = _tbuf[k]
      ex = {}
      try:
          if triton.runtime.driver.active.get_current_target().backend == "hip":
              ex = {"waves_per_eu": 1, "matrix_instr_nonkdim": 16, "kpack": 2}
      except Exception:
          pass
-     _flash_s1[(bs, nsplits)](
+     _flash_s1[(bs, 1, nsplits)](
          q_r, kv_flat, kv_scale, SM_SCALE, kv_indptr, ao, al,
          q_r.stride(0), q_r.stride(1), kv_flat.stride(0),
          ao.stride(0), ao.stride(1), ao.stride(2),
          al.stride(0), al.stride(1),
-         BLOCK_N=BN, BLOCK_H=16, NUM_SPLITS=nsplits,
+         BLOCK_N=32, BLOCK_H=16, NUM_SPLITS=nsplits,
          BLOCK_NOPE=512, BLOCK_ROPE=64, BLOCK_DV=512,
          Lnope=512, Lrope=64, Lv=512,
          num_warps=4, num_stages=1, **ex,
  ⋯ 9 unchanged lines
      return o
   
   
- # ═══════════════════════════════════════════════
- # AITER fp8 (for largest shapes)
- # ═══════════════════════════════════════════════
+ # ═══════════════ AITER ═══════════════
   
+ _meta_cache = {}
+ _idx_cache = {}
+ _out_cache = {}
+  
  def _quantize_fp8(tensor):
      finfo = torch.finfo(FP8_DTYPE)
      amax = tensor.abs().amax().clamp(min=1e-12)
      scale = amax / finfo.max
      return (tensor / scale).clamp(min=finfo.min, max=finfo.max).to(FP8_DTYPE), scale.float().reshape(1)
   
- def _aiter_fp8(q, kv_data, qo_indptr, kv_indptr, config, num_splits):
+ def _aiter(q, kv_data, qo_indptr, kv_indptr, config, num_splits, use_bf16_q=False):
      bs = config["batch_size"]
      q_len = config["q_seq_len"]
      total_kv = int(kv_indptr[-1].item())
      total_q = q.shape[0]
   
-     q_fp8, q_scale = _quantize_fp8(q)
+     if use_bf16_q:
+         q_input = q
+         q_scale = None
+         q_dtype = torch.bfloat16
+     else:
+         q_input, q_scale = _quantize_fp8(q)
+         q_dtype = q_input.dtype
+  
      kv_fp8, kv_scale = kv_data["fp8"]
      kv_4d = kv_fp8.view(kv_fp8.shape[0], PAGE_SIZE, NUM_KV_HEADS, kv_fp8.shape[-1])
   
-     key = (bs, num_splits, str(q_fp8.dtype), str(kv_fp8.dtype))
+     key = (bs, num_splits, str(q_dtype), str(kv_fp8.dtype))
      if key not in _meta_cache:
-         info = get_mla_metadata_info_v1(bs, q_len, NUM_Q_HEADS, q_fp8.dtype, kv_fp8.dtype,
+         info = get_mla_metadata_info_v1(bs, q_len, NUM_Q_HEADS, q_dtype, kv_fp8.dtype,
              is_sparse=False, fast_mode=False, num_kv_splits=num_splits, intra_batch_mode=True)
          _meta_cache[key] = [torch.empty(s, dtype=t, device="cuda") for s, t in info]
      work = _meta_cache[key]
  ⋯ 4 unchanged lines
          page_size=PAGE_SIZE, kv_granularity=max(PAGE_SIZE, 16),
          max_seqlen_qo=q_len, uni_seqlen_qo=q_len,
          fast_mode=False, max_split_per_batch=num_splits,
-         intra_batch_mode=True, dtype_q=q_fp8.dtype, dtype_kv=kv_fp8.dtype)
+         intra_batch_mode=True, dtype_q=q_dtype, dtype_kv=kv_fp8.dtype)
   
      if total_kv not in _idx_cache:
          _idx_cache[total_kv] = torch.arange(total_kv, dtype=torch.int32, device="cuda")
  ⋯ 1 unchanged lines
          _out_cache[total_q] = torch.empty((total_q, NUM_Q_HEADS, V_DIM), dtype=torch.bfloat16, device="cuda")
   
      o = _out_cache[total_q]
-     mla_decode_fwd(q_fp8.view(-1, NUM_Q_HEADS, QK_DIM), kv_4d, o,
+     mla_decode_fwd(q_input.view(-1, NUM_Q_HEADS, QK_DIM), kv_4d, o,
          qo_indptr, kv_indptr, _idx_cache[total_kv], kv_last, q_len,
          page_size=PAGE_SIZE, nhead_kv=NUM_KV_HEADS, sm_scale=SM_SCALE, logit_cap=0.0,
          num_kv_splits=num_splits, q_scale=q_scale, kv_scale=kv_scale,
  ⋯ 3 unchanged lines
      return o
   
   
- # ═══════════════════════════════════════════════
- # BMM bf16 (for small shapes)
- # ═══════════════════════════════════════════════
+ # ═══════════════ BMM ═══════════════
   
  def _bmm(q, kv_data, config):
      bs = config["batch_size"]
  ⋯ 6 unchanged lines
      return torch.bmm(w, V)
   
   
- # ═══════════════════════════════════════════════
- # Shape dispatch based on measured data
- # ═══════════════════════════════════════════════
+ # ═══════════════ Dispatch ═══════════════
   
- # Per-shape best path (from benchmarks):
- # bmm: bs=4/1k(26), bs=4/8k(43), bs=32/1k(42), bs=64/1k(61)
- # Triton fp8: bs=32/8k(120) [AITER was 166]
- # AITER fp8: bs=64/8k(206), bs=256/1k(161), bs=256/8k(353)
+ _warm = set()
   
- _AITER_SPLITS = {
-     (64, 8192): 32,
-     (256, 1024): 16,
-     (256, 8192): 32,
- }
-  
  def custom_kernel(data: input_t) -> output_t:
      q, kv_data, qo_indptr, kv_indptr, config = data
      bs = config["batch_size"]
  ⋯ 2 unchanged lines
      if sk not in _warm:
          _warm.add(sk)
   
-     # bmm: small shapes
+     # bmm bf16: small shapes
      if bs <= 4 or (bs <= 64 and kv_len <= 1024):
          return _bmm(q, kv_data, config)
   
-     # Triton fp8: bs=32/kv=8192 (beats AITER)
-     if bs <= 32 and kv_len <= 8192:
+     # Triton fp8: bs=32/kv=8192 (121μs, beats AITER 166μs)
+     if bs == 32 and kv_len == 8192:
          return _triton_fp8(q, kv_data, kv_indptr, config, 8)
   
-     # AITER fp8: large shapes
-     return _aiter_fp8(q, kv_data, qo_indptr, kv_indptr, config, _AITER_SPLITS.get(sk, 32))
+     # AITER bf16Q+fp8KV: bs=256/kv=1024 (136μs, -24% vs fp8+fp8 179μs)
+     if bs == 256 and kv_len == 1024:
+         return _aiter(q, kv_data, qo_indptr, kv_indptr, config, 16, use_bf16_q=True)
+  
+     # AITER fp8+fp8: remaining large shapes
+     splits = {(64, 8192): 32, (256, 8192): 32}.get(sk, 32)
+     return _aiter(q, kv_data, qo_indptr, kv_indptr, config, splits, use_bf16_q=False)