[ASM] [HIP] [MLA v4 nm] Split plan helper and persistent decode kernel with in-kernel split and merge

aiter/mla.py

@@ -5,6 +5,7 @@
 
 import functools
 import os
+from typing import NamedTuple
 
 import torch
 import triton
@@ -359,6 +360,123 @@ def get_meta_param(
     return num_kv_splits, num_kv_splits_indptr
 
 
+# Per-arch v4 nm decode latency model used to pick `num_kv_splits`:
+#
+#   t(i) ~= L + W(i) * (C + A * kv_len / i) + [i > 1] * (S + D * num_seqs * i)
+#   W(i)  = ceil(num_seqs * tg_factor * i / cu_num)   (workgroup waves)
+#
+# A is the per-token KV cost of one workgroup, C the per-wave fixed cost, and
+# S + D * num_seqs * i the FP32 partial write + stage-2 merge that every split
+# launch pays. The occupancy-only pick (ignore_total_kv=1) instead rewards
+# filling whole multiples of cu_num, e.g. 13-16 splits for 49-119 seqs whose KV
+# fits one wave.
+#
+# gfx950: least-squares fit (relative error) on MI355X, gqa=128, qSeqLen=1,
+# CUDA-graph replay, 1..16 splits over uniform-KV shapes (7..448 seqs, kv_len
+# 192..8320). Archs without an entry keep the occupancy-only pick.
+_V4_NM_SPLIT_COST = {
+    "gfx950": (12.74, 3.21, 0.0300, 3.38, 0.076),  # L, C, A, S, D (us)
+}
+_V4_NM_MAX_SPLITS = 16
+
+
+@functools.lru_cache(maxsize=1024)
+def _v4_nm_pick_num_kv_splits(cost, num_seqs, tg_factor, kv_len, cu_num):
+    """Split count minimizing the latency model above with coefficients `cost`."""
+    L, C, A, S, D = cost
+    kv_len = max(1, kv_len)
+    best_cost, best_splits = None, 1
+    for i in range(1, _V4_NM_MAX_SPLITS + 1):
+        waves = -(-num_seqs * tg_factor * i // cu_num)
+        t = L + waves * (C + A * kv_len / i)
+        if i > 1:
+            t += S + D * num_seqs * i
+        if best_cost is None or t < best_cost:
+            best_cost, best_splits = t, i
+    return best_splits
+
+
+def get_mla_v4_nm_num_kv_splits(num_seqs, num_heads, kv_len) -> int:
+    """`num_kv_splits` for `mla_decode_fwd_v4_nm` (host only, no device work).
+
+    `num_seqs` / `num_heads` are the decode call's `qo_indptr.shape[0] - 1`
+    and `q.size(1)`; `kv_len` is the per-seq KV length to plan for (the
+    longest one the call can see). Minimizes the latency model above where
+    the arch has coefficients, and is the wrapper's own occupancy-only pick
+    elsewhere. The matching `split_indptr` is the uniform
+    `[0, s, 2s, ..., num_seqs * s]`; any prefix of a longer uniform buffer is
+    one, so callers can keep a constant buffer per split count and slice it
+    instead of writing one per call.
+    """
+    tg_factor = max(1, -(-num_heads // 64))  # ceil(num_heads / 64)
+    cost = _V4_NM_SPLIT_COST.get(get_gfx())
+    if cost is not None:
+        num_kv_splits = _v4_nm_pick_num_kv_splits(
+            cost, num_seqs, tg_factor, int(kv_len), get_cu_num()
+        )
+    else:
+        num_kv_splits, _ = get_meta_param(
+            None,
+            num_seqs,
+            num_seqs * max(1, int(kv_len)),
+            num_heads,
+            1,
+            dtypes.fp8,
+            tg_factor,
+            1,  # ignore_total_kv, as in mla_decode_fwd_v4_nm
+        )
+    return num_kv_splits
+
+
+# Workspace of the persistent v4 nm decode kernel (mla_decode_v4_ps_asm).
+# cnt: [0, 2*65536) row counters, [2*65536, +16*512) reserved,
+# [.., +4*1024) group counters.
+_V4_NM_PS_CNT_INTS = 2 * 65536 + 16 * 512 + 4 * 1024
+_V4_NM_PS_ARANGE = 65537
+_V4_NM_PS_MAX_PARTITIONS = 1024
+
+
+class MlaV4NmPsWorkspace(NamedTuple):
+    """Scratch of `mla_decode_fwd_v4_nm_ps`; build it with
+    `get_mla_v4_nm_ps_workspace`. P = `desc.size(0)` partitions."""
+
+    o_acc: torch.Tensor  # [2P, 128, 512] fp32 split partials
+    lse_acc: torch.Tensor  # [2P, 128] fp32 split partials
+    desc: torch.Tensor  # [P, 8] int32 in-kernel plan
+    cnt: torch.Tensor  # int32 merge counters, zero at rest
+    arange: torch.Tensor  # int32 arange, read-only
+
+
+def get_mla_v4_nm_ps_workspace(device="cuda", num_partitions=128) -> MlaV4NmPsWorkspace:
+    """Allocate a workspace for `mla_decode_fwd_v4_nm_ps`.
+
+    One workspace may serve any number of calls that are ordered on the GPU
+    (same stream, or one CUDA graph replayed at a time); calls that can run
+    concurrently need one workspace each. Nothing carries over between calls,
+    so one workspace can be reused across batch sizes and graphs. Allocate it
+    outside CUDA-graph capture.
+    """
+    if not 1 <= num_partitions <= _V4_NM_PS_MAX_PARTITIONS:
+        raise ValueError(
+            f"num_partitions must be in [1, {_V4_NM_PS_MAX_PARTITIONS}], "
+            f"got {num_partitions}"
+        )
+    if torch.cuda.is_current_stream_capturing():
+        raise RuntimeError(
+            "get_mla_v4_nm_ps_workspace: allocate the workspace before "
+            "CUDA-graph capture"
+        )
+    P = num_partitions
+    i32 = {"dtype": dtypes.i32, "device": device}
+    return MlaV4NmPsWorkspace(
+        o_acc=torch.empty(2 * P, 128, 512, dtype=dtypes.fp32, device=device),
+        lse_acc=torch.empty(2 * P, 128, dtype=dtypes.fp32, device=device),
+        desc=torch.empty(P, 8, **i32),
+        cnt=torch.zeros(_V4_NM_PS_CNT_INTS, **i32),
+        arange=torch.arange(_V4_NM_PS_ARANGE, **i32),
+    )
+
+
 # Persistent MLA-decode kernel gate: the persistent kernel
 # ("mla_a16w16_qh16..._ps") is slower than the non-persistent split-KV kernel
 # ("mla_dec_stage1...") above a concurrency threshold (~batch 16-64 on gfx950 bf16
@@ -1738,6 +1856,10 @@ def mla_decode_fwd_v4_nm(
       for the kernel tile) -- identical to `mla_decode_fwd`'s non-persistent
       path. Pass an explicit int to override. Note V4 nm is always
       non-persistent, so only that branch of `get_meta_param` applies.
+      Callers that know their per-seq KV length (and CUDA-graph callers,
+      whose `kv_page_indices` is typically capacity-sized) should pick it
+      with `get_mla_v4_nm_num_kv_splits` and pass it with the uniform
+      `split_indptr` `[0, s, 2s, ...]`.
 
     Multi-pass mode (`num_kv_splits > 1`):
       1. If `split_indptr` is None, build a uniform one:
@@ -1956,7 +2078,12 @@ def mla_decode_fwd_v4_nm(
             0,  # stride_lse_bs (unused, HAS_FINAL_LSE=False)
             page_size=1,  # v4 nm KV cache is page_size=1
             KV_INDPTR_IS_PAGE_LEVEL=False,  # page_size=1 -> token-level indptr
-            MAYBE_FINAL_OUT=True,
+            # The kernel only writes the packed-BF16 result in place when
+            # num_kv_splits == 1, which never reaches this merge. A
+            # caller-supplied split_indptr with one split per seq on a wider
+            # grid still gets FP32 partials, so the in-place copy (keyed on
+            # split_indptr[-1] == num_seqs) must stay off.
+            MAYBE_FINAL_OUT=False,
             HAS_FINAL_LSE=False,
             USE_VALID_SPLIT_COUNT_REDUCE=int(num_kv_splits > 1),
             BATCH_NUM=num_seqs,
@@ -1972,3 +2099,58 @@ def mla_decode_fwd_v4_nm(
         )
 
     return logits, attn_lse
+
+
+def mla_decode_fwd_v4_nm_ps(
+    q_packed,  # [N, 128, 512] FP8 packed Q+e8m0
+    q_rope,  # [N, 128, 64] BF16
+    kv_packed,  # [rows, ..., 512] FP8 packed KV pool, page_size 1
+    kv_rope,  # [rows, ..., 64] BF16
+    kv_indptr,  # [>= N+1] int32
+    kv_page_indices,  # [*] int32
+    sink,  # [128] FP32 attention sink logit
+    workspace: MlaV4NmPsWorkspace,
+    out=None,  # [N, 128, 512] BF16
+    return_lse=False,
+    lse=None,  # [N, 128] FP32, used when return_lse
+):
+    """v4 nm decode (128 heads, one query token per row) with the persistent
+    kernel: one launch that plans the KV split over the workspace's
+    partitions from `kv_indptr`, runs the attention and merges the split
+    partials, so no split plan or stage-2 merge is needed.
+
+    Same math and packed layouts as `mla_decode_fwd_v4_nm` with gqa=128 and
+    max_seqlen_q=1. N = `q_packed.size(0)` (N <= 32768); row j attends to pool
+    rows `kv_page_indices[kv_indptr[j] : kv_indptr[j + 1]]`, so `kv_indptr`
+    may hold more than N+1 entries and need not start at 0. Rows with an
+    empty KV range are left unwritten in `out` (and in the LSE).
+
+    `workspace` comes from `get_mla_v4_nm_ps_workspace`; see its sharing
+    rule. Returns `out`, or `(out, lse)` with the natural-log LSE [N, 128]
+    FP32 (sink included) when `return_lse`; `out` / `lse` are allocated
+    when not given.
+    """
+    require_gfx1250_asm("mla_decode_v4_ps_asm")
+    if out is None:
+        out = torch.empty(
+            (q_packed.size(0), 128, 512), dtype=dtypes.bf16, device=q_packed.device
+        )
+    if not return_lse:
+        lse = None
+    elif lse is None:
+        lse = torch.empty(
+            (q_packed.size(0), 128), dtype=dtypes.fp32, device=q_packed.device
+        )
+    aiter.mla_decode_v4_ps_asm(
+        q_packed,
+        q_rope,
+        kv_packed,
+        kv_rope,
+        kv_indptr,
+        kv_page_indices,
+        sink,
+        *workspace,
+        out,
+        lse,
+    )
+    return (out, lse) if return_lse else out

aiter/ops/attention.py

@@ -915,6 +915,41 @@ def mla_decode_v4_asm(
 ) -> None: ...
 
 
+@compile_ops(MD_NAME_V4, ffi_type="ctypes")
+def mla_decode_v4_ps_asm(
+    # [N, 128, 512] FP8 packed Q + e8m0 scale region
+    Q: torch.Tensor,
+    # [N, 128, 64] BF16
+    qrope: torch.Tensor,
+    # [rows, ..., 512] FP8 packed KV pool (row-dense, page_size 1)
+    KV: torch.Tensor,
+    # [rows, ..., 64] BF16
+    kvrope: torch.Tensor,
+    # [>= N+1] int32
+    kv_indptr: torch.Tensor,
+    # [*] int32
+    kv_page_indices: torch.Tensor,
+    # [128] FP32 attention sink logit
+    sink: torch.Tensor,
+    # workspace (aiter.mla.get_mla_v4_nm_ps_workspace); P = desc.size(0)
+    # [2P, 128, 512] FP32
+    o_acc: torch.Tensor,
+    # [2P, 128] FP32
+    lse_acc: torch.Tensor,
+    # [P, 8] int32
+    desc: torch.Tensor,
+    # int32 counters, zero at rest
+    cnt: torch.Tensor,
+    # int32 arange
+    arange: torch.Tensor,
+    # outputs
+    # [N, 128, 512] BF16
+    output: torch.Tensor,
+    # [N, 128] FP32 natural-log LSE (sink included); None = not written
+    lse: torch.Tensor | None = None,
+) -> None: ...
+
+
 @compile_ops(MD_NAME, ffi_type="ctypes")
 def mla_prefill_asm_fwd(
     # [num_seqs, num_heads, head_size]