submission 606366

  """
- Attempt 130: Direct stage1 ASM + reduce_v1 calls with pre-allocated buffers.
+ Attempt 265: Safe per-case dispatch v2 — NP splits=1 ps=1 for bs=64/kv=1024.
   
- Key insight from reading mla_decode_fwd source:
- - Persistent mode allocates `logits` (fp32) and `attn_lse` (fp32) EVERY CALL
- - These allocations are sized by reduce_partial_map.size(0) * max_seqlen_q
- - By calling mla_decode_stage1_asm_fwd + mla_reduce_v1 directly with
-   pre-allocated buffers, we eliminate per-call allocation overhead.
+ 264 showed NP ps=1 splits=16 is terrible for bs=64/kv=1024 (39μs).
+ Try NP ps=1 splits=1 instead — skip stage2 entirely, just like ps=2 splits=1
+ but with safe accuracy (ps=1).
   
- Based on attempt_094 (current best: 33.5μs ranked).
+ Grid for bs=64/kv=1024 with splits=1: (1, 64, 1) = 64 blocks.
+ Each block processes all 1024 tokens → compute-bound, no reduce.
+  
+ Dispatch:
+ | Case              | Mode       | PS | Splits |
+ |-------------------|------------|----|--------|
+ | bs=4, kv=1024     | NP         | 1  | 16     |
+ | bs=4, kv=8192     | NP         | 8  | 16     |
+ | bs=32, kv=1024    | NP         | 1  | 12     |
+ | bs=32, kv=8192    | NP         | 8  | 12     |
+ | bs=64, kv=1024    | NP         | 1  | 1      | ← splits=1 skip reduce
+ | bs=64, kv=8192    | Persistent | 8  | 12     |
+ | bs=256, kv=1024   | NP         | 2  | 1      |
+ | bs=256, kv=8192   | Persistent | 8  | 12     |
  """
   
  import torch
  from task import input_t, output_t
   
- import aiter
  from aiter.mla import mla_decode_fwd
  from aiter import dtypes as aiter_dtypes
  from aiter import get_mla_metadata_info_v1, get_mla_metadata_v1
  ⋯ 4 unchanged lines
  V_HEAD_DIM = 512
  SM_SCALE = 1.0 / (QK_HEAD_DIM ** 0.5)
  FP8_DTYPE = aiter_dtypes.fp8
- NUM_KV_SPLITS = 16
  Q_SCALE = torch.ones(1, dtype=torch.float32, device="cuda")
   
  _cache = {}
   
+ _CONFIG = {
+     (4, 1024):   (1, 16, False),
+     (4, 8192):   (8, 16, False),
+     (32, 1024):  (1, 12, False),
+     (32, 8192):  (8, 12, False),
+     (64, 1024):  (1, 1, False),    # NP splits=1 ps=1 (safe)
+     (64, 8192):  (8, 12, True),
+     (256, 1024): (2, 1, False),
+     (256, 8192): (8, 12, True),
+ }
   
- def _get_page_size(batch_size, kv_seq_len):
-     if kv_seq_len > 1024:
-         return 8
-     elif batch_size >= 64:
-         return 2
+  
+ def _get_config(batch_size, kv_seq_len):
+     key = (batch_size, kv_seq_len)
+     if key in _CONFIG:
+         return _CONFIG[key]
+     if kv_seq_len <= 1024:
+         if batch_size >= 64:
+             return 1, 1, False
+         else:
+             return 1, 12, False
      else:
