From 193e8d8ea029b936ba8a03997db77179545b339f Mon Sep 17 00:00:00 2001 From: liyjiang Date: Mon, 20 Jul 2026 23:00:30 -0400 Subject: [PATCH 1/2] [MLA v4 nm] V3-style out_16_nosplit + zero-copy single-pass output Adopt V3's approach for the v4 nm decode path: out_16_nosplit is DERIVED internally from the resolved pass count instead of being a caller knob. - asm_mla_v4.cu: out_16_nosplit = (num_kv_splits == 1) ? 1 : 0, mirroring asm_mla.cu. The caller-facing arg is ignored (kept for C-ABI parity). - mla.py: derive the same value in the wrapper so the kernarg and the buffer/return handling stay in lockstep. Single pass now ALWAYS uses the bf16-direct write, and `logits` is aliased onto the caller's `output` (bf16 view) so the kernel lands the final packed-BF16 result straight in `output` -- eliminating the ~5us post-launch output.copy_(packed_bf16). Multi-pass keeps the fp32-split + stage2-merge path unchanged. Tests updated for the new single-pass semantics: byte-level sink diff reads int16 (bf16) bits; multi_split (B) asserts the ignored arg does not raise and still runs the fp32-split path; kernarg slot-15 expectation is now 1 for the single-pass config. Co-authored-by: Cursor --- aiter/mla.py | 44 +++++++++++------------------------ csrc/py_itfs_cu/asm_mla_v4.cu | 5 +++- op_tests/test_mla_v4_nm.py | 43 ++++++++++++++++++++++++---------- 3 files changed, 49 insertions(+), 43 deletions(-) diff --git a/aiter/mla.py b/aiter/mla.py index 1bf72461c43..9746161d678 100644 --- a/aiter/mla.py +++ b/aiter/mla.py @@ -1293,7 +1293,7 @@ def mla_decode_fwd_v4_nm( 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, @@ -1310,9 +1310,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 +1406,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 +1430,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, @@ -1548,22 +1550,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 diff --git a/csrc/py_itfs_cu/asm_mla_v4.cu b/csrc/py_itfs_cu/asm_mla_v4.cu index 4f5f1593960..5ff05c617ac 100644 --- a/csrc/py_itfs_cu/asm_mla_v4.cu +++ b/csrc/py_itfs_cu/asm_mla_v4.cu @@ -244,6 +244,9 @@ AITER_CTYPES_DEFINE_ENTRYPOINT_VOID( 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__, @@ -286,7 +289,7 @@ AITER_CTYPES_DEFINE_ENTRYPOINT_VOID( a.s_kv_split = static_cast(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(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(); diff --git a/op_tests/test_mla_v4_nm.py b/op_tests/test_mla_v4_nm.py index 83929c801a8..e38c3c960af 100644 --- a/op_tests/test_mla_v4_nm.py +++ b/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 @@ -1504,8 +1507,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 +1552,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 +1654,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. From 14de8657bd15b9865ce9741ccdfdf233f0d95380 Mon Sep 17 00:00:00 2001 From: amd-ruitang3 Date: Tue, 21 Jul 2026 04:25:15 -0500 Subject: [PATCH 2/2] update_kv_last_len --- aiter/mla.py | 7 +++++-- aiter/ops/attention.py | 7 +++++-- csrc/py_itfs_cu/asm_mla_v4.cu | 9 ++++++--- op_tests/test_mla_v4_nm.py | 1 - 4 files changed, 16 insertions(+), 8 deletions(-) diff --git a/aiter/mla.py b/aiter/mla.py index 9746161d678..4d8e302bcd5 100644 --- a/aiter/mla.py +++ b/aiter/mla.py @@ -1287,7 +1287,6 @@ 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 @@ -1297,6 +1296,10 @@ def mla_decode_fwd_v4_nm( 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. @@ -1485,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, @@ -1497,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 ------ diff --git a/aiter/ops/attention.py b/aiter/ops/attention.py index 958a594a3b7..496655d2cee 100644 --- a/aiter/ops/attention.py +++ b/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: ... diff --git a/csrc/py_itfs_cu/asm_mla_v4.cu b/csrc/py_itfs_cu/asm_mla_v4.cu index 5ff05c617ac..cfd73118569 100644 --- a/csrc/py_itfs_cu/asm_mla_v4.cu +++ b/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,6 +243,7 @@ AITER_CTYPES_DEFINE_ENTRYPOINT_VOID( output, valid_split_count, use_valid_split_count_reduce, + kv_last_page_lens, stream)) { (void)softmax_scale; @@ -283,7 +286,7 @@ 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(gqa_ratio); a.s_kv_split = static_cast(num_kv_splits); diff --git a/op_tests/test_mla_v4_nm.py b/op_tests/test_mla_v4_nm.py index e38c3c960af..da657bda687 100644 --- a/op_tests/test_mla_v4_nm.py +++ b/op_tests/test_mla_v4_nm.py @@ -904,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"],