[MLA v4 nm] Zero-copy final output for out_16_nosplit=1

aiter/mla.py

@@ -1287,16 +1287,19 @@ def mla_decode_fwd_v4_nm(
     qo_indptr,  # [num_seqs+1]
     kv_indptr,  # [num_seqs+1]
     kv_page_indices,  # [num_page_used]
-    kv_last_page_lens,  # [num_seqs]
     max_seqlen_q,
     *,
     sink,  # REQUIRED [num_heads] FP32 -- attention sink logit
     split_indptr=None,  # [num_seqs+1] int32; auto-built by get_meta_param if None
     sm_scale=None,  # ignored on v4 nm; kernel hardcodes 1/sqrt(512)
-    out_16_nosplit=0,
+    out_16_nosplit=0,  # DEPRECATED/IGNORED: derived internally from num_kv_splits (V3-style)
     num_kv_splits=None,  # None -> auto-pick via V3 get_meta_param heuristic
     logits=None,
     attn_lse=None,
+    kv_last_page_lens=None,  # [num_seqs] int32; OPTIONAL, defaults None. Unused on
+    # the nm path (page_size=1 -> kv_seq_len comes from the token-level kv_indptr).
+    # None flows through to a nullptr kernarg; the host guards the deref and the
+    # kernel never loads through it, so no buffer is allocated.
 ):
     """v4 MLA decode forward.
 
@@ -1310,9 +1313,12 @@ def mla_decode_fwd_v4_nm(
     convention).
 
     Single-pass mode (`num_kv_splits == 1`):
-      Returned `logits[:, 0]` already holds the per-token output (no merge
-      needed). `output[total_q, num_heads, v_head_dim]` BF16 is only written
-      by the kernel when `out_16_nosplit == 1`.
+      V3-style: the kernel ALWAYS writes the final packed-BF16 result straight
+      into `output` (zero-copy). `out_16_nosplit` is derived internally (=1
+      here) and the caller-facing arg is ignored. The returned `logits` is a
+      BF16 view aliased onto `output` (shape [total_q, 1, num_heads,
+      v_head_dim]); read the per-token result from `output` (or `logits[:, 0]`,
+      same bytes). There is no separate FP32 partial to reduce.
 
     Split-count selection (`num_kv_splits`):
       `None` (default) auto-picks the split count via V3's `get_meta_param`
@@ -1403,13 +1409,7 @@ def mla_decode_fwd_v4_nm(
         if split_indptr is None:
             split_indptr = meta_split_indptr.to(device=q.device, dtype=torch.int32)
 
-    # ---- Multi-pass requires the FP32 split-output path ----
-    if num_kv_splits > 1 and out_16_nosplit != 0:
-        raise ValueError(
-            f"mla_decode_fwd_v4_nm: num_kv_splits={num_kv_splits} requires "
-            f"out_16_nosplit=0 (the kernel's bf16-direct-write path only "
-            f"supports passes==1). Got out_16_nosplit={out_16_nosplit}."
-        )
+    out_16_nosplit = 1 if num_kv_splits == 1 else 0
 
     # ---- Allocate splitData / splitLse in KERNEL-NATIVE layout --------------
     # kernel-native layout:   [num_seqs, max_seqlen_q, num_kv_splits, num_heads, v_head_dim]
@@ -1433,7 +1433,12 @@ def mla_decode_fwd_v4_nm(
 
     expected_logits_shape = (total_q, num_kv_splits, num_heads, v_head_dim)
     expected_lse_shape = (total_q, num_kv_splits, num_heads, 1)
-    if logits is None:
+    if out_16_nosplit != 0:
+        # V3-style zero-copy final output. When out_16_nosplit=1 (single-pass),
+        # the kernel writes the final DENSELY-PACKED BF16 result into ptr_R
+        # (= the `logits`/splitData buffer).
+        logits = output.view(expected_logits_shape)
+    elif logits is None:
         logits = torch.empty(
             expected_logits_shape,
             dtype=dtypes.fp32,
@@ -1483,7 +1488,6 @@ def mla_decode_fwd_v4_nm(
         qo_indptr,
         kv_indptr,
         kv_page_indices,
-        kv_last_page_lens,
         split_indptr,
         sink,
         max_seqlen_q,
@@ -1495,6 +1499,7 @@ def mla_decode_fwd_v4_nm(
         output,
         valid_split_count,
         use_valid_split_count_reduce,
+        kv_last_page_lens,  # tail: unused on nm path, None -> nullptr
     )
 
     # ---- Cross-split FlashAttention merge via _fwd_kernel_stage2_asm ------
@@ -1548,22 +1553,4 @@ def mla_decode_fwd_v4_nm(
             waves_per_eu=4,
         )
 
-    # ---- out_16_nosplit=1: resolve the kernel's packed-BF16 output ----------
-    # When out_16_nosplit=1 (single-pass only), the kernel does NOT write the
-    # `output` buffer. Instead it writes the final result as DENSELY-PACKED
-    # BF16 into the `logits` allocation. Global layout is identical to the
-    # fp32 path -- [total_q, nsplit=1, num_heads, v_head_dim] contiguous in
-    # num_heads -- but each element is 2 bytes (R16_BPP=2), so the per-head
-    # stride is v_head_dim*2 bytes = v_head_dim BF16 slots, occupying only the
-    # FIRST half of the fp32 `logits` byte range. Reinterpret those bytes as a
-    # tight [total_q, num_heads, v_head_dim] BF16 tensor and copy into the
-    # caller's `output` so `output` is the single authoritative result buffer
-    # (mirrors the multi-split stage2 path, which also lands in `output`).
-    if out_16_nosplit != 0:
-        n_bf16 = total_q * num_heads * v_head_dim
-        packed_bf16 = (
-            logits.contiguous().view(torch.bfloat16).flatten()[:n_bf16]
-        ).view(total_q, num_heads, v_head_dim)
-        output.copy_(packed_bf16)
-
     return logits, attn_lse

aiter/ops/attention.py

@@ -848,8 +848,6 @@ def mla_decode_v4_asm(
     kv_indptr: torch.Tensor,
     # [num_page_used]
     kv_page_indices: torch.Tensor,
-    # [num_seqs]
-    kv_last_page_lens: torch.Tensor,
     # [num_seqs+1]
     split_indptr: torch.Tensor,
     # [num_heads] FP32 — attention sink logit. Loaded by the kernel via
@@ -876,6 +874,11 @@ def mla_decode_v4_asm(
     # kernel skips the write and nullptr is fine.
     valid_split_count: Optional[torch.Tensor] = None,
     use_valid_split_count_reduce: int = 0,
+    # [num_seqs] int32. Unused on the v4 nm path (page_size=1 -> the kernel derives
+    # kv_seq_len from the token-level kv_indptr). Optional/nullable: None sends a
+    # nullptr; the host guards the deref (asm_mla_v4.cu) and the kernel never loads
+    # through it. Placed at the tail because it carries no data on this path.
+    kv_last_page_lens: Optional[torch.Tensor] = None,
 ) -> None: ...
 
 

csrc/py_itfs_cu/asm_mla_v4.cu

@@ -206,7 +206,6 @@ AITER_CTYPES_DEFINE_ENTRYPOINT_VOID(
      aiter_tensor_t* qo_indptr,       // [num_seqs+1]
      aiter_tensor_t* kv_indptr,       // [num_seqs+1]
      aiter_tensor_t* kv_page_indices, // [num_page_used]
-     aiter_tensor_t* kv_last_page_lens, // [num_seqs]
      aiter_tensor_t* split_indptr,      // [num_seqs+1]
      aiter_tensor_t* sink,              // [num_heads] FP32 — see "ptr_sink" note above
      int max_seqlen_q,
@@ -221,6 +220,10 @@ AITER_CTYPES_DEFINE_ENTRYPOINT_VOID(
          output, // [total_query_len, num_heads, v_head_dim] BF16 (used when out_16_nosplit==1)
      aiter_tensor_t* valid_split_count, // [num_seqs] int32 scratch (slot 19), nullable
      int use_valid_split_count_reduce,  // slot 20 flag; gates the kernel's valid_split write
+     // Moved to the tail: UNUSED on the nm path (page_size=1 -> kv_seq_len comes
+     // from the token-level kv_indptr). Nullable; the host guards the deref below
+     // and the kernel never loads through the pointer.
+     aiter_tensor_t* kv_last_page_lens, // [num_seqs] int32, nullable
      hipStream_t stream),
     (Q,
      qrope,
@@ -229,7 +232,6 @@ AITER_CTYPES_DEFINE_ENTRYPOINT_VOID(
      qo_indptr,
      kv_indptr,
      kv_page_indices,
-     kv_last_page_lens,
      split_indptr,
      sink,
      max_seqlen_q,
@@ -241,9 +243,13 @@ AITER_CTYPES_DEFINE_ENTRYPOINT_VOID(
      output,
      valid_split_count,
      use_valid_split_count_reduce,
+     kv_last_page_lens,
      stream))
 {
     (void)softmax_scale;
+    (void)out_16_nosplit;
+    const unsigned int out_16_nosplit_derived =
+        (num_kv_splits == 1) ? 1u : 0u;
     AITER_CHECK(sink != nullptr, __func__, ": `sink` must not be NULL");
     AITER_CHECK(sink->data_ptr() != nullptr,
                 __func__,
@@ -280,13 +286,13 @@ AITER_CTYPES_DEFINE_ENTRYPOINT_VOID(
         a.ptr_KV      = KV->data_ptr();
         a.ptr_LTP     = kv_indptr->data_ptr();
         a.ptr_LTD     = kv_page_indices->data_ptr();
-        a.ptr_LTL     = kv_last_page_lens->data_ptr();
+        a.ptr_LTL     = kv_last_page_lens ? kv_last_page_lens->data_ptr() : nullptr;
         a.scalar_f    = scalar_f;
         a.s_gqa_ratio = static_cast<unsigned int>(gqa_ratio);
         a.s_kv_split  = static_cast<unsigned int>(num_kv_splits);
         // a.s_total_kv     left as 0 here — set per-arch in fill_gfx1250_kargs below.
         a.ptr_QTP        = qo_indptr->data_ptr();
-        a.out_16_nosplit = static_cast<unsigned int>(out_16_nosplit);
+        a.out_16_nosplit = out_16_nosplit_derived;
         a.ptr_QROPE      = qrope->data_ptr();
         a.ptr_KVROPE     = kvrope->data_ptr();
         a.ptr_sink       = sink->data_ptr();

op_tests/test_mla_v4_nm.py

@@ -259,7 +259,10 @@ def slot_f32(i):
     expected_gqa_ratio = GQA_RATIO  # 16
     expected_kv_split = 1  # num_kv_splits=1
     expected_log2_page = 0  # log2(page_size=1)
-    expected_out16ns = 0  # out_16_nosplit=0
+    # V3-style: out_16_nosplit is DERIVED = (num_kv_splits==1) ? 1 : 0. This
+    # config is single-pass (kv_split=1) so the dispatcher writes 1 into slot 15
+    # regardless of the caller-facing arg (which _build_inputs leaves default).
+    expected_out16ns = 1
     # slots 10 (s_total_kv) and 11 (s_stride_page) are NEVER read — only 17 kernarg
     # loads, none at offsets 0xA0/0xB0). The dispatcher leaves them at 0
     # via `args = {}` zero-init to skip the per-call D2H readback that
@@ -901,7 +904,6 @@ def _run_one_point(
         inputs["qo_indptr"],
         inputs["kv_indptr"],
         inputs["kv_page_indices"],
-        inputs["kv_last_page_lens"],
         split_indptr,
         inputs["sink"],  # per-head [num_heads] sink; req'd positional
         inputs["max_seqlen_q"],
@@ -1504,8 +1506,9 @@ def test_v4_nm_multi_split():
         V3-style, and the kernel doesn't tail-drop any split. Coverage
         invariant: floor(kv/splits) >= SUB_KV (=32); 256/4=64 ✓.
 
-    (B) Rejection: multi-pass + out_16_nosplit=1 is unsupported (bf16-direct
-        is single-pass only); the wrapper must raise BEFORE the kernel.
+    (B) out_16_nosplit is derived internally (V3-style) from num_kv_splits, so
+        the caller-facing arg is ignored: multi-pass + out_16_nosplit=1 must NOT
+        raise and must still run the fp32-split + stage2-merge path.
     """
     # ---- (A) full-KV coverage ----
     NUM_SPLITS = 4
