diff --git a/csrc/cpp_itfs/mha_fwd.cu b/csrc/cpp_itfs/mha_fwd.cu index be7f94b8486..393f43f0f48 100644 --- a/csrc/cpp_itfs/mha_fwd.cu +++ b/csrc/cpp_itfs/mha_fwd.cu @@ -9,6 +9,9 @@ namespace aiter { #if FAV3_ON +// gfx950 group: grid.z = Q tiles, max 65535 per launch (kernel adds s_q_tile_base). +static constexpr int kFmhaV3GroupMaxGridZ = 65535; + int get_cfg_mask_type(const mha_fwd_args& a) { if(a.mask_type == 0) @@ -111,6 +114,7 @@ void init_fmha_fwd_v3_args(fmha_fwd_v3_args& args, args.s_descale_k_Hs = 0; args.s_descale_v_Bs = 0; args.s_descale_v_Hs = 0; + args.s_q_tile_base = 0; int in_bpe = 2; int out_bpe = 2; @@ -258,9 +262,37 @@ float fmha_fwd_v3(mha_fwd_args a, const ck_tile::stream_config& s) int bdx = (a.hdim_q == 192 && a.hdim_v == 128) ? 256 : 512; auto [gdx, gdy, gdz] = get_grid_dim(a, cfg.ts_qo, arch_id); + // gfx950 varlen group only: split grid.z (Q tiles); gfx942 / batch unchanged. + if(a.is_group_mode && arch_id == "gfx950" && gdz > kFmhaV3GroupMaxGridZ) + { + float ret = 0.f; + for(int base = 0; base < gdz; base += kFmhaV3GroupMaxGridZ) + { + const int gdz_chunk = + (gdz - base > kFmhaV3GroupMaxGridZ) ? kFmhaV3GroupMaxGridZ : (gdz - base); + args.s_q_tile_base = static_cast(base); + ret = ck_tile::launch_kernel(s, [=](const ck_tile::stream_config& s_) mutable { + void* args_ptr = &args; + size_t* arg_size_ptr = &arg_size; + impl_ptr->launch_kernel({args_ptr, + arg_size_ptr, + gdx, + gdy, + gdz_chunk, + bdx, + 1, + 1, + s_.stream_id_}); + }); + if(ret < 0.f) + { + return ret; + } + } + return ret; + } + return ck_tile::launch_kernel(s, [=](const ck_tile::stream_config& s_) mutable { - // Explicit assignment forces evaluation order and prevents compiler from - // reordering operations that could lead to accessing uninitialized args void* args_ptr = &args; size_t* arg_size_ptr = &arg_size; impl_ptr->launch_kernel({args_ptr, arg_size_ptr, gdx, gdy, gdz, bdx, 1, 1, s_.stream_id_}); diff --git a/csrc/include/mha_fwd.h b/csrc/include/mha_fwd.h index d191ab652d1..0d22f2d867b 100644 --- a/csrc/include/mha_fwd.h +++ b/csrc/include/mha_fwd.h @@ -355,6 +355,9 @@ struct __attribute__((packed)) fmha_fwd_v3_args p3 _p39; unsigned int s_descale_v_Hs; p3 _p40; + // Global Q-tile index of this launch's blockIdx.z=0 (host gdz chunk for group mode). + unsigned int s_q_tile_base; + p3 _p41; }; __attribute__((visibility("default"))) float mha_fwd(mha_fwd_args args, diff --git a/hsa/gfx950/fmha_v3_fwd/fwd_hd192_hd128_bf16.co b/hsa/gfx950/fmha_v3_fwd/fwd_hd192_hd128_bf16.co index 10ba4080beb..6fe9c768b4b 100755 Binary files a/hsa/gfx950/fmha_v3_fwd/fwd_hd192_hd128_bf16.co and b/hsa/gfx950/fmha_v3_fwd/fwd_hd192_hd128_bf16.co differ diff --git a/hsa/gfx950/fmha_v3_fwd/fwd_hd192_hd128_bf16_causal.co b/hsa/gfx950/fmha_v3_fwd/fwd_hd192_hd128_bf16_causal.co index 00400392dc0..bb1e355b7cf 100755 Binary files a/hsa/gfx950/fmha_v3_fwd/fwd_hd192_hd128_bf16_causal.co and b/hsa/gfx950/fmha_v3_fwd/fwd_hd192_hd128_bf16_causal.co differ diff --git a/hsa/gfx950/fmha_v3_fwd/fwd_hd192_hd128_bf16_causal_group.co b/hsa/gfx950/fmha_v3_fwd/fwd_hd192_hd128_bf16_causal_group.co index 8f730587054..31615c79a4b 100755 Binary files a/hsa/gfx950/fmha_v3_fwd/fwd_hd192_hd128_bf16_causal_group.co and b/hsa/gfx950/fmha_v3_fwd/fwd_hd192_hd128_bf16_causal_group.co differ diff --git a/hsa/gfx950/fmha_v3_fwd/fwd_hd192_hd128_bf16_group.co b/hsa/gfx950/fmha_v3_fwd/fwd_hd192_hd128_bf16_group.co index 100f841ac6a..a497c039d4e 100755 Binary files a/hsa/gfx950/fmha_v3_fwd/fwd_hd192_hd128_bf16_group.co and b/hsa/gfx950/fmha_v3_fwd/fwd_hd192_hd128_bf16_group.co differ diff --git a/op_tests/test_mha_fwd_large_kv.py b/op_tests/test_mha_fwd_large_kv.py new file mode 100644 index 00000000000..6cbac45f9f1 --- /dev/null +++ b/op_tests/test_mha_fwd_large_kv.py @@ -0,0 +1,200 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +"""Forward-only large-KV correctness check for the v3 asm FMHA fwd kernels. + +Purpose: verify that a single batch's K/V slice exceeding 4 GiB is addressed +correctly. `buffer_load` uses a 32-bit offset from the (per-batch) buffer base, +so anything past 4 GiB reads as 0; `global_load` uses a full 64-bit address and +must stay correct. Run this against the buffer-load .co (expect FAIL on the tail) +and the global-load .co (expect PASS). + +Per-batch byte size (bf16): + K = seqlen_k * nheads_k * hdim_q * 2 + V = seqlen_k * nheads_k * hdim_v * 2 +Both must exceed 4 GiB to actually cross the boundary for K and V. + +Only the forward pass is exercised (no autograd), and the reference uses a +flash-style chunked online softmax so host memory stays bounded regardless of +seqlen_k. + +Example (K~7.4GB, V~4.9GB per batch): + python3 op_tests/test_mha_fwd_large_kv.py -b 1 -n 64 -q 256 -k 300000 -c false + +Routing: flash_attn_func only picks fmha_v3_fwd when seqlen_q > 128 (see mha.py); +with -q 128 you get CK mha_fwd (.jit), not hsa/gfx950/fmha_v3_fwd/*.co. + +This script forces v3 asm (fmha_v3_fwd .co), not OPUS: + AITER_DISABLE_FMHA_OPUS=1 is set before importing aiter. +""" + +import argparse +import os + +# hd192 default-on OPUS would bypass fmha_v3_fwd .co; this test targets asm v3. +os.environ.setdefault("AITER_DISABLE_FMHA_OPUS", "1") + +import torch + +import aiter +from aiter import dtypes + +_4GiB = 4 * 1024**3 + +_V3_CO_DIR = "hsa/gfx950/fmha_v3_fwd" + + +def expected_v3_co(causal: bool) -> str: + if causal: + return "fwd_hd192_hd128_bf16_causal.co" + return "fwd_hd192_hd128_bf16.co" + + +def chunked_attn_ref(q, k, v, causal, scale, kv_chunk=8192): + """Flash-style fp32 reference. q/k/v: [b, h, s, d] (h already broadcast for GQA).""" + b, h, sq, dq = q.shape + sk = k.shape[2] + dv = v.shape[3] + qf = q.float() + + row_max = torch.full( + (b, h, sq, 1), float("-inf"), device=q.device, dtype=torch.float32 + ) + softmax_denom = torch.zeros((b, h, sq, 1), device=q.device, dtype=torch.float32) + acc = torch.zeros((b, h, sq, dv), device=q.device, dtype=torch.float32) + + q_idx = torch.arange(sq, device=q.device).view(1, 1, sq, 1) + for start in range(0, sk, kv_chunk): + end = min(start + kv_chunk, sk) + kc = k[:, :, start:end, :].float() + vc = v[:, :, start:end, :].float() + s = torch.matmul(qf, kc.transpose(-1, -2)) * scale # [b,h,sq,chunk] + if causal: + k_idx = torch.arange(start, end, device=q.device).view(1, 1, 1, -1) + # align the causal diagonal to the bottom-right (sk may exceed sq) + s = s.masked_fill(k_idx > (q_idx + (sk - sq)), float("-inf")) + + m_new = torch.maximum(row_max, s.max(dim=-1, keepdim=True).values) + p = torch.exp(s - m_new) + alpha = torch.exp(row_max - m_new) + softmax_denom = softmax_denom * alpha + p.sum(dim=-1, keepdim=True) + acc = acc * alpha + torch.matmul(p, vc) + row_max = m_new + + out = acc / softmax_denom.clamp_min(1e-20) + return out.to(q.dtype) + + +def parse_args(): + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawTextHelpFormatter + ) + p.add_argument("-b", "--batch_size", type=int, default=1) + p.add_argument("-n", "--nheads", type=int, default=64) + p.add_argument("-gr", "--gqa_ratio", type=int, default=1) + p.add_argument( + "-q", + "--seqlen_q", + type=int, + default=256, + help="must be >128 for flash_attn_func to route to fmha_v3_fwd asm", + ) + p.add_argument("-k", "--seqlen_k", type=int, default=300000) + p.add_argument("-d_qk_v", type=dtypes.str2tuple, default=(192, 128)) + p.add_argument("-c", "--causal", type=dtypes.str2bool, default=False) + p.add_argument("--kv_chunk", type=int, default=8192, help="reference kv block size") + return p.parse_args() + + +def main(): + args = parse_args() + dq, dv = args.d_qk_v + assert args.nheads % args.gqa_ratio == 0 + nheads_k = args.nheads // args.gqa_ratio + dtype = dtypes.bf16 + device = "cuda" + + k_bytes = args.seqlen_k * nheads_k * dq * 2 + v_bytes = args.seqlen_k * nheads_k * dv * 2 + print( + f"per-batch K = {k_bytes/1024**3:.2f} GiB " + f"({'>' if k_bytes > _4GiB else '<='} 4 GiB), " + f"V = {v_bytes/1024**3:.2f} GiB " + f"({'>' if v_bytes > _4GiB else '<='} 4 GiB)" + ) + if k_bytes <= _4GiB and v_bytes <= _4GiB: + print( + "WARNING: neither K nor V exceeds 4 GiB; this will NOT exercise the " + ">4GiB path. Increase -k / -n." + ) + + if args.seqlen_q <= 128: + print( + "WARNING: seqlen_q <= 128 -> flash_attn_func uses CK mha_fwd, " + "NOT fmha_v3_fwd .co. Use -q 129 or higher." + ) + + co_name = expected_v3_co(args.causal) + print(f"AITER_DISABLE_FMHA_OPUS={os.environ.get('AITER_DISABLE_FMHA_OPUS', '0')}") + print( + f"expect fmha_v3_fwd (mode=0 batch): {_V3_CO_DIR}/{co_name} " + f"(seqlen_q={args.seqlen_q}, causal={args.causal})" + ) + + torch.manual_seed(0) + scale = dq**-0.5 + + q = torch.randn( + args.batch_size, args.seqlen_q, args.nheads, dq, device=device, dtype=dtype + ) + k = torch.randn( + args.batch_size, args.seqlen_k, nheads_k, dq, device=device, dtype=dtype + ) + v = torch.randn( + args.batch_size, args.seqlen_k, nheads_k, dv, device=device, dtype=dtype + ) + + with torch.no_grad(): + out = aiter.flash_attn_func( + q, + k, + v, + dropout_p=0.0, + softmax_scale=scale, + causal=args.causal, + window_size=(-1, -1), + return_lse=False, + return_attn_probs=False, + ) + if isinstance(out, (tuple, list)): + out = out[0] + + # BSHD -> BHSD; broadcast kv heads for GQA + qb = q.transpose(1, 2) + kb = k.transpose(1, 2) + vb = v.transpose(1, 2) + if args.gqa_ratio > 1: + kb = kb.repeat_interleave(args.gqa_ratio, dim=1) + vb = vb.repeat_interleave(args.gqa_ratio, dim=1) + + with torch.no_grad(): + ref = chunked_attn_ref(qb, kb, vb, args.causal, scale, args.kv_chunk) + ref = ref.transpose(1, 2) # back to BSHD + + diff = (out.float() - ref.float()).abs() + max_diff = diff.max().item() + mean_diff = diff.mean().item() + tol = 0.02 # bf16 + print(f"out max diff: {max_diff:.6f} mean diff: {mean_diff:.6f} tol: {tol}") + print( + f"out abs max: {out.abs().max().item():.4f} ref abs max: {ref.abs().max().item():.4f}" + ) + if max_diff <= tol: + print("#TEST PASSED") + else: + print( + "#TEST FAILED (tail KV likely mis-addressed -> buffer_load 4GiB overflow)" + ) + + +if __name__ == "__main__": + main() diff --git a/op_tests/test_mha_varlen_large_kv.py b/op_tests/test_mha_varlen_large_kv.py new file mode 100644 index 00000000000..062d429e989 --- /dev/null +++ b/op_tests/test_mha_varlen_large_kv.py @@ -0,0 +1,178 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +"""Forward-only large-seqlen check for the v3 asm FMHA *varlen/group* kernels. + +Uses `flash_attn_varlen_func` (cu_seqlens) so it loads the group .co: + causal=False -> fwd_hd192_hd128_bf16_group.co + causal=True -> fwd_hd192_hd128_bf16_causal_group.co + +Builds one sequence of length S and checks only the last `--check_rows` query rows +(global positions near S-1). Reference cost is O(check_rows * S), not O(S^2). + +Example (bracket 2^23 = 8,388,608): + python3 op_tests/test_mha_varlen_large_kv.py -n 1 -s 8000000 -c true + python3 op_tests/test_mha_varlen_large_kv.py -n 1 -s 900000 -c false + +When max_seqlen_q implies Q tiles > 65535, aiter host splits gdz launches and sets +s_q_tile_base on each chunk (still fwd_*_group.co). + +Forces v3 asm group .co (AITER_DISABLE_FMHA_OPUS=1 before import). +""" + +import argparse +import os + +os.environ.setdefault("AITER_DISABLE_FMHA_OPUS", "1") + +import torch + +import aiter +from aiter import dtypes + +_2P23 = 2**23 # 8,388,608 +_V3_CO_DIR = "hsa/gfx950/fmha_v3_fwd" + + +def v3_group_dispatch_mode(seqlen: int, causal: bool) -> int: + """fmha_fwd.csv mode for group on gfx950 hd192 (always mode=1 group .co).""" + del seqlen, causal + return 1 + + +def expected_group_co(causal: bool) -> str: + base = "fwd_hd192_hd128_bf16" + if causal: + base += "_causal" + return base + "_group.co" + + +def ref_last_rows(q, k, v, q0, causal, scale, kv_chunk=8192): + """Exact fp32 attention for query rows [q0:]. q/k/v: [h, S, d] (kv broadcast).""" + h, S, _dq = q.shape + dv = v.shape[2] + qf = q[:, q0:, :].float() + R = qf.shape[1] + + row_max = torch.full((h, R, 1), float("-inf"), device=q.device, dtype=torch.float32) + softmax_denom = torch.zeros((h, R, 1), device=q.device, dtype=torch.float32) + acc = torch.zeros((h, R, dv), device=q.device, dtype=torch.float32) + + q_idx = torch.arange(q0, S, device=q.device).view(1, R, 1) + for start in range(0, S, kv_chunk): + end = min(start + kv_chunk, S) + kc = k[:, start:end, :].float() + vc = v[:, start:end, :].float() + s = torch.matmul(qf, kc.transpose(-1, -2)) * scale + if causal: + k_idx = torch.arange(start, end, device=q.device).view(1, 1, -1) + s = s.masked_fill(k_idx > q_idx, float("-inf")) + m_new = torch.maximum(row_max, s.max(dim=-1, keepdim=True).values) + p = torch.exp(s - m_new) + alpha = torch.exp(row_max - m_new) + softmax_denom = softmax_denom * alpha + p.sum(dim=-1, keepdim=True) + acc = acc * alpha + torch.matmul(p, vc) + row_max = m_new + return (acc / softmax_denom.clamp_min(1e-20)).to(q.dtype) + + +def parse_args(): + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawTextHelpFormatter + ) + p.add_argument("-n", "--nheads", type=int, default=1) + p.add_argument("-gr", "--gqa_ratio", type=int, default=1) + p.add_argument( + "-s", + "--seqlen", + type=int, + default=900000, + help="single-sequence length S (sq == sk == S)", + ) + p.add_argument("-c", "--causal", type=dtypes.str2bool, default=False) + p.add_argument("-d_qk_v", type=dtypes.str2tuple, default=(192, 128)) + p.add_argument("--check_rows", type=int, default=512) + p.add_argument("--kv_chunk", type=int, default=8192) + return p.parse_args() + + +def main(): + args = parse_args() + dq, dv = args.d_qk_v + assert args.nheads % args.gqa_ratio == 0 + nheads_k = args.nheads // args.gqa_ratio + S = args.seqlen + dtype = dtypes.bf16 + device = "cuda" + scale = dq**-0.5 + + v3_mode = v3_group_dispatch_mode(S, args.causal) + co_name = expected_group_co(args.causal) + + print( + f"seqlen S = {S:,} ({'>' if S > _2P23 else '<='} 2^23 = {_2P23:,}) " + f"causal={args.causal} nheads={args.nheads}/{nheads_k}" + ) + print(f"AITER_DISABLE_FMHA_OPUS={os.environ.get('AITER_DISABLE_FMHA_OPUS', '0')}") + print(f"expect fmha_v3_fwd mode={v3_mode}: {_V3_CO_DIR}/{co_name}") + + torch.manual_seed(0) + q = torch.randn(S, args.nheads, dq, device=device, dtype=dtype) + k = torch.randn(S, nheads_k, dq, device=device, dtype=dtype) + v = torch.randn(S, nheads_k, dv, device=device, dtype=dtype) + cu_q = torch.tensor([0, S], device=device, dtype=torch.int32) + cu_k = torch.tensor([0, S], device=device, dtype=torch.int32) + + with torch.no_grad(): + out = aiter.flash_attn_varlen_func( + q, + k, + v, + cu_q, + cu_k, + S, + S, + dropout_p=0.0, + softmax_scale=scale, + causal=args.causal, + window_size=(-1, -1, 0), + return_lse=False, + return_attn_probs=False, + ) + if isinstance(out, (tuple, list)): + out = out[0] + + q0 = max(0, S - args.check_rows) + qb = q.transpose(0, 1) + kb = k.transpose(0, 1) + vb = v.transpose(0, 1) + if args.gqa_ratio > 1: + kb = kb.repeat_interleave(args.gqa_ratio, dim=0) + vb = vb.repeat_interleave(args.gqa_ratio, dim=0) + + with torch.no_grad(): + ref = ref_last_rows(qb, kb, vb, q0, args.causal, scale, args.kv_chunk) + out_last = out[q0:].transpose(0, 1).float() + + diff = (out_last - ref.float()).abs() + max_diff = diff.max().item() + mean_diff = diff.mean().item() + tol = 0.02 + n_bad = int((diff > tol).sum().item()) + finite = torch.isfinite(out_last).all().item() + print(f"checked last {out_last.shape[1]} query rows (global pos {q0:,}..{S-1:,})") + print( + f"out finite: {finite} abs max: {out_last.abs().max().item():.4f} " + f"ref abs max: {ref.float().abs().max().item():.4f}" + ) + print( + f"max diff: {max_diff:.6f} mean diff: {mean_diff:.6f} " + f"bad(>{tol}): {n_bad}/{out_last.numel()} tol: {tol}" + ) + if finite and max_diff <= tol: + print("#TEST PASSED") + else: + print("#TEST FAILED (large-seqlen / global-load group .co mismatch)") + + +if __name__ == "__main__": + main()