-         return 1
+         if batch_size >= 64:
+             return 8, 12, True
+         else:
+             return 8, 12, False
   
   
- def _build_cache(batch_size, kv_seq_len, total_q, total_kv, qo_indptr, kv_indptr):
-     max_q_len = 1
-     q_dtype = FP8_DTYPE
-     kv_dtype = FP8_DTYPE
-     page_size = _get_page_size(batch_size, kv_seq_len)
-     fast_mode = kv_seq_len > 1024
-  
+ def _build_page_info(batch_size, kv_seq_len, total_kv, kv_indptr, page_size):
      seq_lens = kv_indptr[1:] - kv_indptr[:-1]
      if page_size == 1:
          kv_last_page_len = seq_lens.to(torch.int32)
  ⋯ 5 unchanged lines
          kv_indptr_pages = torch.zeros(batch_size + 1, dtype=torch.int32, device="cuda")
          kv_indptr_pages[1:] = torch.cumsum(pages_per_seq, dim=0).to(torch.int32)
          num_pages = int(kv_indptr_pages[-1].item())
-  
      kv_indices = torch.arange(num_pages, dtype=torch.int32, device="cuda")
-     o = torch.empty((total_q, NUM_HEADS, V_HEAD_DIM), dtype=torch.bfloat16, device="cuda")
+     return kv_indices, kv_last_page_len, kv_indptr_pages
   
+  
+ def _build_persistent_metadata(batch_size, qo_indptr, kv_indptr_pages,
+                                 kv_last_page_len, page_size, fast_mode,
+                                 num_kv_splits):
+     q_dtype = FP8_DTYPE
+     kv_dtype = FP8_DTYPE
      info = get_mla_metadata_info_v1(
-         batch_size, max_q_len, NUM_HEADS, q_dtype, kv_dtype,
+         batch_size, 1, NUM_HEADS, q_dtype, kv_dtype,
          is_sparse=False, fast_mode=fast_mode,
-         num_kv_splits=NUM_KV_SPLITS, intra_batch_mode=True,
+         num_kv_splits=num_kv_splits, intra_batch_mode=True,
      )
      work = [torch.empty(s, dtype=t, device="cuda") for s, t in info]
      (work_metadata, work_indptr, work_info_set,
       reduce_indptr, reduce_final_map, reduce_partial_map) = work
-  
      get_mla_metadata_v1(
          qo_indptr, kv_indptr_pages, kv_last_page_len,
          NUM_HEADS // NUM_KV_HEADS, NUM_KV_HEADS, True,
          work_metadata, work_info_set, work_indptr,
          reduce_indptr, reduce_final_map, reduce_partial_map,
          page_size=page_size, kv_granularity=max(page_size, 16),
-         max_seqlen_qo=max_q_len, uni_seqlen_qo=max_q_len,
-         fast_mode=fast_mode, max_split_per_batch=NUM_KV_SPLITS,
-         intra_batch_mode=True, dtype_q=q_dtype, dtype_kv=kv_dtype,
+         max_seqlen_qo=1, uni_seqlen_qo=1, fast_mode=fast_mode,
+         max_split_per_batch=num_kv_splits, intra_batch_mode=True,
+         dtype_q=q_dtype, dtype_kv=kv_dtype,
      )
-  
-     # Pre-allocate intermediate buffers (this is what mla_decode_fwd allocates every call)
-     num_partials = reduce_partial_map.size(0)
-     logits = torch.empty(
-         (num_partials * max_q_len, 1, NUM_HEADS, V_HEAD_DIM),
-         dtype=aiter_dtypes.fp32, device="cuda",
-     )
-     attn_lse = torch.empty(
-         (num_partials * max_q_len, 1, NUM_HEADS, 1),
-         dtype=aiter_dtypes.fp32, device="cuda",
-     )
-  
      return {
          "work_meta_data": work_metadata,
          "work_indptr": work_indptr,
  ⋯ 1 unchanged lines
          "reduce_indptr": reduce_indptr,
          "reduce_final_map": reduce_final_map,
          "reduce_partial_map": reduce_partial_map,
+     }
+  
+  
+ def _build_cache(batch_size, kv_seq_len, total_q, total_kv,
+                  qo_indptr, kv_indptr):
+     page_size, num_kv_splits, use_persistent = _get_config(batch_size, kv_seq_len)
+     fast_mode = kv_seq_len > 1024
+  
+     kv_indices, kv_last_page_len, kv_indptr_pages = \
+         _build_page_info(batch_size, kv_seq_len, total_kv, kv_indptr, page_size)
+  
+     o = torch.empty((total_q, NUM_HEADS, V_HEAD_DIM), dtype=torch.bfloat16, device="cuda")
+  
+     result = {
          "kv_indices": kv_indices,
          "kv_last_page_len": kv_last_page_len,
          "kv_indptr_pages": kv_indptr_pages,
          "o": o,
-         "logits": logits,
-         "attn_lse": attn_lse,
          "total_kv": total_kv,
          "page_size": page_size,
+         "num_kv_splits": num_kv_splits,
+         "use_persistent": use_persistent,
      }
   
