diff --git a/aiter/mla.py b/aiter/mla.py index 84ffc00b48f..a43084f052c 100644 --- a/aiter/mla.py +++ b/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 diff --git a/aiter/ops/attention.py b/aiter/ops/attention.py index db8799ebeb6..46306baa5e7 100644 --- a/aiter/ops/attention.py +++ b/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] diff --git a/csrc/py_itfs_cu/asm_mla_v4.cu b/csrc/py_itfs_cu/asm_mla_v4.cu index cfd73118569..de0a345b1ea 100644 --- a/csrc/py_itfs_cu/asm_mla_v4.cu +++ b/csrc/py_itfs_cu/asm_mla_v4.cu @@ -29,6 +29,7 @@ #include #include #include +#include #include #include #include @@ -542,3 +543,258 @@ AITER_CTYPES_DEFINE_ENTRYPOINT_VOID( impl_ptr->launch_kernel({arg_buf, &arg_size, gdx, gdy, gdz, block_dim, 1, 1, stream}); } + +// ---------------------------------------------------------------------------- +// Persistent-scheduling v4 nm decode (CSV row ps=1) +// +// One launch per call: grid (2 * P, 1, 1), block 256. The kernel plans the +// KV split over P partitions in-kernel from kv_indptr, runs the attention +// and merges the split partials itself (the last arriving partition of a row +// combines it), so there is no host split plan and no stage-2 merge. +// Same math and packed Q / KV / rope layouts as mla_decode_v4_asm with +// gqa=128, max_seqlen_q=1 and page_size=1; rows with K=0 are left unwritten. +// +// The kernarg block keeps the 21 x 16 B layout of MlaV4KernelArgsLegacy; +// slots marked "repurposed" carry a different value than the non-persistent +// kernel, and several upper halves of the pointer slots carry an extra field. +// ---------------------------------------------------------------------------- +struct __attribute__((packed)) MlaV4PsKernelArgs +{ + void* ptr_O_acc; // 0x00: workspace o_acc [2P, 128, 512] FP32 split partials + // (rw) + void* ptr_out; // 0x08: out [N, 128, 512] BF16 (upper half of slot 0) + void* ptr_L_acc; // 0x10: workspace lse_acc [2P, 128] FP32 split partials (rw) + void* ptr_lse; // 0x18: final LSE [N, 128] FP32, nullptr = not written + void* ptr_Q; + p2 _p_q; // 0x20: Q packed FP8 + e8m0 scale + void* ptr_KV; + p2 _p_kv; // 0x30: KV packed FP8 + void* ptr_LTP; + p2 _p_ltp; // 0x40: kv_indptr + void* ptr_LTD; + p2 _p_ltd; // 0x50: kv_page_indices + void* ptr_desc; // 0x60: repurposed: workspace desc [P, 8] int32 (rw) + void* ptr_dbg; // 0x68: unused, nullptr + float scalar_f; + p3 _p_sc; // 0x70: 1.0f/sqrtf(kV4DimNope+kV4DimRope) + unsigned int s_gqa_ratio; + p3 _p_gr; // 0x80: 128 (heads) + unsigned int s_kv_split; // 0x90: 1 + unsigned int _p_ks; // 0x94: pad + void* ptr_cnt; // 0x98: workspace cnt int32 counters (rw, zero at rest) + unsigned int s_num_p; // 0xa0: repurposed: number of partitions P + unsigned int s_num_q; // 0xa4: number of query rows N + p2 _p_nq; // 0xa8: pad + unsigned int s_plan; // 0xb0: repurposed: planner config F | (MT << 8) + unsigned int _p_pl; // 0xb4: pad + void* ptr_unused; // 0xb8: not read by this kernel, nullptr + unsigned int s_zero; + p3 _p_z; // 0xc0: 0 + void* ptr_QTP; + p2 _p_qtp; // 0xd0: repurposed: workspace arange int32 (read-only) + void* ptr_STP; + p2 _p_stp; // 0xe0: repurposed: the same arange + unsigned int out_16_nosplit; + p3 _p_o16; // 0xf0: 1 + void* ptr_QROPE; + p2 _p_qrope; // 0x100: Q rope BF16 + void* ptr_KVROPE; + p2 _p_kvrope; // 0x110: KV rope BF16 + void* ptr_sink; + p2 _p_sink; // 0x120: [128] FP32 attention sink logit + void* ptr_zero; + p2 _p_pz; // 0x130: nullptr + unsigned int s_zero2; + p3 _p_z2; // 0x140: 0 +}; +static_assert(sizeof(MlaV4PsKernelArgs) == 21 * 16, "persistent v4 nm kernarg is 0x150 bytes"); +static_assert(offsetof(MlaV4PsKernelArgs, ptr_cnt) == 0x98 && + offsetof(MlaV4PsKernelArgs, s_num_q) == 0xa4 && + offsetof(MlaV4PsKernelArgs, ptr_unused) == 0xb8 && + offsetof(MlaV4PsKernelArgs, ptr_sink) == 0x120, + "persistent v4 nm kernarg offsets"); + +// Workspace sizes the kernel indexes into (aiter/mla.py +// get_mla_v4_nm_ps_workspace). +// cnt holds [0, 2*65536) row counters, [2*65536, +16*512) reserved (unused) +// and [.., +4*1024) group counters. +static constexpr int kV4PsHeads = 128; +static constexpr int kV4PsDim = kV4DimNope + kV4DimRope; +static constexpr int kV4PsMaxParts = 1024; +static constexpr int kV4PsMaxRows = 32768; // byte offsets of out / q wrap above this +static constexpr int64_t kV4PsArange = 65537; +static constexpr int64_t kV4PsCntInts = 2 * 65536 + 16 * 512 + 4 * 1024; +static constexpr unsigned int kV4PsPlanCfg = 6u | (1u << 8); + +AITER_CTYPES_DEFINE_ENTRYPOINT_VOID( + mla_decode_v4_ps_asm, + (aiter_tensor_t * Q, // [N, 128, 512] FP8 packed Q + e8m0 + aiter_tensor_t* qrope, // [N, 128, 64] BF16 + aiter_tensor_t* KV, // [rows, ..., 512] FP8 packed KV, row-dense + aiter_tensor_t* kvrope, // [rows, ..., 64] BF16, row-dense + aiter_tensor_t* kv_indptr, // [>= N+1] int32 + aiter_tensor_t* kv_page_indices, // [*] int32 + aiter_tensor_t* sink, // [128] FP32 + aiter_tensor_t* o_acc, // workspace [2P, 128, 512] FP32 + aiter_tensor_t* lse_acc, // workspace [2P, 128] FP32 + aiter_tensor_t* desc, // workspace [P, 8] int32 + aiter_tensor_t* cnt, // workspace [kV4PsCntInts] int32, zero at rest + aiter_tensor_t* arange, // workspace arange(kV4PsArange) int32 + aiter_tensor_t* output, // [N, 128, 512] BF16 + aiter_tensor_t* lse, // [N, 128] FP32, nullable + hipStream_t stream), + (Q, + qrope, + KV, + kvrope, + kv_indptr, + kv_page_indices, + sink, + o_acc, + lse_acc, + desc, + cnt, + arange, + output, + lse, + stream)) +{ + auto check_buf = [&](aiter_tensor_t* t, AiterDtype dt, const char* name) { + AITER_CHECK(t != nullptr && t->data_ptr() != nullptr, + "mla_decode_v4_ps_asm", + ": `", + name, + "` is NULL"); + AITER_CHECK(t->dtype() == dt, + "mla_decode_v4_ps_asm", + ": `", + name, + "` has dtype ", + AiterDtype_to_str(t->dtype()), + ", expected ", + AiterDtype_to_str(dt)); + AITER_CHECK( + t->is_contiguous(), "mla_decode_v4_ps_asm", ": `", name, "` must be contiguous"); + AITER_CHECK(t->device_id == Q->device_id, + "mla_decode_v4_ps_asm", + ": `", + name, + "` is on another device"); + }; + AITER_CHECK(Q != nullptr, __func__, ": `Q` is NULL"); + check_buf(Q, AITER_DTYPE_fp8, "Q"); + check_buf(qrope, AITER_DTYPE_bf16, "qrope"); + check_buf(KV, AITER_DTYPE_fp8, "KV"); + check_buf(kvrope, AITER_DTYPE_bf16, "kvrope"); + check_buf(kv_indptr, AITER_DTYPE_i32, "kv_indptr"); + check_buf(kv_page_indices, AITER_DTYPE_i32, "kv_page_indices"); + check_buf(sink, AITER_DTYPE_fp32, "sink"); + check_buf(o_acc, AITER_DTYPE_fp32, "o_acc"); + check_buf(lse_acc, AITER_DTYPE_fp32, "lse_acc"); + check_buf(desc, AITER_DTYPE_i32, "desc"); + check_buf(cnt, AITER_DTYPE_i32, "cnt"); + check_buf(arange, AITER_DTYPE_i32, "arange"); + check_buf(output, AITER_DTYPE_bf16, "output"); + if(lse != nullptr && lse->data_ptr() != nullptr) + check_buf(lse, AITER_DTYPE_fp32, "lse"); + + AITER_CHECK(Q->dim() == 3 && Q->size(1) == kV4PsHeads && Q->size(2) == kV4PsDim, + __func__, + ": Q must be [N, 128, 512]"); + const int64_t num_q = Q->size(0); + AITER_CHECK(num_q <= kV4PsMaxRows, __func__, ": N=", num_q, " exceeds ", kV4PsMaxRows); + AITER_CHECK(qrope->numel() == static_cast(num_q * kV4PsHeads * kV4DimRope), + __func__, + ": qrope must be [N, 128, 64]"); + AITER_CHECK(output->dim() == 3 && output->size(0) == num_q && output->size(1) == kV4PsHeads && + output->size(2) == kV4PsDim, + __func__, + ": output must be [N, 128, 512]"); + if(lse != nullptr && lse->data_ptr() != nullptr) + AITER_CHECK(lse->numel() == static_cast(num_q * kV4PsHeads), + __func__, + ": lse must be [N, 128]"); + AITER_CHECK(KV->size(-1) == kV4PsDim && kvrope->size(-1) == kV4DimRope && + KV->numel() / kV4PsDim == kvrope->numel() / kV4DimRope, + __func__, + ": KV / kvrope must be row-dense [rows, 512] / [rows, 64] with " + "equal rows"); + AITER_CHECK(static_cast(kv_indptr->numel()) >= num_q + 1, + __func__, + ": kv_indptr needs at least N+1 entries"); + AITER_CHECK(sink->numel() == kV4PsHeads, __func__, ": sink must have 128 entries"); + + AITER_CHECK(desc->dim() == 2 && desc->size(1) == 8, __func__, ": desc must be [P, 8]"); + const int64_t num_p = desc->size(0); + AITER_CHECK(num_p >= 1 && num_p <= kV4PsMaxParts, + __func__, + ": num_partitions=", + num_p, + " must be in [1, ", + kV4PsMaxParts, + "]"); + AITER_CHECK(o_acc->numel() == static_cast(2 * num_p * kV4PsHeads * kV4PsDim), + __func__, + ": o_acc must be [2P, 128, 512]"); + AITER_CHECK(lse_acc->numel() == static_cast(2 * num_p * kV4PsHeads), + __func__, + ": lse_acc must be [2P, 128]"); + AITER_CHECK(cnt->numel() >= static_cast(kV4PsCntInts), + __func__, + ": cnt needs ", + kV4PsCntInts, + " entries"); + AITER_CHECK(arange->numel() >= static_cast(kV4PsArange), + __func__, + ": arange needs ", + kV4PsArange, + " entries"); + if(num_q == 0) + return; + + const HipDeviceGuard device_guard(Q->device_id); + + MlaV4PsKernelArgs args = {}; + args.ptr_O_acc = o_acc->data_ptr(); + args.ptr_out = output->data_ptr(); + args.ptr_L_acc = lse_acc->data_ptr(); + args.ptr_lse = (lse != nullptr) ? lse->data_ptr() : nullptr; + args.ptr_Q = Q->data_ptr(); + args.ptr_KV = KV->data_ptr(); + args.ptr_LTP = kv_indptr->data_ptr(); + args.ptr_LTD = kv_page_indices->data_ptr(); + args.ptr_desc = desc->data_ptr(); + args.ptr_dbg = nullptr; + args.scalar_f = 1.0f / std::sqrt(static_cast(kV4PsDim)); + args.s_gqa_ratio = kV4PsHeads; + args.s_kv_split = 1; + args.ptr_cnt = cnt->data_ptr(); + args.s_num_p = static_cast(num_p); + args.s_num_q = static_cast(num_q); + args.s_plan = kV4PsPlanCfg; + args.ptr_unused = nullptr; + args.ptr_QTP = arange->data_ptr(); + args.ptr_STP = arange->data_ptr(); + args.out_16_nosplit = 1; + args.ptr_QROPE = qrope->data_ptr(); + args.ptr_KVROPE = kvrope->data_ptr(); + args.ptr_sink = sink->data_ptr(); + + CFG* config_map = &cfg_mla_v4_asm; + // Separate cache from mla_decode_v4_asm's; keys are the distinct symbol + // names. + static SynchronizedCache impl_ptr_map; + const std::string arch_id = get_gpu_arch(); + const std::string kernelName = get_heuristic_kernel_mla_v4( + "fp8", "fp8", /*gqa=*/64, /*ps=*/1, 0, 0, /*qseqlen=*/1, /*lse=*/1, arch_id, config_map); + auto it = config_map->find(kernelName); + AITER_CHECK(it != config_map->end(), __func__, " not find kernel ", kernelName); + const char* name = it->second.knl_name.c_str(); + const char* co_name = it->second.co_name.c_str(); + AiterAsmKernel* impl_ptr = + &impl_ptr_map.get_or_create(name, [&]() { return AiterAsmKernel(name, co_name); }); + + size_t arg_size = sizeof(args); + impl_ptr->launch_kernel( + {&args, &arg_size, static_cast(2 * num_p), 1, 1, 256, 1, 1, stream}); +} diff --git a/hsa/gfx950/mla_v4/mla_a8w8_qh64_qseqlen1_gqaratio64_nm_ps.co b/hsa/gfx950/mla_v4/mla_a8w8_qh64_qseqlen1_gqaratio64_nm_ps.co new file mode 100755 index 00000000000..f162a7fa95c Binary files /dev/null and b/hsa/gfx950/mla_v4/mla_a8w8_qh64_qseqlen1_gqaratio64_nm_ps.co differ diff --git a/hsa/gfx950/mla_v4/mla_v4_asm.csv b/hsa/gfx950/mla_v4/mla_v4_asm.csv index 0427a938720..5b96c10e846 100644 --- a/hsa/gfx950/mla_v4/mla_v4_asm.csv +++ b/hsa/gfx950/mla_v4/mla_v4_asm.csv @@ -13,16 +13,21 @@ // qType : Q dtype string (fp8 / bf16) // kvType : KV dtype string (fp8 / bf16) // Gqa : gqa_ratio -// ps : persistent scheduling flag (0 = non-persistent; v4 nm only -// ships non-persistent today) +// ps : persistent scheduling flag (0 = non-persistent, served by +// mla_decode_v4_asm; 1 = persistent, served by +// mla_decode_v4_ps_asm: in-kernel split plan + merge, one launch) // qSeqLen : max_seqlen_q (post-padding) baked into the kernel binary // (4 for the qseqlen4 variant) // prefill : 0 for decode, reserved for future prefill variants // causal : 0 for non-causal, reserved -// lse : final-LSE export flag (0 = no, v4 nm always 0 today) +// lse : final-LSE export flag (0 = no; the ps=1 kernel exports it when +// given an lse buffer) // knl_name : kernel symbol in the .co // co_name : .co file path relative to this CSV's directory // +// The ps=1 row is looked up by mla_decode_v4_ps_asm only (gqa=128, qSeqLen=1 +// callers, normalized to the same (Gqa=64, qSeqLen=1) key). +// // Dispatcher keys lookup on (qType, kvType, Gqa, ps, qSeqLen, prefill, // causal, lse) — exactly the same set as v3. // v4 kernel gqa*q_seq_logical=64 tile serves three entry points (gqa=16/qSeqLen=4, gqa=64/qSeqLen=1, gqa=128/qSeqLen=1) @@ -31,4 +36,5 @@ // // Required columns last (per codegen.py convention): knl_name, co_name. qType,kvType,Gqa,ps,qSeqLen,prefill,causal,lse,knl_name,co_name -fp8,fp8,64,0,1,0,0,0,_ZN5aiter36mla_a8w8_qh64_qseqlen1_gqaratio64_nmE,mla_a8w8_qh64_qseqlen1_gqaratio64_nm.co \ No newline at end of file +fp8,fp8,64,0,1,0,0,0,_ZN5aiter36mla_a8w8_qh64_qseqlen1_gqaratio64_nmE,mla_a8w8_qh64_qseqlen1_gqaratio64_nm.co +fp8,fp8,64,1,1,0,0,1,_ZN5aiter39mla_a8w8_qh64_qseqlen1_gqaratio64_nm_psE,mla_a8w8_qh64_qseqlen1_gqaratio64_nm_ps.co diff --git a/op_tests/test_mla_v4_nm.py b/op_tests/test_mla_v4_nm.py index e8c5b759926..32e97fa909e 100644 --- a/op_tests/test_mla_v4_nm.py +++ b/op_tests/test_mla_v4_nm.py @@ -17,7 +17,7 @@ import aiter import aiter.mla # main no longer auto-imports submodules; need explicit from aiter import dtypes -from aiter.jit.utils.chip_info import get_gfx +from aiter.jit.utils.chip_info import get_cu_num, get_gfx from aiter.test_common import checkAllclose, run_perftest # --------------------------------------------------------------------------- @@ -717,6 +717,7 @@ def _gated_allclose(a, b, msg, rtol=3e-2, atol=3e-2, tol_err_ratio=0.02): f"{msg}: {err:.1%} of elements outside rtol={rtol} atol={atol} " f"(limit {tol_err_ratio:.0%})" ) + return err def _run_one_point( @@ -1026,7 +1027,9 @@ def test_v4_nm_accuracy_and_perf(): _run_one_point(batch=2, kv_seq_lens=64, q_seq_logical=1, seed=0) -def _run_varlen_point(kv_lens, gqa_ratio=128, seed=0, attn_sink=True): +def _run_varlen_point( + kv_lens, gqa_ratio=128, seed=0, attn_sink=True, split_kwargs=None +): """Accuracy at a RAGGED (per-seq variable kv_len) decode shape with auto-split (num_kv_splits=None). @@ -1097,12 +1100,11 @@ def _run_varlen_point(kv_lens, gqa_ratio=128, seed=0, attn_sink=True): kv_indptr=kv_indptr, kv_page_indices=kv_page_indices, kv_last_page_lens=kv_last_page_lens, - split_indptr=None, max_seqlen_q=1, sink=sink, sm_scale=sm_scale, out_16_nosplit=0, - num_kv_splits=None, + **(split_kwargs or {"split_indptr": None, "num_kv_splits": None}), ) # Single/multi-split output contract (see ATOM paged_decode.py:986): the # merged result lands in `output` for resolved splits > 1; the single-pass @@ -1114,11 +1116,12 @@ def _run_varlen_point(kv_lens, gqa_ratio=128, seed=0, attn_sink=True): f"\n[v4 nm varlen] gqa={gqa} kv_lens={kv_lens} total_kv={total_kv} " f"resolved_splits={resolved}" ) - _gated_allclose( + err = _gated_allclose( out_ref.float(), out_asm, msg=f"mla_v4_nm varlen [fp8_dequant_ref vs asm] kv_lens={kv_lens}", ) + return float(err or 0), out_asm @needs_gfx950 @@ -1141,6 +1144,173 @@ def test_v4_nm_varlen_ragged_kv_tail_split_guard(): _run_varlen_point([66, 50, 40, 33]) # all-short +@pytest.mark.parametrize( + "num_seqs,kv_len,expected", + [ + # One wave of workgroups: split only when KV is long enough to pay + # for the partial write + stage-2 merge. + (7, 1152, 8), + (7, 384, 5), + (35, 1152, 3), + (56, 1152, 2), + # 49..119 seqs with short KV: the occupancy-only pick chose 13-16 + # splits (5-14 waves) here and ran 3-7x slower than one split. + (49, 384, 1), + (119, 384, 1), + (126, 1152, 1), + # Past one wave, long KV still splits to even out the tail wave. + (140, 8320, 4), + ], +) +def test_v4_nm_split_planner_picks(num_seqs, kv_len, expected): + cost = aiter.mla._V4_NM_SPLIT_COST["gfx950"] + assert ( + aiter.mla._v4_nm_pick_num_kv_splits(cost, num_seqs, 2, kv_len, 256) == expected + ) + + +def _uniform_split(num_seqs, num_heads, kv_len, device="cuda"): + """mla_decode_fwd_v4_nm split kwargs: aiter's split count and its uniform + split_indptr [0, s, 2s, ...].""" + s = aiter.mla.get_mla_v4_nm_num_kv_splits(num_seqs, num_heads, kv_len) + return { + "num_kv_splits": s, + "split_indptr": torch.arange( + 0, (num_seqs + 1) * s, s, dtype=torch.int32, device=device + ), + } + + +@needs_gfx950 +@pytest.mark.parametrize( + "num_seqs,kv_len", [(7, 128), (7, 1152), (56, 1152), (14, 8320)] +) +def test_v4_nm_num_kv_splits_uses_cost_model(num_seqs, kv_len): + cost = aiter.mla._V4_NM_SPLIT_COST["gfx950"] + assert aiter.mla.get_mla_v4_nm_num_kv_splits( + num_seqs, 128, kv_len + ) == aiter.mla._v4_nm_pick_num_kv_splits(cost, num_seqs, 2, kv_len, get_cu_num()) + + +@needs_gfx950 +@pytest.mark.parametrize( + "kv_lens,kv_len", + [ + ([1152, 900, 64, 1152], 1152), # planned for the longest seq + ([8320, 200], 8320), + ([100, 60, 128], 128), # short, window-sized KV + ], +) +def test_v4_nm_split_plan_accuracy(kv_lens, kv_len): + err, out = _run_varlen_point( + kv_lens, split_kwargs=_uniform_split(len(kv_lens), 128, kv_len) + ) + assert torch.isfinite(out).all() + assert err < 0.02 + + +@needs_gfx950 +def test_v4_nm_split_plan_cudagraph_replay(): + """Plan built once into a persistent buffer, captured, then replayed with + different per-seq KV lengths written into the same kv_indptr.""" + device, gqa, n, kv_cap = "cuda", 128, 14, 1152 + torch.manual_seed(0) + sm_scale = 1.0 / (_QUANT_D**0.5) + qp, qr = _native_to_2buff_for_asm( + torch.randn(n, gqa, _QUANT_D, dtype=dtypes.bf16, device=device) + ) + kp, kr = _native_to_2buff_for_asm( + torch.randn( + n * kv_cap, + PAGE_SIZE, + NUM_KV_HEADS, + _QUANT_D, + dtype=dtypes.bf16, + device=device, + ) + ) + qr, kr = qr.contiguous(), kr.contiguous() + sink = torch.randn(gqa, dtype=torch.float32, device=device) + qo_indptr = torch.arange(n + 1, dtype=torch.int32, device=device) + kv_indptr = torch.zeros(n + 1, dtype=torch.int32, device=device) + kv_page_indices = torch.randperm(n * kv_cap, device=device).to(torch.int32) + kv_last_page_lens = torch.ones(n, dtype=torch.int32, device=device) + output = torch.empty((n, gqa, V_HEAD_DIM), dtype=dtypes.bf16, device=device) + + def set_lens(lens): + kv_indptr.copy_(torch.tensor([0] + np.cumsum(lens).tolist(), dtype=torch.int32)) + + split = _uniform_split(n, gqa, kv_cap) + assert split["num_kv_splits"] > 1 # exercise the stage-2 path under the graph + + def run(): + aiter.mla.mla_decode_fwd_v4_nm( + qp, + qr, + kp, + kr, + output, + qo_indptr, + kv_indptr, + kv_page_indices, + 1, + sink=sink, + sm_scale=sm_scale, + **split, + ) + + set_lens([kv_cap] * n) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + + for lens in ([kv_cap] * n, [1 + 83 * i for i in range(n)], [64] * n): + set_lens(lens) + graph.replay() + torch.cuda.synchronize() + ref, _ = _torch_attn_decode_fp8_dequant_ref( + qp, + qr, + kp, + kr, + qo_indptr, + kv_indptr, + kv_page_indices, + kv_last_page_lens, + sm_scale, + attn_sink=sink, + ) + assert torch.isfinite(output).all() + err = checkAllclose( + ref.float(), + output.float(), + rtol=3e-2, + atol=3e-2, + tol_err_ratio=0.02, + msg=f"mla_v4_nm split plan graph replay lens={lens[:3]}...", + ) + assert float(err or 0) < 0.02 + + +@needs_gfx950 +def test_v4_nm_one_split_per_seq_on_wide_grid(): + """Every seq on one split while the grid has 4: the kernel still writes + FP32 partials, so the stage-2 merge must not take its in-place BF16 copy + (it used to, keyed on split_indptr[-1] == num_seqs, and returned NaN).""" + kv_lens = [384, 384, 384, 384] + split_indptr = torch.arange(len(kv_lens) + 1, dtype=torch.int32, device="cuda") + err, out = _run_varlen_point( + kv_lens, split_kwargs={"num_kv_splits": 4, "split_indptr": split_indptr} + ) + assert torch.isfinite(out).all() + assert err < 0.02 + + @needs_gfx950 def test_v4_nm_ragged_short_seq_no_corrupt(): """A short seq must not be corrupted by over-allocated splits (host ragged @@ -2454,6 +2624,369 @@ def test_v4_nm_kv_tail_not_tile_multiple(): ) +# --------------------------------------------------------------------------- +# Persistent kernel: mla_decode_fwd_v4_nm_ps (gqa=128, qlen=1, one launch). +# --------------------------------------------------------------------------- +_PS_HEADS = 128 +_PS_MAX_ROWS = 32768 # mla_decode_fwd_v4_nm_ps row limit +_PS_SENTINEL = -3.0 # rows with K=0 must keep it (the kernel leaves them unwritten) + + +def _ps_inputs(kv_lens, seed=0, pool=None, device="cuda"): + """Packed q / KV pool / tables for kv_lens; row j gathers kv_lens[j] random + rows of a shared pool (page_size 1).""" + g = torch.Generator(device=device).manual_seed(seed) + n, total = len(kv_lens), int(sum(kv_lens)) + pool = pool or max(1, min(total, 20000)) + 97 + qp, qr = _native_to_2buff_for_asm( + torch.randn(n, _PS_HEADS, _QUANT_D, generator=g, device=device).to(dtypes.bf16) + ) + kp, kr = _native_to_2buff_for_asm( + torch.randn( + pool, PAGE_SIZE, NUM_KV_HEADS, _QUANT_D, generator=g, device=device + ).to(dtypes.bf16) + ) + kv_indptr = torch.zeros(n + 1, dtype=torch.int32, device=device) + kv_indptr[1:] = torch.tensor(kv_lens, device=device).cumsum(0) + kv_page_indices = torch.randint( + 0, pool, (max(total, 1),), generator=g, device=device, dtype=torch.int32 + ) + sink = torch.randn(_PS_HEADS, generator=g, device=device) * 2.0 + return { + "q_packed": qp, + "q_rope": qr.contiguous(), + "kv_packed": kp, + "kv_rope": kr.contiguous(), + "kv_indptr": kv_indptr, + "kv_page_indices": kv_page_indices, + "sink": sink, + } + + +def _ps_call(inp, ws, out=None, lse=None, return_lse=True, kv_indptr=None): + n = inp["q_packed"].size(0) + if out is None: + out = torch.full( + (n, _PS_HEADS, V_HEAD_DIM), _PS_SENTINEL, dtype=dtypes.bf16, device="cuda" + ) + if lse is None and return_lse: + lse = torch.full((n, _PS_HEADS), _PS_SENTINEL, device="cuda") + return aiter.mla.mla_decode_fwd_v4_nm_ps( + inp["q_packed"], + inp["q_rope"], + inp["kv_packed"], + inp["kv_rope"], + inp["kv_indptr"] if kv_indptr is None else kv_indptr, + inp["kv_page_indices"], + inp["sink"], + ws, + out=out, + return_lse=return_lse, + lse=lse, + ) + + +def _ps_reference(inp, kv_indptr=None): + """(out bf16, lse fp32) of the fp8-dequant reference; lse includes the sink.""" + n = inp["q_packed"].size(0) + kv_indptr = inp["kv_indptr"] if kv_indptr is None else kv_indptr + ones = torch.ones(n, dtype=torch.int32, device="cuda") + qo_indptr = torch.arange(n + 1, dtype=torch.int32, device="cuda") + out, _ = _torch_attn_decode_fp8_dequant_ref( + inp["q_packed"], + inp["q_rope"], + inp["kv_packed"], + inp["kv_rope"], + qo_indptr, + kv_indptr, + inp["kv_page_indices"], + ones, + 1.0 / (_QUANT_D**0.5), + attn_sink=inp["sink"], + ) + q = _quant_2buff_to_native(inp["q_packed"], inp["q_rope"]).float() + kv = ( + _quant_2buff_to_native(inp["kv_packed"], inp["kv_rope"]) + .float() + .view(-1, _QUANT_D) + ) + lse = torch.empty(n, _PS_HEADS, dtype=torch.float32, device="cuda") + ip = kv_indptr.tolist() + for j in range(n): + k = kv[inp["kv_page_indices"][ip[j] : ip[j + 1]].long()] + s = (q[j] @ k.T) / (_QUANT_D**0.5) + lse[j] = torch.logsumexp(torch.cat([s, inp["sink"][:, None]], 1), 1) + return out, lse + + +def _ps_check(inp, out, lse, ref=None, msg=""): + ref_out, ref_lse = ref or _ps_reference(inp) + ip = inp["kv_indptr"] + live = (ip[1:] - ip[:-1]) > 0 + assert torch.all(out[~live] == _PS_SENTINEL), "K=0 rows must be left unwritten" + if lse is not None: + assert torch.all(lse[~live] == _PS_SENTINEL), "K=0 lse must be left unwritten" + if not bool(live.any()): + return + a, b = out[live].float(), ref_out[live].float() + assert torch.isfinite(a).all() + cos = torch.nn.functional.cosine_similarity(a.flatten(1), b.flatten(1), dim=1) + rel = (a - b).flatten(1).norm(dim=1) / b.flatten(1).norm(dim=1).clamp_min(1e-12) + assert float(cos.min()) >= 0.9999, f"{msg} min row cos {float(cos.min())}" + assert float(rel.max()) <= 1e-2, f"{msg} max row rel-L2 {float(rel.max())}" + err = checkAllclose( + b, a, rtol=3e-2, atol=3e-2, tol_err_ratio=0.02, msg=f"v4 nm ps {msg}" + ) + assert float(err or 0) < 0.02 + if lse is not None: + d = float((lse[live] - ref_lse[live]).abs().max()) + assert d <= 2e-2, f"{msg} lse max abs err {d}" + + +@pytest.fixture(scope="module") +def ps_workspace(): + if not torch.cuda.is_available() or not _on_gfx950(): + pytest.skip("v4 nm shader is shipped only for gfx950") + return aiter.mla.get_mla_v4_nm_ps_workspace() + + +def _ps_ragged(n, hi, seed, zero_every=0): + g = np.random.default_rng(seed) + lens = g.integers(1, hi + 1, size=n).tolist() + if zero_every: + lens[::zero_every] = [0] * len(lens[::zero_every]) + return lens + + +@needs_gfx950 +@pytest.mark.parametrize( + "kv_lens", + [ + [1], + [64] * 7, + [129] * 256, # HCA window + one compressed token + [8320] * 4, + [0, 1, 33, 8320, 0, 500, 7], + _ps_ragged(1000, 2000, 1, zero_every=37), + _ps_ragged(3000, 300, 2), + [8320] * 3 + [5] * 600, + ], + ids=lambda v: f"N{len(v)}_Kmax{max(v)}", +) +def test_v4_nm_ps_accuracy(ps_workspace, kv_lens): + inp = _ps_inputs(kv_lens, seed=len(kv_lens)) + out, lse = _ps_call(inp, ps_workspace) + torch.cuda.synchronize() + _ps_check(inp, out, lse, msg=f"N={len(kv_lens)}") + assert int(ps_workspace.cnt.abs().sum()) == 0 + + +@needs_gfx950 +def test_v4_nm_ps_indptr_offset_and_tail(ps_workspace): + """kv_indptr longer than N+1 and not starting at 0 (N comes from q).""" + kv_lens = [300, 0, 17, 1024, 64] + inp = _ps_inputs(kv_lens, seed=3) + base, base_lse = _ps_call(inp, ps_workspace) + off = 11 + idx = inp["kv_page_indices"] + pad = torch.zeros(off, dtype=torch.int32, device="cuda") + shifted = dict(inp, kv_page_indices=torch.cat([pad, idx, pad])) + kv_indptr = torch.cat( + [inp["kv_indptr"] + off, (inp["kv_indptr"][-1:] + off + 5).repeat(3)] + ) + out, lse = _ps_call(shifted, ps_workspace, kv_indptr=kv_indptr) + torch.cuda.synchronize() + assert torch.equal(out.view(torch.int16), base.view(torch.int16)) + assert torch.equal(lse.view(torch.int32), base_lse.view(torch.int32)) + _ps_check(inp, base, base_lse) + + +@needs_gfx950 +def test_v4_nm_ps_workspace_reuse_across_batches(ps_workspace): + """One workspace across batch sizes: counters return to zero and results + match a fresh workspace bit for bit.""" + for n, seed in ((1, 0), (700, 1), (5, 2), (2048, 3), (33, 4)): + inp = _ps_inputs(_ps_ragged(n, 1500, seed, zero_every=11), seed=seed) + out, lse = _ps_call(inp, ps_workspace) + fresh, fresh_lse = _ps_call(inp, aiter.mla.get_mla_v4_nm_ps_workspace()) + torch.cuda.synchronize() + assert int(ps_workspace.cnt.abs().sum()) == 0 + assert torch.equal(out.view(torch.int16), fresh.view(torch.int16)) + assert torch.equal(lse.view(torch.int32), fresh_lse.view(torch.int32)) + + +@needs_gfx950 +def test_v4_nm_ps_partitions(ps_workspace): + """Non-default partition counts agree with the default one.""" + inp = _ps_inputs(_ps_ragged(300, 3000, 5, zero_every=7), seed=5) + ref = _ps_reference(inp) + for p in (1, 37, 1024): + ws = aiter.mla.get_mla_v4_nm_ps_workspace(num_partitions=p) + out, lse = _ps_call(inp, ws) + torch.cuda.synchronize() + _ps_check(inp, out, lse, ref=ref, msg=f"P={p}") + assert int(ws.cnt.abs().sum()) == 0 + + +@needs_gfx950 +@pytest.mark.filterwarnings("ignore:The CUDA Graph is empty") +def test_v4_nm_ps_cudagraph_replay(ps_workspace): + """Captured once, replayed with changing KV lengths (incl. K=0 rows).""" + n, kv_cap = 96, 2048 + inp = _ps_inputs([kv_cap] * n, seed=6) + out = torch.empty(n, _PS_HEADS, V_HEAD_DIM, dtype=dtypes.bf16, device="cuda") + lse = torch.empty(n, _PS_HEADS, dtype=torch.float32, device="cuda") + + def run(): + _ps_call(inp, ps_workspace, out=out, lse=lse) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + _ps_call(inp, ps_workspace) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + with pytest.raises(RuntimeError), torch.cuda.graph(torch.cuda.CUDAGraph()): + aiter.mla.get_mla_v4_nm_ps_workspace() + + g = np.random.default_rng(7) + for lens in ( + [kv_cap] * n, + g.integers(0, kv_cap + 1, size=n).tolist(), + [1] * n, + [0] * (n // 2) + [kv_cap] * (n - n // 2), + ): + # rows read a dense prefix of the capacity-sized index table + inp["kv_indptr"][1:].copy_(torch.tensor(lens, device="cuda").cumsum(0)) + out.fill_(_PS_SENTINEL) + lse.fill_(_PS_SENTINEL) + graph.replay() + torch.cuda.synchronize() + _ps_check(inp, out, lse, msg=f"graph lens={lens[:3]}") + eager, eager_lse = _ps_call(inp, ps_workspace) + torch.cuda.synchronize() + assert torch.equal(out.view(torch.int16), eager.view(torch.int16)) + assert torch.equal(lse.view(torch.int32), eager_lse.view(torch.int32)) + assert int(ps_workspace.cnt.abs().sum()) == 0 + + +@needs_gfx950 +@pytest.mark.parametrize( + "kv_lens", [[1152, 900, 64, 1152], [8320, 200], [100, 60, 128] * 50] +) +def test_v4_nm_ps_matches_split_kernel(ps_workspace, kv_lens): + """Agrees with mla_decode_fwd_v4_nm (non-persistent kernel + planned split).""" + inp = _ps_inputs(kv_lens, seed=8) + n = len(kv_lens) + out = _ps_call(inp, ps_workspace, return_lse=False) + ref = torch.empty(n, _PS_HEADS, V_HEAD_DIM, dtype=dtypes.bf16, device="cuda") + split = _uniform_split(n, _PS_HEADS, max(kv_lens)) + logits, _ = aiter.mla.mla_decode_fwd_v4_nm( + inp["q_packed"], + inp["q_rope"], + inp["kv_packed"], + inp["kv_rope"], + ref, + torch.arange(n + 1, dtype=torch.int32, device="cuda"), + inp["kv_indptr"], + inp["kv_page_indices"], + 1, + sink=inp["sink"], + **split, + ) + if split["num_kv_splits"] == 1: + ref = logits[:, 0] + torch.cuda.synchronize() + assert _cos_diff(out, ref) < 1e-5 + err = checkAllclose( + ref.float(), + out.float(), + rtol=3e-2, + atol=3e-2, + tol_err_ratio=0.02, + msg="v4 nm ps vs split kernel", + ) + assert float(err or 0) < 0.02 + + +@needs_gfx950 +def test_v4_nm_ps_rejects_bad_args(ps_workspace): + inp = _ps_inputs([64, 64], seed=9) + with pytest.raises(RuntimeError): # 64 heads + aiter.mla.mla_decode_fwd_v4_nm_ps( + inp["q_packed"][:, :64].contiguous(), + inp["q_rope"][:, :64].contiguous(), + inp["kv_packed"], + inp["kv_rope"], + inp["kv_indptr"], + inp["kv_page_indices"], + inp["sink"], + ps_workspace, + ) + with pytest.raises(RuntimeError): # kv_indptr shorter than N+1 + _ps_call(inp, ps_workspace, kv_indptr=inp["kv_indptr"][:2]) + with pytest.raises(ValueError): + aiter.mla.get_mla_v4_nm_ps_workspace(num_partitions=1025) + + +def _ps_oob_worker_main(): + """Subprocess body of test_v4_nm_ps_oob_guardpage. Every buffer the + persistent kernel reads or writes (q, qrope, KV pools, kv_indptr with + exactly N+1 entries, kv_page_indices, sink, out, lse) is placed at the + very end of its own allocation, rounded to 2 MiB, under the non-caching + allocator, so an access even one element past any of them crosses into an + unmapped page and faults the GPU. (A one-entry over-read of + kv_page_indices faults with this layout; a plain clone of the same tensor + leaves up to 2 MiB of slack and does not.)""" + granule = 2 << 20 + + def tail(t): + nbytes = t.numel() * t.element_size() + total = -(-nbytes // granule) * granule + buf = torch.empty(total, dtype=torch.uint8, device=t.device) + v = buf[total - nbytes :].view(t.dtype).view(t.shape) + v.copy_(t) + return v + + for n, hi, zero_every, return_lse in ( + (8192, 700, 13, True), + (_PS_MAX_ROWS, 64, 0, True), # largest N: out ends at the 4 GiB offset limit + (3000, 4000, 0, False), + ): + inp = _ps_inputs(_ps_ragged(n, hi, n, zero_every=zero_every), seed=n) + inp = {k: tail(v) for k, v in inp.items()} + out = tail( + torch.empty(n, _PS_HEADS, V_HEAD_DIM, dtype=dtypes.bf16, device="cuda") + ) + lse = tail(torch.empty(n, _PS_HEADS, device="cuda")) if return_lse else None + ws = aiter.mla.get_mla_v4_nm_ps_workspace() + _ps_call(inp, ws, out=out, lse=lse, return_lse=return_lse) + torch.cuda.synchronize() + assert int(ws.cnt.abs().sum()) == 0 + del inp, out, lse, ws + + +@needs_gfx950 +def test_v4_nm_ps_oob_guardpage(): + """Guard-page out-of-bounds check for mla_decode_fwd_v4_nm_ps. + + The persistent kernel runs the head 64..127 path of the ASM (the one the + gqa=128 Q_rope over-read regression lived on) and adds its own reads: + the kv_indptr scans of the in-kernel planner, the split partials and the + merge. An over-read there can be numerically silent, so this runs large + batches in a subprocess with every input and output in a tight + allocation and asserts a clean exit, like + test_v4_nm_gqa128_qrope_oob_guardpage. + """ + _run_oob_guardpage_worker( + ["ps"], + "v4 nm ps OOB detected: the persistent kernel accessed memory past " + "one of its input / output buffers.", + ) + + if __name__ == "__main__": import argparse import itertools @@ -2472,6 +3005,10 @@ def test_v4_nm_kv_tail_not_tile_multiple(): # non-caching allocator so any tg_idx=1 Q_rope over-read crosses into an # unmapped page and faults the GPU (killing THIS process). A clean exit 0 # means no OOB. See the test's docstring for the full rationale. + if sys.argv[1:3] == ["--oob-worker", "ps"]: + _ps_oob_worker_main() + print("[v4 nm][oob-worker] COMPLETED no fault") + sys.exit(0) if len(sys.argv) >= 2 and sys.argv[1] == "--oob-worker": _gqa = int(sys.argv[2]) if len(sys.argv) >= 3 else 128 _msq = int(sys.argv[3]) if len(sys.argv) >= 4 else 1