From 42387e2c05f1378ea20c6d89a6619a76fa0035dc Mon Sep 17 00:00:00 2001 From: adlashab Date: Fri, 26 Jun 2026 18:21:00 +0200 Subject: [PATCH 1/2] [Kernel][Triton] sliding-window decode over shuffled fp8 paged KV The shuffled fp8 KV layout is written by reshape_and_cache_shuffle and read at decode by the asm paged-attention path, which has no window argument. So a model with sliding-window attention can't use the shuffled cache: the only windowed fp8 decode is the NHD paged_attention_v1 path, on a different layout. There was no kernel that reads the shuffled layout and honors a window. Add paged_attention_decode_shuffle_swa, a triton paged decode that mirrors the shuffled K/V indexing, dequantizes fp8 KV to fp32 with per-tensor or per-token scales (A16W8, query stays 16-bit), runs an online-softmax decode, and masks keys outside [seq_len - window, seq_len). One program per (sequence, kv head), decode only. Validated against an independent fp8 reference on MI300X (gfx942) and MI355X (gfx950), both e4m3fn and e4m3fnuz, across head 128, GQA 32/8 and MQA, block 16/32, windowed and full, per-tensor and per-token scales. Max abs error stays below the fp8 quantization floor on both arches. Signed-off-by: adlashab --- .../ops/triton/attention/pa_decode_shuffle.py | 196 ++++++++++++++++++ 1 file changed, 196 insertions(+) create mode 100644 aiter/ops/triton/attention/pa_decode_shuffle.py diff --git a/aiter/ops/triton/attention/pa_decode_shuffle.py b/aiter/ops/triton/attention/pa_decode_shuffle.py new file mode 100644 index 00000000000..55588e8b8a3 --- /dev/null +++ b/aiter/ops/triton/attention/pa_decode_shuffle.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import triton +import triton.language as tl + +from aiter.ops.triton.utils.logger import AiterTritonLogger + +_LOGGER = AiterTritonLogger() + +# Sliding-window paged decode over the shuffled fp8 KV cache. The asm pa_fwd +# decode path takes no window arg, so SWA layers can't use the shuffled cache; +# this reads the same layout and masks the window. +# +# Shuffled layout (x = 16 // elem_size, = 16 for fp8): +# K: [num_blocks, num_kv_heads, head_size // x, block_size, x] +# V: [num_blocks, num_kv_heads, block_size // x, head_size, x] +# K[d, slot] -> (d // x) * block_size * x + slot * x + (d % x) +# V[slot, d] -> (slot // x) * head_size * x + d * x + (slot % x) +# A16W8: q is 16-bit, fp8 KV dequantized (val * scale) before the dot. + + +@triton.jit +def _pa_decode_shuffle_swa_kernel( + out_ptr, + q_ptr, + k_cache_ptr, + v_cache_ptr, + blk_tbl_ptr, + seq_lens_ptr, + k_scale_ptr, + v_scale_ptr, + scale, + sliding_window, + q_s0, + q_s1, + q_s2, + o_s0, + o_s1, + o_s2, + bt_s0, + ks_s0, + ks_s1, + ks_s2, + k_blk_stride, + k_head_stride, + v_blk_stride, + v_head_stride, + HEAD: tl.constexpr, + HEAD_POW2: tl.constexpr, + BLOCK: tl.constexpr, + X: tl.constexpr, + GRP: tl.constexpr, + GRP_POW2: tl.constexpr, + SCALE_SCALAR: tl.constexpr, +): + seq = tl.program_id(0) + kvh = tl.program_id(1) + + seq_len = tl.load(seq_lens_ptr + seq) + win_start = 0 + if sliding_window > 0: + ws = seq_len - sliding_window + win_start = tl.where(ws > 0, ws, 0) + + grp = tl.arange(0, GRP_POW2) + d = tl.arange(0, HEAD_POW2) + q_head = kvh * GRP + grp + q_mask = (grp[:, None] < GRP) & (d[None, :] < HEAD) + q_ptrs = q_ptr + seq * q_s0 + q_head[:, None] * q_s1 + d[None, :] * q_s2 + q = tl.load(q_ptrs, mask=q_mask, other=0.0).to(tl.float32) + + m_i = tl.full([GRP_POW2], -float("inf"), tl.float32) + l_i = tl.zeros([GRP_POW2], tl.float32) + acc = tl.zeros([GRP_POW2, HEAD_POW2], tl.float32) + + s = tl.arange(0, BLOCK) + n_blk = tl.cdiv(seq_len, BLOCK) + start_blk = win_start // BLOCK + + if SCALE_SCALAR: + ks_scalar = tl.load(k_scale_ptr) + vs_scalar = tl.load(v_scale_ptr) + + for b in range(start_blk, n_blk): + block_id = tl.load(blk_tbl_ptr + seq * bt_s0 + b).to(tl.int64) + pos = b * BLOCK + s + valid = (pos < seq_len) & (pos >= win_start) + + k_base = block_id * k_blk_stride + kvh * k_head_stride + v_base = block_id * v_blk_stride + kvh * v_head_stride + ld_mask = (s[:, None] < BLOCK) & (d[None, :] < HEAD) + + k_off = ( + k_base + (d[None, :] // X) * (BLOCK * X) + s[:, None] * X + (d[None, :] % X) + ) + v_off = ( + v_base + (s[:, None] // X) * (HEAD * X) + d[None, :] * X + (s[:, None] % X) + ) + + k = tl.load(k_cache_ptr + k_off, mask=ld_mask, other=0.0).to(tl.float32) + v = tl.load(v_cache_ptr + v_off, mask=ld_mask, other=0.0).to(tl.float32) + + if SCALE_SCALAR: + k = k * ks_scalar + v = v * vs_scalar + else: + sc_off = block_id * ks_s0 + kvh * ks_s1 + s * ks_s2 + ks = tl.load(k_scale_ptr + sc_off, mask=s < BLOCK, other=0.0) + vs = tl.load(v_scale_ptr + sc_off, mask=s < BLOCK, other=0.0) + k = k * ks[:, None] + v = v * vs[:, None] + + scores = tl.dot(q, tl.trans(k)) * scale + scores = tl.where(valid[None, :], scores, -float("inf")) + + m_new = tl.maximum(m_i, tl.max(scores, axis=1)) + p = tl.exp(scores - m_new[:, None]) + alpha = tl.exp(m_i - m_new) + l_i = l_i * alpha + tl.sum(p, axis=1) + acc = acc * alpha[:, None] + tl.dot(p.to(tl.float32), v) + m_i = m_new + + out = acc / l_i[:, None] + o_ptrs = out_ptr + seq * o_s0 + q_head[:, None] * o_s1 + d[None, :] * o_s2 + tl.store(o_ptrs, out.to(out_ptr.dtype.element_ty), mask=q_mask) + + +def paged_attention_decode_shuffle_swa( + out, # [num_seqs, num_q_heads, head_size] + query, # [num_seqs, num_q_heads, head_size] + key_cache, # [num_blocks, num_kv_heads, head_size // x, block_size, x] + value_cache, # [num_blocks, num_kv_heads, block_size // x, head_size, x] + block_tables, # [num_seqs, max_num_blocks] + seq_lens, # [num_seqs] + scale, # softmax scale + k_scale, # scalar or [num_blocks, num_kv_heads, block_size] + v_scale, + sliding_window, # number of keys attended (<= 0 means full causal) +): + """Sliding-window decode over the shuffled fp8 KV layout. One program per + (sequence, kv head); decode only (one query token per sequence).""" + num_seqs, num_q_heads, head = query.shape + num_kv_heads = key_cache.shape[1] + block_size = key_cache.shape[3] + x = key_cache.shape[4] + grp = num_q_heads // num_kv_heads + grp_pow2 = max(16, triton.next_power_of_2(grp)) + head_pow2 = triton.next_power_of_2(head) + + _LOGGER.info( + f"PA_DECODE_SHUFFLE_SWA: q={tuple(query.shape)} " + f"k_cache={tuple(key_cache.shape)} window={sliding_window}" + ) + + scale_scalar = k_scale.numel() == 1 + if scale_scalar: + ks_s0 = ks_s1 = ks_s2 = 0 + else: + ks_s0, ks_s1, ks_s2 = k_scale.stride() + + grid = (num_seqs, num_kv_heads) + _pa_decode_shuffle_swa_kernel[grid]( + out, + query, + key_cache, + value_cache, + block_tables, + seq_lens, + k_scale, + v_scale, + scale, + sliding_window, + query.stride(0), + query.stride(1), + query.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + block_tables.stride(0), + ks_s0, + ks_s1, + ks_s2, + key_cache.stride(0), + key_cache.stride(1), + value_cache.stride(0), + value_cache.stride(1), + HEAD=head, + HEAD_POW2=head_pow2, + BLOCK=block_size, + X=x, + GRP=grp, + GRP_POW2=grp_pow2, + SCALE_SCALAR=scale_scalar, + ) + return out From 3ea37a3a5e149a16222a307b91aed7e7e6511034 Mon Sep 17 00:00:00 2001 From: adlashab Date: Tue, 20 Oct 2026 05:05:17 -0400 Subject: [PATCH 2/2] [Triton] pa_decode_shuffle: zero unused V slots, zero-length rows, add test and bench Unused slots in the first and last block of the window get p = 0, but a NaN in their V or V scale still went through p @ v as 0 * NaN and made the whole row NaN (seq_len 1 with 16-slot pages is enough). Zero those V rows. Only the edge blocks have unused slots, so interior blocks skip the select. Rows with seq_len 0 (cudagraph padding) now return zeros, not 0 / 0. V scales use their own strides, so K and V scale tensors can differ. Log with % args and give the kernel a make_kernel_repr name. Add op_tests/triton_tests/attention/test_pa_decode_shuffle.py (both fp8 types, scalar and per-token scales, 16 and 32 slot pages, GQA, window edges, NaN in every unused slot, a seq_len 0 row) and op_tests/op_benchmarks/triton/bench_pa_decode_shuffle.py. Signed-off-by: adlashab --- .../ops/triton/attention/pa_decode_shuffle.py | 54 +++++- .../triton/bench_pa_decode_shuffle.py | 143 ++++++++++++++++ .../attention/test_pa_decode_shuffle.py | 159 ++++++++++++++++++ 3 files changed, 348 insertions(+), 8 deletions(-) create mode 100644 op_tests/op_benchmarks/triton/bench_pa_decode_shuffle.py create mode 100644 op_tests/triton_tests/attention/test_pa_decode_shuffle.py diff --git a/aiter/ops/triton/attention/pa_decode_shuffle.py b/aiter/ops/triton/attention/pa_decode_shuffle.py index 55588e8b8a3..8304c077ae2 100644 --- a/aiter/ops/triton/attention/pa_decode_shuffle.py +++ b/aiter/ops/triton/attention/pa_decode_shuffle.py @@ -4,6 +4,7 @@ import triton import triton.language as tl +from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr from aiter.ops.triton.utils.logger import AiterTritonLogger _LOGGER = AiterTritonLogger() @@ -19,8 +20,13 @@ # V[slot, d] -> (slot // x) * head_size * x + d * x + (slot % x) # A16W8: q is 16-bit, fp8 KV dequantized (val * scale) before the dot. +_pa_decode_shuffle_swa_repr = make_kernel_repr( + "_pa_decode_shuffle_swa_kernel", + ["HEAD", "BLOCK", "GRP", "SCALE_SCALAR"], +) -@triton.jit + +@triton.jit(repr=_pa_decode_shuffle_swa_repr) def _pa_decode_shuffle_swa_kernel( out_ptr, q_ptr, @@ -42,6 +48,9 @@ def _pa_decode_shuffle_swa_kernel( ks_s0, ks_s1, ks_s2, + vs_s0, + vs_s1, + vs_s2, k_blk_stride, k_head_stride, v_blk_stride, @@ -105,11 +114,17 @@ def _pa_decode_shuffle_swa_kernel( k = k * ks_scalar v = v * vs_scalar else: - sc_off = block_id * ks_s0 + kvh * ks_s1 + s * ks_s2 - ks = tl.load(k_scale_ptr + sc_off, mask=s < BLOCK, other=0.0) - vs = tl.load(v_scale_ptr + sc_off, mask=s < BLOCK, other=0.0) + ks_off = block_id * ks_s0 + kvh * ks_s1 + s * ks_s2 + vs_off = block_id * vs_s0 + kvh * vs_s1 + s * vs_s2 + ks = tl.load(k_scale_ptr + ks_off, mask=s < BLOCK, other=0.0) + vs = tl.load(v_scale_ptr + vs_off, mask=s < BLOCK, other=0.0) k = k * ks[:, None] v = v * vs[:, None] + # unused slots get p = 0, but 0 * NaN is NaN in the p @ v dot, so a NaN + # left in an unused V slot or its scale has to be zeroed. Only the + # window edge blocks have unused slots, interior blocks skip the select. + if (b * BLOCK < win_start) | (b * BLOCK + BLOCK > seq_len): + v = tl.where(valid[:, None], v, 0.0) scores = tl.dot(q, tl.trans(k)) * scale scores = tl.where(valid[None, :], scores, -float("inf")) @@ -121,7 +136,8 @@ def _pa_decode_shuffle_swa_kernel( acc = acc * alpha[:, None] + tl.dot(p.to(tl.float32), v) m_i = m_new - out = acc / l_i[:, None] + # seq_len == 0 (cudagraph padding rows) never enters the loop: write zeros, not 0 / 0 + out = tl.where(l_i[:, None] > 0.0, acc / l_i[:, None], 0.0) o_ptrs = out_ptr + seq * o_s0 + q_head[:, None] * o_s1 + d[None, :] * o_s2 tl.store(o_ptrs, out.to(out_ptr.dtype.element_ty), mask=q_mask) @@ -139,7 +155,22 @@ def paged_attention_decode_shuffle_swa( sliding_window, # number of keys attended (<= 0 means full causal) ): """Sliding-window decode over the shuffled fp8 KV layout. One program per - (sequence, kv head); decode only (one query token per sequence).""" + (sequence, kv head), one query token per sequence. + + Args: + out: output, same shape as query, written in place. + query: 16-bit query, [num_seqs, num_q_heads, head_size]. + key_cache, value_cache: fp8 caches in the shuffled layout above. + block_tables: int32 [num_seqs, max_num_blocks]. + seq_lens: int32 [num_seqs]. A row with seq_len 0 gets zeros. + scale: softmax scale. + k_scale, v_scale: one scalar each, or per token [num_blocks, + num_kv_heads, block_size]. K and V scales can have different strides. + sliding_window: keys attended per sequence, <= 0 means all. + + Returns: + out. + """ num_seqs, num_q_heads, head = query.shape num_kv_heads = key_cache.shape[1] block_size = key_cache.shape[3] @@ -149,15 +180,19 @@ def paged_attention_decode_shuffle_swa( head_pow2 = triton.next_power_of_2(head) _LOGGER.info( - f"PA_DECODE_SHUFFLE_SWA: q={tuple(query.shape)} " - f"k_cache={tuple(key_cache.shape)} window={sliding_window}" + "PA_DECODE_SHUFFLE_SWA: q=%s k_cache=%s window=%s", + tuple(query.shape), + tuple(key_cache.shape), + sliding_window, ) scale_scalar = k_scale.numel() == 1 if scale_scalar: ks_s0 = ks_s1 = ks_s2 = 0 + vs_s0 = vs_s1 = vs_s2 = 0 else: ks_s0, ks_s1, ks_s2 = k_scale.stride() + vs_s0, vs_s1, vs_s2 = v_scale.stride() grid = (num_seqs, num_kv_heads) _pa_decode_shuffle_swa_kernel[grid]( @@ -181,6 +216,9 @@ def paged_attention_decode_shuffle_swa( ks_s0, ks_s1, ks_s2, + vs_s0, + vs_s1, + vs_s2, key_cache.stride(0), key_cache.stride(1), value_cache.stride(0), diff --git a/op_tests/op_benchmarks/triton/bench_pa_decode_shuffle.py b/op_tests/op_benchmarks/triton/bench_pa_decode_shuffle.py new file mode 100644 index 00000000000..085628c5281 --- /dev/null +++ b/op_tests/op_benchmarks/triton/bench_pa_decode_shuffle.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +""" +Benchmark for sliding-window decode over the shuffled fp8 paged KV cache. +""" + +import argparse +import sys + +import torch +import triton + +from aiter.ops.triton.attention.pa_decode_shuffle import ( + paged_attention_decode_shuffle_swa, +) +from op_tests.op_benchmarks.triton.utils.benchmark_utils import ( + get_caller_name_no_ext, + print_vgpr, +) +from op_tests.triton_tests.attention.test_pa_decode_shuffle import ( + make_shuffle_inputs, +) + +# (BS, SEQ_LEN, WINDOW) +_POINTS = ( + (1, 8192, 4096), + (8, 8192, 4096), + (32, 8192, 4096), + (32, 2048, 1024), + (32, 512, 256), + (64, 4096, 0), +) + +arg_to_torch_dtype = { + "e4m3fn": torch.float8_e4m3fn, + "e4m3fnuz": torch.float8_e4m3fnuz, +} + + +def get_benchmark_shapes(args): + if args.b: + return [(args.b, args.sq, args.window)] + return list(_POINTS) + + +def bench_fn(BS, SEQ_LEN, WINDOW, metric, args): + q, kc, vc, bt, sl, ks, vs, _, _ = make_shuffle_inputs( + [SEQ_LEN] * BS, + args.hq, + args.hk, + args.d, + args.block, + arg_to_torch_dtype[args.kv_cache_dtype], + args.per_token_scale, + ) + out = torch.empty_like(q) + + def fn(): + paged_attention_decode_shuffle_swa( + out, q, kc, vc, bt, sl, args.d**-0.5, ks, vs, WINDOW + ) + + ms = triton.testing.do_bench(fn, warmup=args.warmup, rep=args.rep) + if metric == "time": + return ms * 1000 # us + if metric == "bandwidth": + keys = SEQ_LEN if WINDOW <= 0 else min(WINDOW, SEQ_LEN) + # fp8 K + V for the window, plus query and output + mem = BS * args.hk * keys * args.d * 2 * kc.element_size() + mem += 2 * q.numel() * q.element_size() + return mem / (ms * 1e-3) * 1e-9 # GB/s + raise ValueError(f"unknown metric {metric}") + + +def run_benchmark(args): + metrics = ["time", "bandwidth"] if args.metric == "all" else [args.metric] + + benchmark = triton.testing.Benchmark( + x_names=["BS", "SEQ_LEN", "WINDOW"], + x_vals=get_benchmark_shapes(args), + line_arg="metric", + line_vals=metrics, + line_names=[ + "Time_(us)" if m == "time" else "Bandwidth_(GB/s)" for m in metrics + ], + styles=[("red", "-"), ("blue", "-")], + ylabel="us / GB/s", + plot_name=get_caller_name_no_ext() + f"_HQ{args.hq}_HK{args.hk}_D{args.d}", + args={}, + ) + + @triton.testing.perf_report([benchmark]) + def bench(BS, SEQ_LEN, WINDOW, metric): + return bench_fn(BS, SEQ_LEN, WINDOW, metric, args) + + bench.run(save_path="." if args.o else None, print_data=True) + + +def parse_args(): + parser = argparse.ArgumentParser( + prog="Benchmark pa_decode_shuffle", + description="Sliding-window decode over the shuffled fp8 KV cache", + allow_abbrev=False, + ) + parser.add_argument("-b", type=int, default=0, help="Batch size (0: default list)") + parser.add_argument("-sq", type=int, default=8192, help="Context length") + parser.add_argument("--window", type=int, default=4096, help="<= 0: full") + parser.add_argument("-hq", type=int, default=32) + parser.add_argument("-hk", type=int, default=8) + parser.add_argument("-d", type=int, default=128) + parser.add_argument("--block", type=int, default=16) + parser.add_argument( + "-kv_cache_dtype", default="e4m3fn", choices=list(arg_to_torch_dtype) + ) + parser.add_argument("--per_token_scale", action="store_true") + parser.add_argument("--metric", default="all", choices=["all", "time", "bandwidth"]) + parser.add_argument("--warmup", type=int, default=25) + parser.add_argument("--rep", type=int, default=100) + parser.add_argument( + "-print_vgpr", + action="store_true", + default=False, + help="Print VGPR usage for Triton kernels", + ) + parser.add_argument( + "-o", action="store_true", default=False, help="Write results to a CSV file" + ) + return parser.parse_args() + + +def main(): + args = parse_args() + if args.print_vgpr: + print("Retrieving VGPR usage for pa_decode_shuffle Triton kernels...") + print_vgpr(lambda: run_benchmark(args), get_caller_name_no_ext()) + return 0 + run_benchmark(args) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/op_tests/triton_tests/attention/test_pa_decode_shuffle.py b/op_tests/triton_tests/attention/test_pa_decode_shuffle.py new file mode 100644 index 00000000000..15bf6a7215e --- /dev/null +++ b/op_tests/triton_tests/attention/test_pa_decode_shuffle.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import pytest +import torch + +from aiter.ops.triton.attention.pa_decode_shuffle import ( + paged_attention_decode_shuffle_swa, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires a ROCm GPU" +) + +FP8_DTYPES = [torch.float8_e4m3fn, torch.float8_e4m3fnuz] +X = 16 # 16 // fp8 elem size + + +def make_shuffle_inputs( + seq_lens, + num_q_heads, + num_kv_heads, + head, + block_size, + fp8_dtype, + per_token_scale, + device="cuda", + seed=0, +): + """Shuffled fp8 KV cache for the given seq_lens. Unused slots and their + per-token scales hold NaN. Shared with the benchmark. + + Returns (query, key_cache, value_cache, block_tables, seq_lens, k_scale, + v_scale, k_deq, v_deq), k_deq / v_deq are the dequantized fp32 tokens per + sequence for the reference.""" + gen = torch.Generator(device=device).manual_seed(seed) + num_seqs = len(seq_lens) + blocks = [(n + block_size - 1) // block_size for n in seq_lens] + max_blocks = max(1, max(blocks)) + num_blocks = sum(blocks) + 2 + perm = torch.randperm(num_blocks, generator=gen, device=device).to(torch.int32) + block_tables = torch.zeros((num_seqs, max_blocks), dtype=torch.int32, device=device) + nb = 0 + for i, n in enumerate(blocks): + block_tables[i, :n] = perm[nb : nb + n] + nb += n + + key_cache = torch.full( + (num_blocks, num_kv_heads, head // X, block_size, X), + float("nan"), + dtype=fp8_dtype, + device=device, + ) + value_cache = torch.full( + (num_blocks, num_kv_heads, block_size // X, head, X), + float("nan"), + dtype=fp8_dtype, + device=device, + ) + if per_token_scale: + # V scale is a strided view so K and V scales have different strides + k_scale = torch.full( + (num_blocks, num_kv_heads, block_size), float("nan"), device=device + ) + v_scale = torch.full( + (num_blocks, num_kv_heads, block_size, 2), float("nan"), device=device + )[..., 1] + else: + k_scale = torch.tensor([0.7], device=device) + v_scale = torch.tensor([1.3], device=device) + + fp8_max = torch.finfo(fp8_dtype).max + k_deq, v_deq = [], [] + for i, n in enumerate(seq_lens): + pos = torch.arange(n, device=device) + blk = block_tables[i, pos // block_size].long() + slot = pos % block_size + k = 2 * torch.randn((n, num_kv_heads, head), generator=gen, device=device) + v = 2 * torch.randn((n, num_kv_heads, head), generator=gen, device=device) + if per_token_scale: + ks = 0.5 + torch.rand((n, num_kv_heads), generator=gen, device=device) + vs = 0.5 + torch.rand((n, num_kv_heads), generator=gen, device=device) + k_scale[blk, :, slot] = ks + v_scale[blk, :, slot] = vs + ks, vs = ks[..., None], vs[..., None] + else: + ks, vs = k_scale, v_scale + kq = (k / ks).clamp(-fp8_max, fp8_max).to(fp8_dtype) + vq = (v / vs).clamp(-fp8_max, fp8_max).to(fp8_dtype) + key_cache[blk, :, :, slot, :] = kq.view(n, num_kv_heads, head // X, X) + value_cache[blk, :, slot // X, :, slot % X] = vq + k_deq.append(kq.float() * ks) + v_deq.append(vq.float() * vs) + + query = torch.randn((num_seqs, num_q_heads, head), generator=gen, device=device).to( + torch.bfloat16 + ) + seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=device) + return ( + query, + key_cache, + value_cache, + block_tables, + seq_lens_t, + k_scale, + v_scale, + k_deq, + v_deq, + ) + + +def ref_decode(query, k_deq, v_deq, seq_lens, scale, sliding_window): + num_seqs, num_q_heads, head = query.shape + out = torch.zeros((num_seqs, num_q_heads, head), device=query.device) + for i, n in enumerate(seq_lens): + if n == 0: + continue + lo = 0 if sliding_window <= 0 else max(0, n - sliding_window) + grp = num_q_heads // k_deq[i].shape[1] + k = k_deq[i][lo:n].repeat_interleave(grp, dim=1) + v = v_deq[i][lo:n].repeat_interleave(grp, dim=1) + p = torch.softmax(torch.einsum("hd,khd->hk", query[i].float(), k) * scale, -1) + out[i] = torch.einsum("hk,khd->hd", p, v) + return out + + +@pytest.mark.parametrize("fp8_dtype", FP8_DTYPES) +@pytest.mark.parametrize("per_token_scale", [False, True]) +@pytest.mark.parametrize( + "seq_lens, num_q_heads, num_kv_heads, block_size, window", + [ + ([300, 130, 600, 65], 32, 8, 16, 256), + ([300, 130, 600, 65], 32, 8, 16, 0), + ([500, 200, 90], 32, 8, 32, 256), + # one valid slot in a 16-slot page, window edges inside a page + ([1, 17, 50, 200], 32, 8, 16, 64), + ([0, 31, 255, 1000], 16, 1, 16, 100), + ([128, 256], 8, 8, 16, 128), + ], +) +def test_pa_decode_shuffle( + seq_lens, num_q_heads, num_kv_heads, block_size, window, per_token_scale, fp8_dtype +): + head = 128 + q, kc, vc, bt, sl, ks, vs, k_deq, v_deq = make_shuffle_inputs( + seq_lens, + num_q_heads, + num_kv_heads, + head, + block_size, + fp8_dtype, + per_token_scale, + ) + scale = head**-0.5 + out = torch.full_like(q, float("nan")) + paged_attention_decode_shuffle_swa(out, q, kc, vc, bt, sl, scale, ks, vs, window) + ref = ref_decode(q, k_deq, v_deq, seq_lens, scale, window) + assert torch.isfinite(out).all(), "NaN or inf in the output" + torch.testing.assert_close(out.float(), ref, atol=2e-2, rtol=2e-2)