+     if use_persistent:
+         meta = _build_persistent_metadata(
+             batch_size, qo_indptr, kv_indptr_pages,
+             kv_last_page_len, page_size, fast_mode, num_kv_splits,
+         )
+         result.update(meta)
   
+     return result
+  
+  
  def custom_kernel(data: input_t) -> output_t:
      q, kv_data, qo_indptr, kv_indptr, config = data
   
  ⋯ 16 unchanged lines
      q_fp8_view = q.to(FP8_DTYPE).view(-1, NUM_HEADS, QK_HEAD_DIM)
      kv_4d = kv_buffer_fp8.view(total_kv // ps, ps, NUM_KV_HEADS, QK_HEAD_DIM)
   
-     # Direct stage1 ASM call (bypasses mla_decode_fwd Python overhead)
-     aiter.mla_decode_stage1_asm_fwd(
-         q_fp8_view,
-         kv_4d,
-         qo_indptr,
-         c["kv_indptr_pages"],
-         c["kv_indices"],
-         c["kv_last_page_len"],
-         None,  # num_kv_splits_indptr (None for persistent mode)
-         c["work_meta_data"],
-         c["work_indptr"],
-         c["work_info_set"],
-         1,  # max_seqlen_q
-         ps,  # page_size
-         NUM_KV_HEADS,
-         SM_SCALE,
-         c["logits"],  # pre-allocated
-         c["attn_lse"],  # pre-allocated
-         c["o"],
-         Q_SCALE,
-         kv_scale,
-     )
+     if c["use_persistent"]:
+         mla_decode_fwd(
+             q_fp8_view, kv_4d, c["o"],
+             qo_indptr, c["kv_indptr_pages"], c["kv_indices"],
+             c["kv_last_page_len"], 1,
+             page_size=ps, nhead_kv=NUM_KV_HEADS,
+             sm_scale=SM_SCALE, logit_cap=0.0,
+             num_kv_splits=c["num_kv_splits"],
+             q_scale=Q_SCALE, kv_scale=kv_scale,
+             intra_batch_mode=True,
+             work_meta_data=c["work_meta_data"],
+             work_indptr=c["work_indptr"],
+             work_info_set=c["work_info_set"],
+             reduce_indptr=c["reduce_indptr"],
+             reduce_final_map=c["reduce_final_map"],
+             reduce_partial_map=c["reduce_partial_map"],
+         )
+     else:
+         mla_decode_fwd(
+             q_fp8_view, kv_4d, c["o"],
+             qo_indptr, c["kv_indptr_pages"], c["kv_indices"],
+             c["kv_last_page_len"], 1,
+             page_size=ps, nhead_kv=NUM_KV_HEADS,
+             sm_scale=SM_SCALE, logit_cap=0.0,
+             num_kv_splits=c["num_kv_splits"],
+             q_scale=Q_SCALE, kv_scale=kv_scale,
+         )
   
-     # Direct reduce call
-     aiter.mla_reduce_v1(
-         c["logits"],
-         c["attn_lse"],
-         c["reduce_indptr"],
-         c["reduce_final_map"],
-         c["reduce_partial_map"],
-         1,  # max_seqlen_q
-         c["o"],
-         None,  # final_lse (not needed)
-     )
-  
      return c["o"]