Add forward compatibility to get_mla_metadata_v1

aiter/ops/attention.py

@@ -1271,8 +1271,8 @@ def get_mla_metadata_info_v1(
         )
 
 
-@compile_ops("module_mla_metadata")
-def get_mla_metadata_v1(
+@compile_ops("module_mla_metadata", fc_name="get_mla_metadata_v1")
+def _get_mla_metadata_v1_impl(
     seqlens_qo_indptr: torch.Tensor,
     seqlens_kv_indptr: torch.Tensor,
     kv_last_page_lens: torch.Tensor,
@@ -1300,52 +1300,138 @@ def get_mla_metadata_v1(
     dtype_kv_nope: Optional[torch.dtype] = None,
     dtype_kv_rope: Optional[torch.dtype] = None,
 ) -> None:
-    """
-    Inputs:
-        cumulated seqlens of q/o: (batch_size + 1), dtype torch.int32.
-        cumulated page indices of k/v: (batch_size + 1), dtype torch.int32.
-        Length of last page of k/v: (batch_size), dtype torch.int32.
-        num_heads_per_head_k: Equals to num_heads_q // num_heads_k.
-        num_heads_k: num_heads_k.
-        is_causal: Whether causal mask is enabled.
-        Options: Detailed settings for spliting. All of them are optional.
-            page_size: default=1. The size of a page.
-            kv_granularity: default=16. The granularity on kv page nums when cutting batch.
-            max_seqlen_qo: default=-1. Used to check lds usage and save time. value less than 1 means unknown.
-            uni_seqlen_qo: default=-1. Sequence length of qo is uniform across batches. value less than 1 means the
-                           length is not fixed.
-            fast_mode: default=True. Whether user wants metadata become as fast as possible. Note that fast
-                       mode may lead to bad overall performance.
-            intra_batch_mode: default=False. Fake non persistent mode. Same splits for each batch.
-            topk: default=-1. Top-k tokens selected for sparse attention. -1 means non-sparse attention.
-    Outputs:
-        [0] work_metadata_ptrs  (2)                 Two 64-bits pointers point to the 1st element of work_indptr and
-                                                    work_info.
-        [1] work_indptr:        (#cu_part + 1),     The IDs of work handled by each cu_part.
-        [2] work_info           (#work, 8)
-        [2.0] bs_index:         (#work),            The index of batch handled by each work.
-        [2.1] partial_index:    (#work),            The index of tile in output buffer when splits. -1 means no split.
-        [2.2] q_start:          (#work),            The global index in seq where q/o starts. Use global index here can
-                                                    reduce memory access count in kernel.
-        [2.3] q_end:            (#work),            The global index in seq where q/o ends (not included).
-        [2.4] kv_start:         (#work),            The global index in page where k/v starts.
-        [2.5] kv_end:           (#work),            The global index in page where k/v ends (not included). Note that
-                                                    this value indicates the end of last qo sequence if there are
-                                                    multiple qo sequences included in the current work and causal mask
-                                                    is enabled when page_size is 1.
-        [2.6] kv_offset:        (#work),            Remaining length in seq from kv_end to the end of current batch.
-        [2.7] pad               (#work, 1),         Pad to 8 DWs.
-        [3] reduce_indptr:      (sum(qo_seqlen_blk_count) + 1),
-                                                    The IDs in reduce_partial_map indicates the tiles should be merged
-                                                    together.
-        [4] reduce_final_map:   (sum(qo_seqlen_blk_count)),
-                                                    The final output location of each group of tiles.
-        [5] reduce_partial_map: (#partial_tiles),   The locations in partial buffer of partial tiles waiting for being
-                                                    reduced.
+    """Compiled binding for ``get_mla_metadata_v1`` (bound via ``fc_name``).
+
+    Internal: call the ``get_mla_metadata_v1`` wrapper below instead, which also
+    accepts the legacy ``dtype_q`` / ``dtype_kv`` aliases. This entry takes the
+    per-component dtypes directly.
+
+    Args:
+        seqlens_qo_indptr: (batch_size + 1) int32. Cumulative q/o seqlens.
+        seqlens_kv_indptr: (batch_size + 1) int32. Cumulative k/v page indices.
+        kv_last_page_lens: (batch_size) int32. Length of the last k/v page.
+        num_heads_per_head_k: num_heads_q // num_heads_k.
+        num_heads_k: Number of k heads.
+        is_causal: Whether the causal mask is enabled.
+        work_metadata_ptrs / work_info_set / work_indptr / reduce_indptr /
+            reduce_final_map / reduce_partial_map: Pre-allocated output buffers
+            (see get_mla_metadata_info_v1 for shapes/dtypes) filled in place.
+        page_size: Page size (default 1).
+        kv_granularity: Granularity of kv page counts when splitting a batch (default 16).
+        max_seqlen_qo: Max qo seqlen; used to bound LDS usage. < 1 means unknown (default -1).
+        uni_seqlen_qo: Uniform qo seqlen across batches. < 1 means not fixed (default -1).
+        fast_mode: Prefer fastest metadata generation; may hurt overall perf (default True).
+        topk: Top-k tokens for sparse attention. -1 means non-sparse (default -1).
+        max_split_per_batch: Cap on KV splits per batch. -1 means unbounded (default -1).
+        intra_batch_mode: Fake non-persistent mode; same split count per batch (default False).
+        is_cp_round_robin: Context-parallel round-robin work assignment (default False).
+        mla_version: MLA kernel version (default V32).
+        dtype_q_nope / dtype_kv_nope: Q/KV NoPE dtypes. None -> bf16 (C++ default).
+        dtype_q_rope / dtype_kv_rope: Q/KV RoPE dtypes. None -> the matching NoPE dtype.
+
+    Outputs (written into the buffers above):
+        [0] work_metadata_ptrs  (2)             Two 64-bit pointers to the 1st elements of
+                                                work_indptr and work_info.
+        [1] work_indptr         (#cu_part + 1)  IDs of the work handled by each cu_part.
+        [2] work_info           (#work, 8)      Per-work descriptor, 8 DWs each:
+              [.0] bs_index        Batch index handled by this work.
+              [.1] partial_index   Output-buffer tile index when split; -1 means no split.
+              [.2] q_start         Global seq index where q/o starts.
+              [.3] q_end           Global seq index where q/o ends (exclusive).
+              [.4] kv_start        Global page index where k/v starts.
+              [.5] kv_end          Global page index where k/v ends (exclusive). With causal
+                                   mask and page_size==1, this is the end of the last qo seq
+                                   when the work covers multiple qo sequences.
+              [.6] kv_offset       Remaining seq length from kv_end to the end of the batch.
+              [.7] pad             Padding to 8 DWs.
+        [3] reduce_indptr       (sum(qo_seqlen_blk_count) + 1)  Groups reduce_partial_map
+                                                entries that must be merged together.
+        [4] reduce_final_map    (sum(qo_seqlen_blk_count))      Final output location per group.
+        [5] reduce_partial_map  (#partial_tiles)                Partial-buffer locations of the
+                                                tiles awaiting reduction.
     """
     ...
 
 
+def get_mla_metadata_v1(
+    seqlens_qo_indptr: torch.Tensor,
+    seqlens_kv_indptr: torch.Tensor,
+    kv_last_page_lens: torch.Tensor,
+    num_heads_per_head_k: int,
+    num_heads_k: int,
+    is_causal: bool,
+    work_metadata_ptrs: torch.Tensor,
+    work_info_set: torch.Tensor,
+    work_indptr: torch.Tensor,
+    reduce_indptr: torch.Tensor,
+    reduce_final_map: torch.Tensor,
+    reduce_partial_map: torch.Tensor,
+    page_size: int = 1,
+    kv_granularity: int = 16,
+    max_seqlen_qo: int = -1,
+    uni_seqlen_qo: int = -1,
+    fast_mode: bool = True,
+    topk: int = -1,
+    max_split_per_batch: int = -1,
+    intra_batch_mode: bool = False,
+    is_cp_round_robin: bool = False,
+    mla_version: int = MlaVersion.V32.value,
+    dtype_q_nope: Optional[torch.dtype] = None,
+    dtype_q_rope: Optional[torch.dtype] = None,
+    dtype_kv_nope: Optional[torch.dtype] = None,
+    dtype_kv_rope: Optional[torch.dtype] = None,
+    dtype_q: Optional[torch.dtype] = None,
+    dtype_kv: Optional[torch.dtype] = None,
+) -> None:
+    """Forward-compatible wrapper around the compiled ``get_mla_metadata_v1``.
+
+    ``dtype_q`` / ``dtype_kv`` are legacy aliases kept for callers written
+    before the NoPE/RoPE dtype split. When provided, each fills in the
+    corresponding NoPE and RoPE dtype only if that specific one was not already
+    given (an explicit ``dtype_q_nope`` / ``dtype_q_rope`` etc. always wins).
+    See the compiled op's docstring for the full input/output description.
+    """
+    if dtype_q is not None:
+        if dtype_q_nope is None:
+            dtype_q_nope = dtype_q
+        if dtype_q_rope is None:
+            dtype_q_rope = dtype_q
+    if dtype_kv is not None:
+        if dtype_kv_nope is None:
+            dtype_kv_nope = dtype_kv
+        if dtype_kv_rope is None:
+            dtype_kv_rope = dtype_kv
+
+    return _get_mla_metadata_v1_impl(
+        seqlens_qo_indptr,
+        seqlens_kv_indptr,
+        kv_last_page_lens,
+        num_heads_per_head_k,
+        num_heads_k,
+        is_causal,
+        work_metadata_ptrs,
+        work_info_set,
+        work_indptr,
+        reduce_indptr,
+        reduce_final_map,
+        reduce_partial_map,
+        page_size=page_size,
+        kv_granularity=kv_granularity,
+        max_seqlen_qo=max_seqlen_qo,
+        uni_seqlen_qo=uni_seqlen_qo,
+        fast_mode=fast_mode,
+        topk=topk,
+        max_split_per_batch=max_split_per_batch,
+        intra_batch_mode=intra_batch_mode,
+        is_cp_round_robin=is_cp_round_robin,
+        mla_version=mla_version,
+        dtype_q_nope=dtype_q_nope,
+        dtype_q_rope=dtype_q_rope,
+        dtype_kv_nope=dtype_kv_nope,
+        dtype_kv_rope=dtype_kv_rope,
+    )
+
+
 @compile_ops("module_mla_metadata")
 def get_mla_metadata_v1_no_redundant(
     seqlens_qo_indptr: torch.Tensor,

csrc/kernels/mla/hk/doc/hk_mla_v40_gen1_spec.md

@@ -123,16 +123,18 @@ attention. The kernel's job is to make those many Q rows reuse each KV load.
 
 ### 2.2 The mfma tile and "m=16"
 
-We use the gfx950 fp8 mfma `v_mfma_f32_16x16x32_fp8_fp8` of shape
+We use the gfx950 bf16 mfma `v_mfma_f32_16x16x32_bf16` of shape
 
 $$
 (16 \times 32) \cdot (32 \times 16) \to (16 \times 16)
 $$
 
-for QK and the bf16 mfma `v_mfma_f32_16x16x32_f16` of the same shape for
-PV (after the fp8 P from softmax is cast to bf16). Both have the same
-m=16 / n=16 shape, so the per-ptile accumulator is a 16-row × 16-column
-fp32 tile.
+for **both** QK and PV. Although Q/K/V arrive as fp8 on disk, the QManager /
+KvManager cvt the operands to bf16 *before* the GEMM (`mfma_ab_t = hk::bf16`;
+P from softmax is likewise cast to bf16), so there is **no fp8 MFMA in this
+kernel** — QK and PV emit the identical `v_mfma_f32_16x16x32_bf16`. Both have
+the same m=16 / n=16 shape, so the per-ptile accumulator is a 16-row ×
+16-column fp32 tile.
 
 The "**m=16**" in the kernel name refers to this m dim: each ptile holds
 16 rows of $Q$ in its mfma accumulators. Those 16 rows are the work items
@@ -282,7 +284,7 @@ Critical sp3 / mfma instruction at each edge is annotated in parentheses.
    │  q_lds window Ph.B)│  │  3 tiles for N=64)   │
    └─────────┬──────────┘  └──────────┬───────────┘
              │                        │
-             └─── v_mfma_f32_16x16x32_fp8_fp8 ───┘   ← QK
+             └─── v_mfma_f32_16x16x32_bf16 ───┘   ← QK (operands cvt fp8→bf16 first)
                               │
                               ▼
                   ┌─────────────────────┐
@@ -300,7 +302,7 @@ Critical sp3 / mfma instruction at each edge is annotated in parentheses.
                   └─────────┬───────────┘
                             │   ── reuses KV-LDS pong as V via transpose-read
                             ▼
-              ── v_mfma_f32_16x16x32_f16 ──   ← PV
+              ── v_mfma_f32_16x16x32_bf16 ──   ← PV (same instr as QK)
                             │   (interleaved with v_mul_f32 rescale)
                             ▼
                   ┌─────────────────────┐
@@ -323,8 +325,7 @@ Critical sp3 / mfma instruction at each edge is annotated in parentheses.
 
 Key sp3 ops to remember:
 
-- **`v_mfma_f32_16x16x32_fp8_fp8`** — QK (single 16×32 × 32×16 → 16×16 fp32).
-- **`v_mfma_f32_16x16x32_f16`** — PV (same shape, bf16 inputs).
+- **`v_mfma_f32_16x16x32_bf16`** — **both** QK and PV (single 16×32 × 32×16 → 16×16 fp32). Operands are cvt fp8→bf16 before the GEMM, so this is the only MFMA in the kernel; there is no fp8 MFMA.
 - **`ds_read_b128`** — vanilla bf16 read for QK A-tile.
 - **`ds_read_b64_tr_b16`** — transpose-read for V → mfma A-operand layout (Ch. 10.2).
 - **`buffer_load_lds`** — direct vmem → LDS bypassing VGPRs (RoPE path + Q Phase 1 staging).
@@ -1656,7 +1657,7 @@ and read `oaccu` as $O^\top$ (col-major), `kv` as $V^\top$ in A-operand
 layout, `p_mfma` as $P^\top$ in B-operand layout. This is identical to
 how QK was already running ($K^\top Q^\top = S^\top$).
 
-The mfma is `v_mfma_f32_16x16x32_f16` of shape
+The mfma is `v_mfma_f32_16x16x32_bf16` of shape
 
 $$
 (16 \times 32) \cdot (32 \times 16) \to (16 \times 16)
@@ -1723,7 +1724,7 @@ One PV iter, canonical case (`kIsFirstIter == false`, `has_next == true`):
 So per PV iter:
 
 - 4× `ds_read_b64_tr_b16` (V load)
-- 2× `v_mfma_f32_16x16x32_f16` (the PV mfmas)
+- 2× `v_mfma_f32_16x16x32_bf16` (the PV mfmas)
 - 4× `v_mul_f32` rescaling the NEXT iter's oaccu base tile +0/+1 (interleaved with ds_read)
 - 4× `v_mul_f32` rescaling the NEXT iter's oaccu base tile +2/+3 (1 mul_pair per mfma slot)
 - 2× `s_waitcnt`
@@ -1776,7 +1777,7 @@ beyond the one ballot).
 ### 10.5 Why this pattern hides everything
 
 The numbers work out because gfx950's mfma occupancy budget per lane is
-generous: each `v_mfma_f32_16x16x32_f16` issue spends ~32 cycles in the
+generous: each `v_mfma_f32_16x16x32_bf16` issue spends ~32 cycles in the
 mfma pipe, during which the VALU is free. The pattern fills both halves:
 
 | Time slice within one PV iter | mfma pipe | VALU | LDS |
@@ -1792,6 +1793,65 @@ The schedule is dense — there's no slot where mfma is idle waiting on
 VALU or LDS. The rescale, naïvely a 128-mul standalone phase before PV,
 is fully hidden.
 
+### 10.5.1 The cross-wave PV ping-pong (why 6-deep PV + warp specialization + deferred B+C exist)
+
+§10.5 hides the rescale *within* one wave. But PV itself — a deep mfma
+sequence — saturates the **shared MFMA + LDS pipeline** for the wave running
+it. On m16x8 two waves share each SIMD (8 waves/workgroup over 4 SIMDs =
+2 waves/SIMD; `kOccupancy_ = 1`, so this is intra-workgroup, *not* two
+resident workgroups). While one wave is in PV, its SIMD-mate physically
+cannot issue MFMA/LDS work — it can only run non-MFMA work (softmax VALU,
+dispatch-ladder SALU, next-KV vmem loads). So the *only* way to hide a PV is
+to have the partner wave doing that other work underneath it.
+
+Three design choices exist purely to make this cross-wave overlap work — they
+are not independent optimizations, they all serve the same ping-pong:
+
+1. **PV is made deep (contracts both sub-tiles A+B in one merged call, §10.1)**
+   so one wave's PV window is *long enough to cover the partner's entire
+   non-MFMA workload* (softmax + dispatch SALU + next-KV load). The depth is
+   chosen for the hiding-window length, not that wave's own throughput; too
+   shallow and the partner's work spills past the window and becomes exposed.
+
+2. **Warp specialization (`WarpTypeM16x8` Lo/Hi, §3.5) keeps the two
+   SIMD-mate waves out of phase.** `kPvAtEnd = (kWarpType == LoNoPEWarp)`:
+   lo runs PV at call end, hi *defers* PV by one tile (`kHasPv`, runs the
+   previous tile's PV at the next call's start). If both waves ran the same
+   schedule they would hit PV simultaneously, contend for the one MFMA+LDS
+   pipe, and have nothing to overlap. The offset guarantees that at any moment
+   one wave is in PV while the other is in non-MFMA work.
+
+3. **Deferring part of lo's Phase B+C behind PV (deferred strip-3 cvt+store,
+   §8.13)** does two things: (a) it packs that non-MFMA KV cvt/store work
+   *into* the partner's PV window instead of leaving it as exposed serial
+   time; and (b) with less B+C ahead of lo's PV, **lo starts its PV earlier**,
+   returning the shared MFMA+LDS pipe to hi sooner. hi's wait on lo's PV is a
+   **resource dependency on the shared pipe, not an `s_barrier`** — hi stalls
+   because the pipe is busy, so pulling lo's PV earlier lets hi acquire the
+   pipe sooner *and* slides lo's PV window earlier so it overlaps more of hi's
+   softmax. It is a critical-path move, not just window-packing.
+
+Net per-tile picture (the two waves offset by ~half a tile): while one wave
+runs PV (owning MFMA+LDS), its partner runs softmax + dispatch-ladder
+SALU/branch + next-KV load + deferred B+C; the next tile the roles swap. This
+is why an instruction trace shows a *higher* per-wave stall% yet a busier CU
+(lower idle) — the "stall" is a wave parked on the shared pipe, and that gap
+is filled by the partner's useful work.
+
+**Why m16x4 does not use this scheme.** m16x4 also puts 2 waves on each SIMD,
+but differently: `kOccupancy_ = 2` runs **two whole workgroups** per CU
+(4 waves/tg × 2 tgs = 8 waves), so a SIMD's two waves belong to *different
+workgroups*. `s_barrier` only synchronizes waves within one workgroup, so
+there is no way to keep those two SIMD-mates in a deliberate PV/non-PV phase
+relationship — the mate's scheduling is uncontrollable across the workgroup
+boundary. m16x4 therefore forgoes the ping-pong; it just issues the next-KV
+load earlier. That is the right trade for m16x4 because, *compared to m16x8*,
+it leans more toward memory-bound — though whether either partition is
+actually memory-bound in a given run depends on the real situation (HW state,
+and what other programs / tgs are co-resident in the grid), not always. The
+deliberate cross-wave PV overlap is an m16x8-only design, enabled by both
+SIMD-mates living in the *same* workgroup.
+
 ### 10.6 `pv_v_aux` is dead in Gen.1
 
 The pinned VGPR map in Ch. 5 lists `pv_v_aux` v104..v111 as "second

csrc/kernels/mla/reduce.cu

@@ -932,6 +932,7 @@ __launch_bounds__(Traits::kNumThreads, Traits::kOccupancy) __global__
     MLA_REDUCE_CASE_EF(NUM_HEAD, 32, HEAD_DIM, 128, NUM_WG_PER_BH, NAME, __VA_ARGS__)  \
     MLA_REDUCE_CASE_EF(NUM_HEAD, 32, HEAD_DIM, 512, NUM_WG_PER_BH, NAME, __VA_ARGS__)  \
     MLA_REDUCE_CASE_EF(NUM_HEAD, 40, HEAD_DIM, 128, NUM_WG_PER_BH, NAME, __VA_ARGS__)  \
+    MLA_REDUCE_CASE_EF(NUM_HEAD, 64, HEAD_DIM, 64, NUM_WG_PER_BH, NAME, __VA_ARGS__)   \
     MLA_REDUCE_CASE_EF(NUM_HEAD, 64, HEAD_DIM, 128, NUM_WG_PER_BH, NAME, __VA_ARGS__)  \
     MLA_REDUCE_CASE_EF(NUM_HEAD, 64, HEAD_DIM, 512, NUM_WG_PER_BH, NAME, __VA_ARGS__)  \
     MLA_REDUCE_CASE_EF(NUM_HEAD, 128, HEAD_DIM, 128, NUM_WG_PER_BH, NAME, __VA_ARGS__) \