submission 607363

  """
- Attempt 265: Safe per-case dispatch v2 — NP splits=1 ps=1 for bs=64/kv=1024.
+ Attempt 305: Try is_sparse=True for metadata scheduling.
   
- 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).
+ Never explicitly tested is_sparse mode. Might use different work scheduling
+ that better handles our 16:1 MQA + non-standard 576 head dim.
   
- 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     |
+ Uses same proven config as 094/272: ps=2 for bs>=64/kv=1024, ps=8 for kv>1024,
+ 12 splits for kv=8192, 16 splits for kv=1024.
  """
   
  import torch
  ⋯ 14 unchanged lines
  _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),
+     (4, 1024):   (1, 16),
+     (4, 8192):   (8, 12),
+     (32, 1024):  (1, 16),
+     (32, 8192):  (8, 12),
+     (64, 1024):  (2, 16),
+     (64, 8192):  (8, 12),
+     (256, 1024): (2, 16),
+     (256, 8192): (8, 12),
  }
   
   
  ⋯ 3 unchanged lines
          return _CONFIG[key]
      if kv_seq_len <= 1024:
          if batch_size >= 64:
-             return 1, 1, False
-         else:
-             return 1, 12, False
-     else:
-         if batch_size >= 64:
-             return 8, 12, True
-         else:
-             return 8, 12, False
+             return 2, 16
+         return 1, 16
+     return 8, 12
   
   
- def _build_page_info(batch_size, kv_seq_len, total_kv, kv_indptr, page_size):
+ def _build_cache(batch_size, kv_seq_len, total_q, total_kv, qo_indptr, kv_indptr):
+     page_size, num_kv_splits = _get_config(batch_size, kv_seq_len)
+     fast_mode = kv_seq_len > 1024
+  
      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")
-     return kv_indices, kv_last_page_len, kv_indptr_pages
+     o = torch.empty((total_q, NUM_HEADS, V_HEAD_DIM), dtype=torch.bfloat16, device="cuda")
   
-  
- 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, 1, NUM_HEADS, q_dtype, kv_dtype,
-         is_sparse=False, fast_mode=fast_mode,
+         batch_size, 1, NUM_HEADS, FP8_DTYPE, FP8_DTYPE,
+         is_sparse=True, fast_mode=fast_mode,
          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
+     (wm, wi, wis, ri, rfm, rpm) = 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,
+         wm, wis, wi, ri, rfm, rpm,
          page_size=page_size, kv_granularity=max(page_size, 16),
          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,
+         dtype_q=FP8_DTYPE, dtype_kv=FP8_DTYPE,
      )
      return {
-         "work_meta_data": work_metadata,
-         "work_indptr": work_indptr,
-         "work_info_set": work_info_set,
-         "reduce_indptr": reduce_indptr,
-         "reduce_final_map": reduce_final_map,
-         "reduce_partial_map": reduce_partial_map,
+         "work_meta_data": wm, "work_indptr": wi, "work_info_set": wis,
+         "reduce_indptr": ri, "reduce_final_map": rfm, "reduce_partial_map": rpm,
+         "kv_indices": kv_indices, "kv_last_page_len": kv_last_page_len,
+         "kv_indptr_pages": kv_indptr_pages, "o": o, "total_kv": total_kv,
+         "page_size": page_size, "num_kv_splits": num_kv_splits,
      }
   
   
- 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,
-         "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
   
  ⋯ 7 unchanged lines
      cache_key = (batch_size, kv_seq_len)
      if cache_key not in _cache or _cache[cache_key]["total_kv"] != total_kv:
          _cache[cache_key] = _build_cache(
-             batch_size, kv_seq_len, total_q, total_kv,
-             qo_indptr, kv_indptr,
+             batch_size, kv_seq_len, total_q, total_kv, qo_indptr, kv_indptr,
          )
      c = _cache[cache_key]
   
  ⋯ 1 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)
   
-     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,
-         )
+     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"],
+     )
   
      return c["o"]