@@ -1548,13 +1551,25 @@ def test_v4_nm_multi_split():
             f"geometry / split_indptr stride / kernel early-exit)."
         )
 
-    # ---- (B) multi-pass + out_16_nosplit=1 must be rejected ----
-    args_rej = _build_inputs(batch=1, kv_seq_lens=128, q_seq_logical=1, seed=0)
-    args_rej["num_kv_splits"] = 2
-    args_rej["out_16_nosplit"] = 1
-    args_rej.pop("split_indptr")
-    with pytest.raises(ValueError, match="out_16_nosplit"):
-        aiter.mla.mla_decode_fwd_v4_nm(**args_rej)
+    # ---- (B) out_16_nosplit is now DERIVED (V3-style), so the caller-facing
+    # arg is ignored: passing out_16_nosplit=1 with multi-pass must NOT raise
+    # and must still run the fp32-split + stage2-merge path (the wrapper/kernel
+    # derive out_16_nosplit=0 from num_kv_splits>1). ----
+    args_ign = _build_inputs(batch=1, kv_seq_lens=128, q_seq_logical=1, seed=0)
+    args_ign["num_kv_splits"] = 2
+    args_ign["out_16_nosplit"] = 1  # ignored; derived to 0 for multi-pass
+    args_ign.pop("split_indptr")
+    out_ign, lse_ign = aiter.mla.mla_decode_fwd_v4_nm(**args_ign)
+    torch.cuda.synchronize()
+    # Multi-pass returns fp32 split logits (NOT the bf16 single-pass alias).
+    assert out_ign.dtype == torch.float32, (
+        f"multi-pass logits should be fp32 split partials, got {out_ign.dtype} "
+        f"— out_16_nosplit=1 was NOT correctly overridden to 0."
+    )
+    assert out_ign.shape[1] == 2, (
+        f"multi-pass logits should keep the num_kv_splits=2 axis, got "
+        f"shape {tuple(out_ign.shape)}."
+    )
 
 
 # ---------------------------------------------------------------------------
@@ -1638,11 +1653,14 @@ def test_v4_nm_sink():
 
     logits_a, _ = aiter.mla.mla_decode_fwd_v4_nm(**args_a)
     torch.cuda.synchronize()
-    logits_a_bits = logits_a.view(torch.int32).clone()
+    # Single-pass (V3-style) now returns packed-BF16 (2-byte) logits aliased
+    # onto `output`, so reinterpret the raw bits as int16 (not int32) to keep
+    # the last-dim aligned with the finite mask for the byte-level sink diff.
+    logits_a_bits = logits_a.view(torch.int16).clone()
 
     logits_b, _ = aiter.mla.mla_decode_fwd_v4_nm(**args_b)
     torch.cuda.synchronize()
-    logits_b_bits = logits_b.view(torch.int32)
+    logits_b_bits = logits_b.view(torch.int16)
 
     # logits is 4D [total_q, num_kv_splits=1, num_heads, dv]; only [:, 0]
     # is kernel-written.