diff --git a/aiter/jit/optCompilerConfig.json b/aiter/jit/optCompilerConfig.json index 29a18801eb5..f86937bfc22 100644 --- a/aiter/jit/optCompilerConfig.json +++ b/aiter/jit/optCompilerConfig.json @@ -53,6 +53,17 @@ "verbose": "False", "blob_gen_cmd": "f'{AITER_META_DIR}/hsa/codegen.py -m pa --output_dir {{}}'" }, + "module_mtp_verify_attn_asm": { + "srcs": [ + "f'{AITER_CSRC_DIR}/py_itfs_cu/asm_mtp_verify_attn.cu'" + ], + "flags_extra_cc": [], + "flags_extra_hip": [], + "extra_ldflags": "None", + "extra_include": [], + "verbose": "False", + "blob_gen_cmd": "''" + }, "module_pa_metadata": { "srcs": [ "f'{AITER_CSRC_DIR}/pybind/pa_metadata_pybind.cu'", diff --git a/aiter/ops/attention.py b/aiter/ops/attention.py index 47365bb999a..58a9e0fac63 100644 --- a/aiter/ops/attention.py +++ b/aiter/ops/attention.py @@ -1864,3 +1864,101 @@ def mla_decode_fwd_opus_stage1( kv_scale: torch.Tensor | None = None, # float[1] per-tensor descale causal: bool = True, # mask across the max_seqlen_q query tokens ) -> None: ... + + +@compile_ops( + "module_mtp_verify_attn_asm", fc_name="mtp_verify_attn_fwd", ffi_type="ctypes" +) +def _mtp_verify_attn_fwd( + Q: torch.Tensor, + K: torch.Tensor, + V: torch.Tensor, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + cu_seqlens_q: torch.Tensor, + k_descale: torch.Tensor, + v_descale: torch.Tensor, + segm_out: torch.Tensor, + segm_max: torch.Tensor, + segm_expsum: torch.Tensor, + out: torch.Tensor, + softmax_scale: float, +) -> None: ... + + +def mtp_verify_attn_num_segments(num_seqs: int, num_kv_heads: int) -> int: + """Split-KV segments per sequence so the grid (segments, seqs, kv heads) + fills the GPU in one round of 1-work-group-per-CU. Any value in [1, 64] is + valid; empty segments are written neutrally by the kernel.""" + num_cus = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + segs = max(1, num_cus // max(1, num_seqs * num_kv_heads)) + segs = 1 << (segs.bit_length() - 1) + return max(1, min(64, segs)) + + +def mtp_verify_attn_fwd_asm( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + cu_seqlens_q: torch.Tensor, + k_descale: torch.Tensor, + v_descale: torch.Tensor, + softmax_scale: float, + num_segments: int | None = None, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """MTP-verify attention on gfx950 (hand-written assembly). + + Speculative-decoding verify step: every sequence has exactly 4 query tokens + (q_len 4, causal against the last 4 KV positions) attending to its paged + fp8 KV cache. Shapes: + q: [num_seqs * 4, num_q_heads, 256] bf16 (token-major, cu_seqlens_q order) + k_cache/v_cache: [num_pages, 16, num_kv_heads, 256] fp8 e4m3 (NHD, page size 16) + block_tables: [num_seqs, max_pages] int32 + seq_lens: [num_seqs] int64, KV length including the 4 query tokens + cu_seqlens_q: [num_seqs + 1] int32 + k_descale / v_descale: [1] fp32 device tensors (per-tensor KV scales) + GQA ratio (num_q_heads / num_kv_heads) must be 16 or 8. Returns bf16 + [num_seqs * 4, num_q_heads, 256]. + """ + num_tokens, num_q_heads, head_size = q.shape + num_seqs = seq_lens.shape[0] + num_kv_heads = k_cache.shape[2] + if num_segments is None: + num_segments = mtp_verify_attn_num_segments(num_seqs, num_kv_heads) + segm_out = torch.empty( + num_tokens, + num_q_heads, + num_segments, + head_size, + dtype=torch.float32, + device=q.device, + ) + segm_max = torch.empty( + num_tokens, num_q_heads, num_segments, dtype=torch.float32, device=q.device + ) + segm_expsum = torch.empty_like(segm_max) + if out is None: + out = torch.empty( + num_tokens, num_q_heads, head_size, dtype=q.dtype, device=q.device + ) + _mtp_verify_attn_fwd( + q, + k_cache, + v_cache, + block_tables, + seq_lens, + cu_seqlens_q, + k_descale, + v_descale, + segm_out, + segm_max, + segm_expsum, + out, + float(softmax_scale), + ) + return out diff --git a/csrc/py_itfs_cu/asm_mtp_verify_attn.cu b/csrc/py_itfs_cu/asm_mtp_verify_attn.cu new file mode 100644 index 00000000000..57242ea6337 --- /dev/null +++ b/csrc/py_itfs_cu/asm_mtp_verify_attn.cu @@ -0,0 +1,206 @@ +// SPDX-License-Identifier: MIT +// Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +// +// MTP-verify attention for gfx950: the speculative-decoding verify step +// (uniform q_len = 4 query tokens per sequence) over an NHD page-16 fp8 KV +// cache with head_dim 256, GQA ratio 16 or 8. Split-KV: the main kernel +// writes per-segment partial outputs / max / expsum, the reduce kernel merges +// them into the final bf16 output. +#include "aiter_tensor.h" +#include +#include + +namespace { + +// Must match the kernarg layout of hsa/gfx950/mtp_verify_attn/src/vattn3_core.s +struct __attribute__((packed)) VattnKernelArgs +{ + void* ptr_K; // [num_pages, 16, num_kv_heads, 256] fp8 + void* ptr_V; + void* ptr_Q; // [num_tokens, num_q_heads, 256] bf16 + void* ptr_BT; // block tables [num_seqs, max_pages] int32 + void* ptr_SL; // kv seq lens [num_seqs] int64 + void* ptr_CU; // cu_seqlens_q [num_seqs + 1] int32 + void* ptr_SO; // segm out [num_tokens, num_q_heads, segs, 256] fp32 + void* ptr_SM; // segm max [num_tokens, num_q_heads, segs] fp32 + void* ptr_SE; // segm expsum [num_tokens, num_q_heads, segs] fp32 + int32_t num_segments; + int32_t bt_stride0; // block table row stride (elements) + int32_t q_stride0; // Q token stride (elements) + int32_t num_kv_heads; + int32_t num_q_heads; + float scale_log2; // softmax_scale * log2(e); kernel multiplies by *k_descale + float v_scale; // 1.0; kernel multiplies by *v_descale + uint32_t magic_m; // magic division by (num_segments * 16) + int32_t magic_sh; + int32_t _pad; + void* ptr_KD; // k_descale [1] fp32 (device) + void* ptr_VD; // v_descale [1] fp32 (device) +}; +static_assert(sizeof(VattnKernelArgs) == 128, "kernarg layout must match vattn3_core.s"); + +// Must match hsa/gfx950/mtp_verify_attn/src/vred.s +struct __attribute__((packed)) VredKernelArgs +{ + void* ptr_O; // [num_tokens, num_q_heads, 256] bf16 + void* ptr_SO; + void* ptr_SM; + void* ptr_SE; + uint32_t num_q_heads; + uint32_t num_segments; + uint32_t o_stride0; // elements + uint32_t o_stride1; // elements + uint32_t magic_m; // magic division by num_q_heads + uint32_t magic_sh; +}; +static_assert(sizeof(VredKernelArgs) == 56, "kernarg layout must match vred.s"); + +// (M, sh) such that ((n * M) >> 32) >> sh == n / d for all n < 2^31. +void magic_u32(uint32_t d, uint32_t& M, uint32_t& sh) +{ + if((d & (d - 1)) == 0) + { + M = 1u << 31; + sh = 0; + for(uint32_t x = d; x > 1; x >>= 1) + sh++; + sh -= 1; + return; + } + uint32_t s = 0; + for(uint32_t x = d; x > 1; x >>= 1) + s++; + uint64_t m = ((1ull << (32 + s)) + d - 1) / d; + AITER_CHECK(m < (1ull << 32), "magic_u32: divisor ", d, " out of range"); + M = static_cast(m); + sh = s; +} + +const float k_log2e = 1.4426950408889634f; + +} // namespace + +AITER_C_ITFS void mtp_verify_attn_fwd( + aiter_tensor_t* Q, // [num_tokens, num_q_heads, 256] bf16, num_tokens = 4 * num_seqs + aiter_tensor_t* K, // [num_pages, 16, num_kv_heads, 256] fp8 e4m3 (NHD page 16) + aiter_tensor_t* V, // same layout as K + aiter_tensor_t* block_tables, // [num_seqs, max_pages] int32 + aiter_tensor_t* seq_lens, // [num_seqs] int64, kv length including the 4 query tokens + aiter_tensor_t* cu_seqlens_q, // [num_seqs + 1] int32 + aiter_tensor_t* k_descale, // [1] fp32, device + aiter_tensor_t* v_descale, // [1] fp32, device + aiter_tensor_t* segm_out, // [num_tokens, num_q_heads, num_segments, 256] fp32 workspace + aiter_tensor_t* segm_max, // [num_tokens, num_q_heads, num_segments] fp32 workspace + aiter_tensor_t* segm_expsum, // [num_tokens, num_q_heads, num_segments] fp32 workspace + aiter_tensor_t* out, // [num_tokens, num_q_heads, 256] bf16 + float softmax_scale, + hipStream_t stream) +{ + const int num_tokens = static_cast(Q->size(0)); + const int num_q_heads = static_cast(Q->size(1)); + const int head_size = static_cast(Q->size(2)); + const int page_size = static_cast(K->size(1)); + const int num_kv_heads = static_cast(K->size(2)); + const int num_seqs = static_cast(seq_lens->size(0)); + const int num_segments = static_cast(segm_max->size(2)); + + AITER_CHECK(Q->dtype() == AITER_DTYPE_bf16, __func__, ": Q must be bf16"); + AITER_CHECK(K->dtype() == AITER_DTYPE_fp8 && V->dtype() == AITER_DTYPE_fp8, + __func__, ": K/V must be fp8 e4m3"); + AITER_CHECK(head_size == 256 && K->size(3) == 256, __func__, ": head_dim must be 256"); + AITER_CHECK(page_size == 16, __func__, ": page size must be 16"); + AITER_CHECK(num_tokens == 4 * num_seqs, __func__, + ": uniform q_len of 4 required (num_tokens = 4 * num_seqs)"); + AITER_CHECK(num_q_heads % num_kv_heads == 0, __func__, ": bad GQA layout"); + const int gqa = num_q_heads / num_kv_heads; + AITER_CHECK(gqa == 16 || gqa == 8, __func__, ": GQA ratio must be 16 or 8, got ", gqa); + AITER_CHECK(Q->stride(2) == 1 && Q->stride(1) == head_size, __func__, ": Q heads must be contiguous"); + AITER_CHECK(K->stride(3) == 1 && K->stride(2) == head_size && + K->stride(1) == num_kv_heads * head_size && + K->stride(0) == page_size * num_kv_heads * head_size, + __func__, ": K must be a contiguous NHD page-16 cache"); + AITER_CHECK(V->stride(3) == 1 && V->stride(2) == head_size && + V->stride(1) == num_kv_heads * head_size && + V->stride(0) == page_size * num_kv_heads * head_size, + __func__, ": V must be a contiguous NHD page-16 cache"); + AITER_CHECK(seq_lens->element_size() == 8, __func__, ": seq_lens must be int64"); + AITER_CHECK(cu_seqlens_q->element_size() == 4 && block_tables->element_size() == 4, + __func__, ": cu_seqlens_q / block_tables must be int32"); + AITER_CHECK(k_descale->element_size() == 4 && v_descale->element_size() == 4, + __func__, ": k/v descale must be fp32"); + AITER_CHECK(segm_out->size(0) == num_tokens && segm_out->size(1) == num_q_heads && + segm_out->size(2) == num_segments && segm_out->size(3) == head_size && + segm_out->is_contiguous(), + __func__, ": segm_out must be [num_tokens, num_q_heads, num_segments, 256] contiguous"); + AITER_CHECK(segm_max->is_contiguous() && segm_expsum->is_contiguous(), __func__, + ": segm_max / segm_expsum must be contiguous"); + AITER_CHECK(num_segments >= 1 && num_segments <= 64, __func__, ": num_segments must be in [1, 64]"); + AITER_CHECK(out->dtype() == AITER_DTYPE_bf16 && out->stride(2) == 1, __func__, + ": out must be bf16 with a contiguous head dim"); + + const HipDeviceGuard device_guard(Q->device_id); + + VattnKernelArgs args = {}; + size_t arg_size = sizeof(args); + args.ptr_K = K->data_ptr(); + args.ptr_V = V->data_ptr(); + args.ptr_Q = Q->data_ptr(); + args.ptr_BT = block_tables->data_ptr(); + args.ptr_SL = seq_lens->data_ptr(); + args.ptr_CU = cu_seqlens_q->data_ptr(); + args.ptr_SO = segm_out->data_ptr(); + args.ptr_SM = segm_max->data_ptr(); + args.ptr_SE = segm_expsum->data_ptr(); + args.num_segments = num_segments; + args.bt_stride0 = static_cast(block_tables->stride(0)); + args.q_stride0 = static_cast(Q->stride(0)); + args.num_kv_heads = num_kv_heads; + args.num_q_heads = num_q_heads; + args.scale_log2 = softmax_scale * k_log2e; + args.v_scale = 1.0f; + uint32_t m, sh; + magic_u32(static_cast(num_segments) * 16u, m, sh); + args.magic_m = m; + args.magic_sh = static_cast(sh); + args._pad = 0; + args.ptr_KD = k_descale->data_ptr(); + args.ptr_VD = v_descale->data_ptr(); + + static AiterAsmKernel impl_gqa16("vattn_asm", "/mtp_verify_attn/vattn_hd256_fp8_gqa16.co"); + static AiterAsmKernel impl_gqa8("vattn_asm", "/mtp_verify_attn/vattn_hd256_fp8_gqa8.co"); + static AiterAsmKernel impl_reduce("vred_asm", "/mtp_verify_attn/vattn_reduce.co"); + + AiterAsmKernel* impl = (gqa == 16) ? &impl_gqa16 : &impl_gqa8; + impl->launch_kernel({&args, + &arg_size, + num_segments, // gdx: KV segments + num_seqs, // gdy + num_kv_heads, // gdz + 512, // bdx: 8 waves + 1, + 1, + stream}); + + VredKernelArgs rargs = {}; + size_t rarg_size = sizeof(rargs); + rargs.ptr_O = out->data_ptr(); + rargs.ptr_SO = segm_out->data_ptr(); + rargs.ptr_SM = segm_max->data_ptr(); + rargs.ptr_SE = segm_expsum->data_ptr(); + rargs.num_q_heads = static_cast(num_q_heads); + rargs.num_segments = static_cast(num_segments); + rargs.o_stride0 = static_cast(out->stride(0)); + rargs.o_stride1 = static_cast(out->stride(1)); + magic_u32(static_cast(num_q_heads), m, sh); + rargs.magic_m = m; + rargs.magic_sh = sh; + impl_reduce.launch_kernel({&rargs, + &rarg_size, + num_tokens * num_q_heads, // one wave per (token, head, 64-dim quarter) + 1, + 1, + 256, + 1, + 1, + stream}); +} diff --git a/hsa/gfx950/mtp_verify_attn/vattn_hd256_fp8_gqa16.co b/hsa/gfx950/mtp_verify_attn/vattn_hd256_fp8_gqa16.co new file mode 100755 index 00000000000..959aca2beda Binary files /dev/null and b/hsa/gfx950/mtp_verify_attn/vattn_hd256_fp8_gqa16.co differ diff --git a/hsa/gfx950/mtp_verify_attn/vattn_hd256_fp8_gqa8.co b/hsa/gfx950/mtp_verify_attn/vattn_hd256_fp8_gqa8.co new file mode 100755 index 00000000000..604e8045f5c Binary files /dev/null and b/hsa/gfx950/mtp_verify_attn/vattn_hd256_fp8_gqa8.co differ diff --git a/hsa/gfx950/mtp_verify_attn/vattn_reduce.co b/hsa/gfx950/mtp_verify_attn/vattn_reduce.co new file mode 100755 index 00000000000..815b9bfcc83 Binary files /dev/null and b/hsa/gfx950/mtp_verify_attn/vattn_reduce.co differ diff --git a/op_tests/test_mtp_verify_attn_asm.py b/op_tests/test_mtp_verify_attn_asm.py new file mode 100644 index 00000000000..dc67525c230 --- /dev/null +++ b/op_tests/test_mtp_verify_attn_asm.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +# +# Correctness + performance test for the gfx950 assembly MTP-verify attention +# (aiter.mtp_verify_attn_fwd_asm): 4 query tokens per sequence over an NHD +# page-16 fp8 KV cache, head_dim 256, GQA ratio 16 (TP2 shape) or 8 (TP4). +import argparse +import math + +import pandas as pd +import torch + +import aiter +from aiter import dtypes +from aiter.jit.utils.chip_info import get_gfx +from aiter.test_common import checkAllclose, run_perftest + +torch.set_default_device("cuda") +SUPPORTED_GFX = ["gfx950"] +QLEN, HD, PAGE = 4, 256, 16 + + +def make_inputs( + ctx, + num_seqs, + num_q_heads, + num_kv_heads, + k_descale=1.0, + v_descale=1.0, + replicas=1, + seed=0, +): + """Random paged KV cache with a shuffled page table (pages of one sequence + are scattered over the pool, as a real allocator would leave them).""" + torch.manual_seed(seed) + kvlen = ctx + QLEN + npages = (kvlen + PAGE - 1) // PAGE + perm = torch.randperm(num_seqs * npages) + bt = perm.to(torch.int32).view(num_seqs, npages) + caches = [] + for _ in range(replicas): + k = (torch.randn(num_seqs * npages, PAGE, num_kv_heads, HD) / 4).to(dtypes.fp8) + v = (torch.randn(num_seqs * npages, PAGE, num_kv_heads, HD) / 4).to(dtypes.fp8) + k2, v2 = torch.empty_like(k), torch.empty_like(v) + k2[perm], v2[perm] = k, v + caches.append((k2, v2)) + q = (torch.randn(num_seqs * QLEN, num_q_heads, HD) / 4).to(dtypes.bf16) + cu_q = torch.arange(0, (num_seqs + 1) * QLEN, QLEN, dtype=torch.int32) + seq_lens = torch.full((num_seqs,), kvlen, dtype=torch.int64) + kd = torch.full((1,), k_descale, dtype=torch.float32) + vd = torch.full((1,), v_descale, dtype=torch.float32) + return caches, bt, q, cu_q, seq_lens, kd, vd, kvlen + + +def torch_ref(k2, v2, bt, q, seq_lens, kd, vd, num_q_heads, num_kv_heads): + num_seqs = bt.shape[0] + gqa = num_q_heads // num_kv_heads + outs = [] + for s in range(num_seqs): + kvlen = int(seq_lens[s]) + pages = bt[s].long() + kk = k2[pages].reshape(-1, num_kv_heads, HD)[:kvlen].float() * kd + vv = v2[pages].reshape(-1, num_kv_heads, HD)[:kvlen].float() * vd + qq = q[s * QLEN : (s + 1) * QLEN].float() + o = torch.empty(QLEN, num_q_heads, HD) + for t in range(QLEN): + L = kvlen - QLEN + t + 1 # causal: query t sees the first L keys + for h in range(num_q_heads): + kvh = h // gqa + sc = (qq[t, h] @ kk[:L, kvh].T) / math.sqrt(HD) + o[t, h] = torch.softmax(sc, dim=-1) @ vv[:L, kvh] + outs.append(o) + return torch.cat(outs).to(dtypes.bf16) + + +def test_correctness(ctx, num_seqs, num_q_heads, num_kv_heads, k_descale, v_descale): + caches, bt, q, cu_q, seq_lens, kd, vd, _kvlen = make_inputs( + ctx, num_seqs, num_q_heads, num_kv_heads, k_descale, v_descale + ) + k2, v2 = caches[0] + out = aiter.mtp_verify_attn_fwd_asm( + q, k2, v2, bt, seq_lens, cu_q, kd, vd, 1.0 / math.sqrt(HD) + ) + ref = torch_ref(k2, v2, bt, q, seq_lens, kd, vd, num_q_heads, num_kv_heads) + tag = f"ctx={ctx} seqs={num_seqs} q_heads={num_q_heads} kv_heads={num_kv_heads} kd={k_descale} vd={v_descale}" + # Q and P are fp8-quantized inside the kernel: error is ~0.13 of the output std. + checkAllclose(ref.float(), out.float(), atol=4e-3, rtol=0.2, msg=tag) + return (out.float() - ref.float()).abs().max().item() / ref.float().std().item() + + +def test_perf(ctx, num_seqs, num_q_heads, num_kv_heads, replicas=8): + caches, bt, q, cu_q, seq_lens, kd, vd, kvlen = make_inputs( + ctx, num_seqs, num_q_heads, num_kv_heads, replicas=replicas + ) + scale = 1.0 / math.sqrt(HD) + num_segments = aiter.mtp_verify_attn_num_segments(num_seqs, num_kv_heads) + state = {"i": 0} + + def run(): + k2, v2 = caches[state["i"] % replicas] # rotate replicas so the KV read is cold + state["i"] += 1 + return aiter.mtp_verify_attn_fwd_asm( + q, k2, v2, bt, seq_lens, cu_q, kd, vd, scale, num_segments=num_segments + ) + + _, us = run_perftest(run, num_iters=64, num_warmup=4) + kv_bytes = num_seqs * kvlen * num_kv_heads * HD * 2 + return us, kv_bytes / us / 1e6, num_segments + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--perf", action="store_true", help="also run the bandwidth benchmark" + ) + args = parser.parse_args() + if get_gfx() not in SUPPORTED_GFX: + print( + f"skip: mtp_verify_attn asm kernel is only shipped for {SUPPORTED_GFX}, got {get_gfx()}" + ) + return + rows = [] + for ctx, seqs, hq, hkv, kd, vd in [ + (1000, 2, 16, 1, 1.0, 1.0), # TP2 shape, GQA 16 + (1000, 2, 16, 1, 0.5, 2.0), # per-tensor KV descales + (1000, 2, 8, 1, 1.0, 1.0), # TP4 shape, GQA 8 + (777, 2, 32, 2, 1.0, 1.0), # 2 KV heads, ragged page count + (8000, 3, 16, 1, 1.0, 1.0), # empty tail segments + ]: + err = test_correctness(ctx, seqs, hq, hkv, kd, vd) + rows.append( + { + "ctx": ctx, + "seqs": seqs, + "q_heads": hq, + "kv_heads": hkv, + "k_descale": kd, + "v_descale": vd, + "max_err_over_std": round(err, 3), + } + ) + print(pd.DataFrame(rows).to_markdown(index=False)) + if args.perf: + rows = [] + for ctx, seqs, hq, hkv in [ + (35000, 16, 16, 1), + (70000, 16, 16, 1), + (35000, 16, 8, 1), + (35000, 4, 16, 1), + ]: + us, tbps, segs = test_perf(ctx, seqs, hq, hkv) + rows.append( + { + "ctx": ctx, + "seqs": seqs, + "q_heads": hq, + "kv_heads": hkv, + "segments": segs, + "us": round(us, 1), + "kv_TBps": round(tbps, 2), + } + ) + print(pd.DataFrame(rows).to_markdown(index=False)) + + +if __name__ == "__main__": + main()