diff --git a/.github/scripts/select_triton_tests.py b/.github/scripts/select_triton_tests.py index 3cc68a2fb82..3d31ee23d38 100755 --- a/.github/scripts/select_triton_tests.py +++ b/.github/scripts/select_triton_tests.py @@ -21,6 +21,14 @@ KERNELS = SRC + "_triton_kernels/" GLUON_KERNELS = SRC + "_gluon_kernels/" CONFIGS = SRC + "configs/" +# Ops whose Triton launch path is only exercised on MI300X (MI35X takes a +# Gluon branch instead), so a change to them needs the MI300X job even without +# a gfx942 config or Gluon-kernel change. +GFX942_LAUNCH_PATHS = ( + SRC + "attention/pa_prefill_sparse.py", + SRC + "_triton_kernels/attention/sparse_attention_dsv4.py", + "op_tests/triton_tests/attention/test_pa_prefill_sparse.py", +) TESTS = "op_tests/triton_tests/" BENCH = "op_tests/op_benchmarks/triton/" ROOT = Path(__file__).resolve().parents[2] @@ -436,6 +444,7 @@ def main(): diff = changed_files(args.merge_ref) needs_mi300x = any( path.startswith((CONFIGS + "gfx942/", GLUON_KERNELS + "gfx942/")) + or path in GFX942_LAUNCH_PATHS for path in diff ) if not args.all: diff --git a/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4.py b/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4.py index f5530a5195b..91462c63cc7 100644 --- a/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4.py +++ b/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4.py @@ -227,6 +227,13 @@ def _combine_topk_swa_indices_ragged_kernel( # --------------------------------------------------------------------------- +@triton.jit +def _exp(x, USE_EXP2: tl.constexpr): + if USE_EXP2: + return tl.math.exp2(x) + return tl.exp(x) + + def _prefill_prune_configs(configs, named_args, **kwargs): BLOCK_D = kwargs.get("BLOCK_D", named_args.get("BLOCK_D")) pruned = [] @@ -271,6 +278,9 @@ def _get_prefill_autotune_configs(): "BLOCK_H", "BLOCK_D", "BLOCK_K", + "USE_EXP2", + "EVEN_H", + "EVEN_D", ], ) @@ -304,6 +314,11 @@ def _sparse_attn_prefill_kernel( BLOCK_H: tl.constexpr, BLOCK_D: tl.constexpr, BLOCK_K: tl.constexpr, + USE_EXP2: tl.constexpr = False, + # num_heads % BLOCK_H == 0: no head mask. + EVEN_H: tl.constexpr = False, + # head_dim == BLOCK_D: no dim mask; the KV gather is masked by row only. + EVEN_D: tl.constexpr = False, ): # 64-bit before the multiply, same reasoning as `slot_off` below: the # program id fits 32 bits, but `query_idx * q_stride_t` does not once @@ -318,6 +333,11 @@ def _sparse_attn_prefill_kernel( head_mask = head_offsets < num_heads dim_mask = dim_offsets < head_dim + # Constant all-true masks fold away in the loads and the out store. + if EVEN_H: + head_mask = tl.full((BLOCK_H,), True, tl.int1) + if EVEN_D: + dim_mask = tl.full((BLOCK_D,), True, tl.int1) q = tl.load( q_ptr + query_idx * q_stride_t @@ -327,6 +347,8 @@ def _sparse_attn_prefill_kernel( other=0.0, ) + if USE_EXP2: + scale = scale * 1.4426950408889634 # log2(e): softmax runs in base 2 m_i = tl.full((BLOCK_H,), float("-inf"), dtype=tl.float32) l_i = tl.zeros((BLOCK_H,), dtype=tl.float32) acc = tl.zeros((BLOCK_H, BLOCK_D), dtype=tl.float32) @@ -346,20 +368,24 @@ def _sparse_attn_prefill_kernel( in_range = k_pos < kv_len valid = in_range & (slot >= 0) & (slot < num_kv) - # 64-bit before the multiply, same as the decode kernel: a slot index - # fits 32 bits (the index buffer is int32 by ABI) but `slot * + # 64-bit before the multiply, same as the decode kernel: an int32 + # slot index fits 32 bits but `slot * # kv_stride_n` does not once the unified V4 pool runs to ~150M rows. # The wrapped offset still lands inside the same allocation, so the # bad read is silent. slot_off = slot.to(tl.int64) - kv = tl.load( + kv_ptrs = ( kv_ptr + slot_off[:, None] * kv_stride_n - + dim_offsets[None, :] * kv_stride_d, - mask=valid[:, None] & dim_mask[None, :], - other=0.0, + + dim_offsets[None, :] * kv_stride_d ) + # Invalid rows are never loaded: a zero softmax weight does not + # neutralize a non-finite value in the PV dot (0 * NaN = NaN). + if EVEN_D: + kv = tl.load(kv_ptrs, mask=valid[:, None], other=0.0) + else: + kv = tl.load(kv_ptrs, mask=valid[:, None] & dim_mask[None, :], other=0.0) # Prefetch next tile's indices before heavy compute on current tile. next_k_pos = k_start + BLOCK_K + k_offsets @@ -374,9 +400,11 @@ def _sparse_attn_prefill_kernel( m_block = tl.max(scores, axis=1) m_new = tl.maximum(m_i, m_block) - alpha = tl.where(m_new == float("-inf"), 0.0, tl.exp(m_i - m_new)) + alpha = tl.where(m_new == float("-inf"), 0.0, _exp(m_i - m_new, USE_EXP2)) p = tl.where( - m_new[:, None] == float("-inf"), 0.0, tl.exp(scores - m_new[:, None]) + m_new[:, None] == float("-inf"), + 0.0, + _exp(scores - m_new[:, None], USE_EXP2), ) p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0) l_new = l_i * alpha + tl.sum(p, axis=1) @@ -389,9 +417,11 @@ def _sparse_attn_prefill_kernel( sink = tl.load( attn_sink_ptr + head_offsets, mask=head_mask, other=float("-inf") ).to(tl.float32) + if USE_EXP2: + sink = sink * 1.4426950408889634 m_final = tl.maximum(m_i, sink) - alpha = tl.where(m_final == float("-inf"), 0.0, tl.exp(m_i - m_final)) - exp_sink = tl.where(sink == float("-inf"), 0.0, tl.exp(sink - m_final)) + alpha = tl.where(m_final == float("-inf"), 0.0, _exp(m_i - m_final, USE_EXP2)) + exp_sink = tl.where(sink == float("-inf"), 0.0, _exp(sink - m_final, USE_EXP2)) l_final = l_i * alpha + exp_sink denom = tl.maximum(l_final, 1.0e-30) out = tl.where( diff --git a/aiter/ops/triton/attention/pa_prefill_sparse.py b/aiter/ops/triton/attention/pa_prefill_sparse.py index 118378c1c65..c47ee3a4911 100644 --- a/aiter/ops/triton/attention/pa_prefill_sparse.py +++ b/aiter/ops/triton/attention/pa_prefill_sparse.py @@ -18,6 +18,8 @@ else -> triton ``_sparse_attn_prefill_kernel`` (single source) """ +import functools + import torch import triton @@ -32,12 +34,35 @@ ) from aiter.ops.triton.utils._triton import arch_info from aiter.ops.triton.utils.logger import AiterTritonLogger +from aiter.ops.triton.utils.tuned_config_utils import get_tuned_kernel_config DEVICE_ARCH = arch_info.get_arch() _LOGGER = AiterTritonLogger() +# Returned by the lookup where nothing is published for the device and shape. +_NO_PINNED_CONFIG = triton.Config({}) + + +@functools.lru_cache +def _pinned_prefill_config(num_heads: int, head_dim: int) -> triton.Config | None: + """Published single-source prefill tile for this device and shape, if any. + + Entries are keyed ``_sparse_attn_prefill_kernel_H{num_heads}_D{head_dim}`` + in the arch's ``sparse_attention_dsv4/DEFAULT.json``. Devices and shapes + without an entry return None, quietly, and use the autotuned launch. + """ + cfg = get_tuned_kernel_config( + "attention", + "SPARSE_ATTENTION_DSV4", + f"_sparse_attn_prefill_kernel_H{num_heads}_D{head_dim}", + fallback=_NO_PINNED_CONFIG, + warn_on_miss=False, + ) + return None if cfg is _NO_PINNED_CONFIG else cfg + + def pa_prefill_sparse( q: torch.Tensor, unified_kv: torch.Tensor, @@ -49,6 +74,7 @@ def pa_prefill_sparse( attn_sink: torch.Tensor | None, softmax_scale: float, has_invalid: bool | None = None, + out: torch.Tensor | None = None, ) -> torch.Tensor: """Sparse prefill attention over two KV sources with sink. @@ -56,11 +82,17 @@ def pa_prefill_sparse( otherwise the Triton kernel. Only gfx1250 reads a second KV source; the other two serve the prefix source alone and reject a non-empty extend. + Meant for prefill-sized calls. The grid is one program per query token with + no split over the KV slots, so a call with only a few query rows (a decode + step) leaves most of the GPU idle; use a split-KV decode kernel for those. + Args: q: [T, H, D] BF16/FP16 — queries. unified_kv: [total_pages, D] — prefix KV source (paged). kv_indices_prefix: [total_prefix] int32 — flat per-token slot lists - into unified_kv. ``-1`` sentinels skipped. + into unified_kv. The Triton branch also takes int64 as is; gfx950 + casts other dtypes to int32, so slots there must be below 2^31. Invalid-slot handling depends on the branch; see + ``has_invalid``. kv_indptr_prefix: [T+1] int32 — true prefix sum. kv: [total_tokens, D] — extend KV source (this fwd's input K, not yet in paged buffer). ``None`` for no extend source. @@ -69,10 +101,28 @@ def pa_prefill_sparse( kv_indptr_extend: [T+1] int32 — true prefix sum. ``None`` for none. attn_sink: [H] fp32 — per-head softmax-denom bias. softmax_scale: float. + has_invalid: gfx1250 only: whether index lists may hold ``-1`` + sentinels (``None`` picks a heuristic). The Triton fallback always + skips ``-1`` and out-of-pool slots. gfx950 (``mla_gluon``) does not + check slot values: pass only valid slots there. + out: optional contiguous [T, H, D] buffer on q's device, + same dtype as q; written in place and returned. Returns: [T, H, D] attention output, same dtype as q. """ + if out is None: + out = torch.empty_like(q) + else: + assert ( + out.shape == q.shape and out.dtype == q.dtype + ), f"out {tuple(out.shape)} {out.dtype} != q {tuple(q.shape)} {q.dtype}" + assert out.device == q.device, f"out on {out.device}, q on {q.device}" + # Every branch writes [T, H, D] in place: overlapping views (e.g. an + # expand() over heads) would race, and gfx950 ignores out.stride(2). + assert ( + out.is_contiguous() + ), f"out must be contiguous, got strides {out.stride()}" if DEVICE_ARCH == "gfx1250": if not q.is_cuda: raise RuntimeError("pa_prefill_sparse requires CUDA/HIP tensors") @@ -98,7 +148,6 @@ def pa_prefill_sparse( kv_indices_extend.shape[0], ) - out = torch.empty_like(q) assert ( kv_indices_prefix.dtype == torch.int32 and kv_indices_prefix.is_contiguous() ) @@ -178,7 +227,6 @@ def pa_prefill_sparse( kv_indices_extend, kv_indptr_extend, ) - out = torch.empty_like(q) gluon_mla_sparse_prefill( q, # q_nope = combined-D query (RoPE folded in) None, # q_pe unused in prefill mode @@ -195,24 +243,32 @@ def pa_prefill_sparse( else: # Portable Triton fallback. + # Same up-front checks as gfx1250: an fp8 KV pool would otherwise fail + # inside tl.dot at compile time with an opaque CompilationError. + if q.dtype not in (torch.bfloat16, torch.float16): + raise RuntimeError(f"pa_prefill_sparse expects fp16/bf16 q, got {q.dtype}") + if unified_kv.dtype != q.dtype: + raise RuntimeError( + f"unified_kv dtype mismatch: kv={unified_kv.dtype}, q={q.dtype}" + ) + # The Triton kernel takes int32 or int64 indices as passed; converting + # int64 to int32 would silently wrap slots past 2^31 - 1. kv_indices_prefix, kv_indptr_prefix = _prep_single_source( kv_indices_prefix, kv_indptr_prefix, kv, kv_indices_extend, kv_indptr_extend, + index_dtypes=(torch.int32, torch.int64), ) - attn_sink = attn_sink or torch.empty(1, device="cuda", dtype=torch.float32) has_attn_sink = attn_sink is not None + if has_attn_sink: + attn_sink = attn_sink.contiguous() + else: + attn_sink = torch.empty(1, device=q.device, dtype=torch.float32) num_queries, num_heads, head_dim = q.shape block_d = triton.next_power_of_2(head_dim) - out = torch.empty_like(q) - - grid = lambda META: ( - num_queries, - triton.cdiv(num_heads, META["BLOCK_H"]), - ) - _sparse_attn_prefill_kernel[grid]( + args = ( q, unified_kv, kv_indices_prefix, @@ -231,6 +287,35 @@ def pa_prefill_sparse( head_dim, unified_kv.shape[0], float(softmax_scale), + ) + + pinned = _pinned_prefill_config(num_heads, head_dim) + if pinned is not None: + # Published per-shape tile (configs//.../sparse_attention_dsv4). + block_h = pinned.kwargs["BLOCK_H"] + _sparse_attn_prefill_kernel.fn[ + (num_queries, triton.cdiv(num_heads, block_h)) + ]( + *args, + HAS_ATTN_SINK=has_attn_sink, + BLOCK_D=block_d, + # Drop the head/dim masks only; invalid KV rows are still + # masked out of the gather (never loaded), so an empty pool + # or a non-finite unused row is safe. + EVEN_H=num_heads % block_h == 0, + EVEN_D=block_d == head_dim, + num_warps=pinned.num_warps, + num_stages=pinned.num_stages, + **pinned.kwargs, + ) + return out + + grid = lambda META: ( + num_queries, + triton.cdiv(num_heads, META["BLOCK_H"]), + ) + _sparse_attn_prefill_kernel[grid]( + *args, HAS_ATTN_SINK=has_attn_sink, BLOCK_D=block_d, ) @@ -248,9 +333,13 @@ def _prep_single_source( kv: torch.Tensor | None, kv_indices_extend: torch.Tensor | None, kv_indptr_extend: torch.Tensor | None, + index_dtypes: tuple[torch.dtype, ...] = (torch.int32,), ) -> tuple[torch.Tensor, torch.Tensor]: """Normalize the KV pool and indices for the gfx950 / Triton kernels. + Indices and indptr keep their dtype if it is in ``index_dtypes`` and are + converted to int32 otherwise. + Rejects an extend KV source outright: only the gfx1250 gluon kernel reads a second pool. """ @@ -269,11 +358,19 @@ def _prep_single_source( ) return ( - _as_int32_contiguous_1d(kv_indices_prefix), - _as_int32_contiguous_1d(kv_indptr_prefix), + _as_index_contiguous_1d(kv_indices_prefix, index_dtypes), + _as_index_contiguous_1d(kv_indptr_prefix, index_dtypes), ) +def _as_index_contiguous_1d( + x: torch.Tensor, index_dtypes: tuple[torch.dtype, ...] +) -> torch.Tensor: + if x.dtype not in index_dtypes: + return _as_int32_contiguous_1d(x) + return x.reshape(-1).contiguous() + + def _as_int32_contiguous_1d(x: torch.Tensor) -> torch.Tensor: if x.dtype == torch.int32 and x.ndim == 1 and x.is_contiguous(): return x diff --git a/aiter/ops/triton/configs/gfx942/triton/attention/sparse_attention_dsv4/DEFAULT.json b/aiter/ops/triton/configs/gfx942/triton/attention/sparse_attention_dsv4/DEFAULT.json new file mode 100644 index 00000000000..c2cb5a9b8ce --- /dev/null +++ b/aiter/ops/triton/configs/gfx942/triton/attention/sparse_attention_dsv4/DEFAULT.json @@ -0,0 +1,22 @@ +{ + "_sparse_attn_prefill_kernel_H8_D512": { + "BLOCK_H": 16, + "BLOCK_K": 16, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 2, + "USE_EXP2": true, + "num_warps": 1, + "num_stages": 2 + }, + "_sparse_attn_prefill_kernel_H16_D512": { + "BLOCK_H": 16, + "BLOCK_K": 16, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "kpack": 2, + "USE_EXP2": true, + "num_warps": 1, + "num_stages": 2 + } +} diff --git a/aiter/ops/triton/utils/tuned_config_utils.py b/aiter/ops/triton/utils/tuned_config_utils.py index 1ee41c9fe7a..f706d28a1de 100644 --- a/aiter/ops/triton/utils/tuned_config_utils.py +++ b/aiter/ops/triton/utils/tuned_config_utils.py @@ -82,6 +82,7 @@ def get_tuned_kernel_config( kernel_name: str, fallback: triton.Config, backend: str = "triton", + warn_on_miss: bool = True, ) -> triton.Config: """The tile pinned for this device, or ``fallback`` where none is published. @@ -96,6 +97,8 @@ def get_tuned_kernel_config( kernel_name: Key of the kernel's entry within the config file. fallback: Config to register when this device has no published entry. backend: ``"triton"`` or ``"gluon"``. + warn_on_miss: Log a warning when no entry is published. Pass ``False`` + where a miss is expected, e.g. a lookup keyed by shape. """ try: config_path, entry = _get_tuned_kernel_entry( @@ -113,6 +116,8 @@ def get_tuned_kernel_config( ) return fallback if not entry: + if not warn_on_miss: + return fallback logger.warning( "No tuned Triton config for kernel '%s' in '%s'; using fallback %s", kernel_name, diff --git a/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4.py b/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4.py index 34c9ca15a3d..8a16f6a753c 100644 --- a/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4.py +++ b/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4.py @@ -3,8 +3,10 @@ """Benchmark for the DSV4 sparse MLA prefill kernels. -Benchmarks the Triton prefill kernel against the Gluon (CDNA4 / gfx950) prefill -backend and reports the speedup over Triton. The Gluon backend is the unified +Benchmarks the autotuned Triton prefill kernel, the public `pa_prefill_sparse` +entry (which takes the pinned per-arch launch, e.g. on gfx942 for H=8/16 at +D=512), and the Gluon (CDNA4 / gfx950) prefill backend, and reports the speedup +over Triton. The Gluon backend is the unified `mla_gluon(..., has_pe=False)` entry (HAS_PE=False over the shared `_mla_gluon` kernel) — the backend the production entrance uses. When Gluon is unavailable (non-gfx950 arch, or Triton < 3.6), only the Triton perf is reported. @@ -15,6 +17,7 @@ """ import argparse +import functools import torch import triton @@ -22,6 +25,8 @@ from aiter.ops.triton._triton_kernels.attention.sparse_attention_dsv4 import ( _sparse_attn_prefill_kernel as csa_prefill_tl, ) +from aiter.ops.triton.attention.pa_prefill_sparse import pa_prefill_sparse +from aiter.ops.triton.utils._triton import arch_info # The Gluon prefill kernel is opt-in (gfx950 + Triton >= 3.6). Probe it once at # import time; the benchmark falls back to Triton-only when unavailable. @@ -46,7 +51,29 @@ # --------------------------------------------------------------------------- # Bench data builder # --------------------------------------------------------------------------- +def _alloc_kv(num_kv: int, device: str): + """KV pool as a view into a >2 GiB buffer, as vLLM passes it. + + Triton on AMD specializes pointers whose storage is under 2 GiB for buffer + ops; a standalone small allocation would compile a different variant than + the one production runs. + """ + rows = max(num_kv, (2**31 + 2**20) // (HEAD_DIM * 2)) + pool = torch.empty(rows, HEAD_DIM, dtype=torch.bfloat16, device=device) + kv = pool[:num_kv] + kv.normal_() + return kv + + def _build_csr(num_q: int, max_slots: int, max_topk: int, device: str): + """Ragged slot lists: per query a random length in [max_topk // 4, max_topk] + and that many distinct slots in [0, max_slots). + + Vectorized, O(num_q * max_topk) with one host sync: sorted uniform draws over + [0, max_slots - max_topk] plus 0..max_topk-1 are strictly increasing, so the + slots in a row are distinct and stay in range. + """ + assert max_topk <= max_slots, f"{max_topk=} > {max_slots=}" lens = torch.randint( max(1, max_topk // 4), max_topk + 1, @@ -54,12 +81,13 @@ def _build_csr(num_q: int, max_slots: int, max_topk: int, device: str): dtype=torch.int32, device=device, ) - flat, ptr = [], [0] - for i in range(num_q): - L = int(lens[i].item()) - flat.append(torch.randperm(max_slots, device=device, dtype=torch.int32)[:L]) - ptr.append(ptr[-1] + L) - return torch.cat(flat), torch.tensor(ptr, dtype=torch.int32, device=device), lens + draws = torch.randint(0, max_slots - max_topk + 1, (num_q, max_topk), device=device) + slots = draws.sort(dim=1).values + torch.arange(max_topk, device=device) + keep = torch.arange(max_topk, device=device)[None, :] < lens[:, None] + flat = slots[keep].to(torch.int32) + ptr = torch.zeros(num_q + 1, dtype=torch.int32, device=device) + ptr[1:] = lens.cumsum(0) + return flat, ptr, lens def _ref_prefill(q, kv, indices, indptr, scale, attn_sink=None): @@ -86,6 +114,19 @@ def _ref_prefill(q, kv, indices, indptr, scale, attn_sink=None): # --------------------------------------------------------------------------- # Kernel launchers # --------------------------------------------------------------------------- +@functools.lru_cache +def _extend_source(device, dtype, num_queries): + """Second KV source for `pa_prefill_sparse`: gfx1250 requires one, so pass an + empty one there (built once per shape, outside the timed loop).""" + if arch_info.get_arch() != "gfx1250": + return None, None, None + return ( + torch.empty(1, HEAD_DIM, dtype=dtype, device=device), + torch.empty(0, dtype=torch.int32, device=device), + torch.zeros(num_queries + 1, dtype=torch.int32, device=device), + ) + + def _launch_prefill( backend, q, @@ -101,6 +142,18 @@ def _launch_prefill( scale, ): block_d = triton.next_power_of_2(head_dim) + if backend == "pa_prefill_sparse": + pa_prefill_sparse( + q, + kv, + indices, + indptr, + *_extend_source(q.device, q.dtype, num_queries), + attn_sink if has_sink else None, + scale, + out=out, + ) + return if backend == "gluon": mla_gluon( q, # q_nope = combined-D query @@ -208,7 +261,7 @@ def run_prefill_bench(args, device: str): q = torch.randn( num_queries, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device ) - kv = torch.randn(num_kv, HEAD_DIM, dtype=torch.bfloat16, device=device) + kv = _alloc_kv(num_kv, device) indices, indptr, _ = _build_csr(num_queries, num_kv, topk, device) nnz = int(indptr[-1].item()) # number of non-zeros scale = 1.0 / (HEAD_DIM**0.5) @@ -229,6 +282,11 @@ def _perf(ms, bytes_moved=bytes_moved, flops=flops): "triton", q, kv, indices, indptr, num_queries, num_heads, scale ) tri_tflops, tri_gbps = _perf(tri_ms) + pa_ms = _time_backend( + "pa_prefill_sparse", q, kv, indices, indptr, num_queries, num_heads, scale + ) + pa_tflops, _pa_gbps = _perf(pa_ms) + pa_cols = (pa_ms, pa_tflops, f"{tri_ms / pa_ms:.2f}x") if HAS_GLUON: glu_ms = _time_backend( @@ -244,6 +302,7 @@ def _perf(ms, bytes_moved=bytes_moved, flops=flops): topk, tri_ms, tri_tflops, + *pa_cols, glu_ms, glu_tflops, f"{speedup:.2f}x", @@ -251,7 +310,16 @@ def _perf(ms, bytes_moved=bytes_moved, flops=flops): ) else: rows.append( - (num_queries, num_heads, num_kv, topk, tri_ms, tri_tflops, tri_gbps) + ( + num_queries, + num_heads, + num_kv, + topk, + tri_ms, + tri_tflops, + tri_gbps, + *pa_cols, + ) ) if HAS_GLUON: @@ -262,12 +330,26 @@ def _perf(ms, bytes_moved=bytes_moved, flops=flops): "topk", "triton ms", "triton TFLOPS", + "pa_prefill_sparse ms", + "pa_prefill_sparse TFLOPS", + "vs triton", "gluon ms", "gluon TFLOPS", "speedup", ] else: - headers = ["Q", "H", "Kv", "topk", "triton ms", "triton TFLOPS", "triton GB/s"] + headers = [ + "Q", + "H", + "Kv", + "topk", + "triton ms", + "triton TFLOPS", + "triton GB/s", + "pa_prefill_sparse ms", + "pa_prefill_sparse TFLOPS", + "vs triton", + ] _print_table("PREFILL", headers, rows) @@ -304,7 +386,7 @@ def check_correctness(device: str): indices, indptr, _ = _build_csr(num_queries, num_kv, topk, device) scale = 1.0 / (HEAD_DIM**0.5) - backends = ["triton", "gluon"] if HAS_GLUON else ["triton"] + backends = ["triton", "pa_prefill_sparse"] + (["gluon"] if HAS_GLUON else []) for backend in backends: for has_sink in (False, True): attn_sink = ( @@ -355,6 +437,9 @@ def _parse_args(): "4096,128,4096,1024", "8192,128,8192,512", "8192,128,8192,1024", + # DSv4.1-Flash TP4 prefill on gfx942: top-512 + 128 SWA + "16384,16,65536,640", + "16384,8,65536,640", ], ) args = p.parse_args() diff --git a/op_tests/triton_tests/attention/test_pa_prefill_sparse.py b/op_tests/triton_tests/attention/test_pa_prefill_sparse.py index b39166bd5f2..5f513feea89 100644 --- a/op_tests/triton_tests/attention/test_pa_prefill_sparse.py +++ b/op_tests/triton_tests/attention/test_pa_prefill_sparse.py @@ -14,6 +14,16 @@ # --------------------------------------------------------------------------- +def _gather_valid(pool, indices, start, end): + """Rows ``indices[start:end]`` of ``pool`` in fp32, with slots outside + ``[0, len(pool))`` zeroed and reported invalid.""" + slots = indices[start:end].to(torch.int64) + valid = (slots >= 0) & (slots < pool.shape[0]) + rows = torch.zeros(slots.numel(), pool.shape[1], device=pool.device) + rows[valid] = pool[slots[valid]].float() + return rows, valid + + def _sparse_prefill_attn_torch( q, unified_kv, @@ -33,47 +43,35 @@ def _sparse_prefill_attn_torch( unified_kv: [total_pages, D] kv_indices_prefix: [total_prefix] int32 kv_indptr_prefix: [T+1] int32 - kv: [total_tokens, D] - kv_indices_extend: [total_extend] int32 - kv_indptr_extend: [T+1] int32 - attn_sink: [H] fp32 + kv: [total_tokens, D] — ``None`` for no extend source + kv_indices_extend: [total_extend] int32 — ``None`` for no extend source + kv_indptr_extend: [T+1] int32 — ``None`` for no extend source + attn_sink: [H] fp32 — ``None`` for no sink + Slots outside ``[0, len(pool))`` (``-1`` sentinels or past the pool) are + masked out. Returns: [T, H, D] """ T, H, D = q.shape device = q.device + has_extend = kv is not None p_indptr = kv_indptr_prefix.to(torch.int64) - e_indptr = kv_indptr_extend.to(torch.int64) + if has_extend: + e_indptr = kv_indptr_extend.to(torch.int64) out = torch.zeros(T, H, D, dtype=torch.float32, device=device) for t in range(T): - # Gather prefix KV - ps = int(p_indptr[t].item()) - pe = int(p_indptr[t + 1].item()) - p_slots = kv_indices_prefix[ps:pe].to(torch.int64) - p_valid = p_slots >= 0 - p_safe = p_slots.clamp(min=0) - p_kv = unified_kv[p_safe] # [P, D] - p_kv = torch.where( - p_valid[:, None], p_kv.float(), torch.zeros_like(p_kv.float()) - ) - - # Gather extend KV - es = int(e_indptr[t].item()) - ee = int(e_indptr[t + 1].item()) - e_slots = kv_indices_extend[es:ee].to(torch.int64) - e_valid = e_slots >= 0 - e_safe = e_slots.clamp(min=0) - e_kv = kv[e_safe] # [E, D] - e_kv = torch.where( - e_valid[:, None], e_kv.float(), torch.zeros_like(e_kv.float()) + all_kv, all_valid = _gather_valid( + unified_kv, kv_indices_prefix, p_indptr[t], p_indptr[t + 1] ) - - # Concatenate - all_kv = torch.cat([p_kv, e_kv], dim=0) # [K, D] - all_valid = torch.cat([p_valid, e_valid], dim=0) # [K] + if has_extend: + e_kv, e_valid = _gather_valid( + kv, kv_indices_extend, e_indptr[t], e_indptr[t + 1] + ) + all_kv = torch.cat([all_kv, e_kv], dim=0) # [K, D] + all_valid = torch.cat([all_valid, e_valid], dim=0) # [K] K = all_kv.shape[0] if K == 0: @@ -84,15 +82,16 @@ def _sparse_prefill_attn_torch( scores = scores.masked_fill(~all_valid[None, :], float("-inf")) # Add sink as virtual K with V=0 - sink = attn_sink.float().unsqueeze(1) # [H, 1] - combined = torch.cat([scores, sink], dim=-1) # [H, K+1] - cmax = combined.amax(dim=-1, keepdim=True) + if attn_sink is not None: + sink = attn_sink.float().unsqueeze(1) # [H, 1] + scores = torch.cat([scores, sink], dim=-1) # [H, K+1] + cmax = scores.amax(dim=-1, keepdim=True) cmax = torch.where( cmax == float("-inf"), torch.zeros_like(cmax), cmax, ) - weights = (combined - cmax).exp() + weights = (scores - cmax).exp() denom = weights.sum(dim=-1, keepdim=True).clamp(min=1e-30) weights = weights / denom weights_kv = weights[:, :K] # [H, K] @@ -467,3 +466,283 @@ def test_pa_prefill_sparse_gfx950(T, H, D, prefix_len): ) torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2) + + +def _triton_branch_only(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + if DEVICE_ARCH in ("gfx950", "gfx1250"): + pytest.skip("covers the Triton single-source branch") + + +def _inject_invalid_slots(indices, num_kv): + """Mark slots invalid in place: 1/8 to -1, 1/32 to exactly num_kv (the + first row past the pool) and 1/32 to far past it, up to the int32 max.""" + n = indices.numel() + if not n: + return + perm = torch.randperm(n, device=indices.device) + indices[perm[: n // 8]] = -1 + indices[perm[n // 8 : n // 8 + n // 32]] = num_kv + far = perm[n // 8 + n // 32 : n // 8 + 2 * (n // 32)] + indices[far] = torch.randint( + num_kv + 1, + 2**31 - 1, + (far.numel(),), + dtype=indices.dtype, + device=indices.device, + ) + + +# DSv4.1-Flash TP4: H=16 per rank, D=512, top-512 + 128 SWA = 640 slots. +# Lengths vary per row so the last BLOCK_K tile is usually partial. +# H=32 exercises the autotuned launch; H<=16 the gfx942 pinned config, in +# both dtypes since that config is keyed by shape only. +@pytest.mark.parametrize( + "H, dtype", + [ + (8, torch.float16), + (16, torch.bfloat16), + (16, torch.float16), + (32, torch.bfloat16), + ], +) +@pytest.mark.parametrize("with_sink", [True, False]) +def test_pa_prefill_sparse_single_source(H, dtype, with_sink): + _triton_branch_only() + + T, D, max_len = 300, 512, 640 + torch.manual_seed(0) + dev = "cuda" + num_kv = 4096 + q = torch.randn(T, H, D, dtype=dtype, device=dev) + kv = torch.randn(num_kv, D, dtype=dtype, device=dev) + sink = torch.randn(H, dtype=torch.float32, device=dev) if with_sink else None + lens = torch.randint(1, max_len + 1, (T,), device=dev) + lens[::7] = 0 # empty rows: output is 0 (sink only) + lens[1::7] = torch.randint(1, 17, (lens[1::7].numel(),), device=dev) # short rows + indptr = torch.zeros(T + 1, dtype=torch.int32, device=dev) + indptr[1:] = lens.cumsum(0) + indices = torch.randint( + 0, num_kv, (int(indptr[-1]),), dtype=torch.int32, device=dev + ) + # Every other row keeps only valid slots; the rest get invalid ones. + flat_row = torch.repeat_interleave(torch.arange(T, device=dev), lens) + with_invalid = indices.clone() + _inject_invalid_slots(with_invalid, num_kv) + indices = torch.where(flat_row % 2 == 0, with_invalid, indices) + row = int(torch.nonzero((lens >= 32) & (torch.arange(T, device=dev) % 2 == 0))[0]) + indices[int(indptr[row]) : int(indptr[row]) + 16] = -1 # a whole leading tile + scale = D**-0.5 + + ref = _sparse_prefill_attn_torch( + q, kv, indices, indptr, None, None, None, sink, scale + ) + out = torch.full_like(q, float("nan")) + ret = pa_prefill_sparse( + q, + kv, + indices, + indptr, + None, + None, + None, + sink, + scale, + out=out, + ) + assert ret is out + torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2) + + +# Head dims other than 512 take the autotuned launch on every arch; 576 also +# needs the dim mask (BLOCK_D=1024). +def test_pa_prefill_sparse_single_source_other_head_dims(): + _triton_branch_only() + + torch.manual_seed(1) + D = 576 + T, H, num_kv, dev = 257, 16, 2048, "cuda" + q = torch.randn(T, H, D, dtype=torch.bfloat16, device=dev) + kv = torch.randn(num_kv, D, dtype=torch.bfloat16, device=dev) + sink = torch.randn(H, dtype=torch.float32, device=dev) + lens = torch.randint(1, 200, (T,), device=dev) + indptr = torch.zeros(T + 1, dtype=torch.int32, device=dev) + indptr[1:] = lens.cumsum(0) + indices = torch.randint( + 0, num_kv, (int(indptr[-1]),), dtype=torch.int32, device=dev + ) + _inject_invalid_slots(indices, num_kv) + scale = D**-0.5 + + ref = _sparse_prefill_attn_torch( + q, kv, indices, indptr, None, None, None, sink, scale + ) + out = pa_prefill_sparse(q, kv, indices, indptr, None, None, None, sink, scale) + torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2) + + +# An empty pool leaves nothing to gather: every query gets the sink-only (zero) output. +def test_pa_prefill_sparse_empty_kv_pool(): + _triton_branch_only() + T, H, D, dev = 64, 16, 512, "cuda" + q = torch.randn(T, H, D, dtype=torch.bfloat16, device=dev) + kv = torch.empty(0, D, dtype=torch.bfloat16, device=dev) + indptr = torch.arange(0, (T + 1) * 16, 16, dtype=torch.int32, device=dev) + indices = torch.full((T * 16,), -1, dtype=torch.int32, device=dev) + sink = torch.randn(H, dtype=torch.float32, device=dev) + out = pa_prefill_sparse(q, kv, indices, indptr, None, None, None, sink, D**-0.5) + torch.cuda.synchronize() + assert torch.equal(out, torch.zeros_like(out)) + + +# Strided or overlapping out= buffers would be written wrongly (gfx950 ignores +# out.stride(2)) or raced on (an expand() over heads), so they are rejected. +# An fp8 KV pool is rejected up front with a clear error instead of failing to +# compile inside tl.dot. +def test_pa_prefill_sparse_rejects_bad_inputs(): + _triton_branch_only() + T, H, D = 4, 16, 512 + q = torch.randn(T, H, D, dtype=torch.bfloat16, device="cuda") + kv = torch.randn(8, D, dtype=torch.bfloat16, device="cuda") + indptr = torch.arange(0, T + 1, dtype=torch.int32, device="cuda") + indices = torch.zeros(T, dtype=torch.int32, device="cuda") + outs = ( + torch.empty(T, D, H, dtype=q.dtype, device="cuda").transpose(1, 2), + torch.empty(T, 1, D, dtype=q.dtype, device="cuda").expand(T, H, D), + ) + for out in outs: + with pytest.raises(AssertionError, match="contiguous"): + pa_prefill_sparse( + q, kv, indices, indptr, None, None, None, None, D**-0.5, out=out + ) + with pytest.raises(RuntimeError, match="dtype mismatch"): + pa_prefill_sparse( + q, + kv.to(torch.float8_e4m3fnuz), + indices, + indptr, + None, + None, + None, + None, + D**-0.5, + ) + + +# out= on whichever branch this device takes (Triton fallback, gfx950 or +# gfx1250 Gluon): the supplied buffer is returned and holds the same result. +def test_pa_prefill_sparse_out_buffer_every_branch(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + T, H, D = 512, 64, 512 + two_sources = DEVICE_ARCH == "gfx1250" + q, ukv, p_idx, p_indptr, kv, e_idx, e_indptr, sink, scale = _make_inputs( + T, H, D, 128, 128 if two_sources else 1, T * 128, T * 128 if two_sources else 1 + ) + extend = (kv, e_idx, e_indptr) if two_sources else (None, None, None) + ref = pa_prefill_sparse(q, ukv, p_idx, p_indptr, *extend, sink, scale) + out = torch.full_like(q, float("nan")) + ret = pa_prefill_sparse(q, ukv, p_idx, p_indptr, *extend, sink, scale, out=out) + assert ret is out + assert torch.equal(out, ref) + + +# int64 index buffers stay int64 on the Triton branch: slots past 2^31 - 1 must +# be rejected as out of pool, not wrapped into range by an int32 cast. +def test_pa_prefill_sparse_int64_indices(): + _triton_branch_only() + torch.manual_seed(3) + H, T, D, num_kv, dev = 16, 512, 512, 4096, "cuda" + q = torch.randn(T, H, D, dtype=torch.bfloat16, device=dev) + kv = torch.randn(num_kv, D, dtype=torch.bfloat16, device=dev) + sink = torch.randn(H, dtype=torch.float32, device=dev) + lens = torch.randint(1, 640, (T,), device=dev) + indptr = torch.zeros(T + 1, dtype=torch.int64, device=dev) + indptr[1:] = lens.cumsum(0) + indices = torch.randint( + 0, num_kv, (int(indptr[-1]),), dtype=torch.int64, device=dev + ) + n = indices.numel() + perm = torch.randperm(n, device=dev) + indices[perm[: n // 8]] = -1 + # int32 casts of these would wrap to 0, 5 and num_kv - 1: valid-looking rows. + indices[perm[n // 8 : n // 8 + n // 16]] = 2**32 + indices[perm[n // 8 + n // 16 : n // 8 + n // 8]] = 2**32 + 5 + indices[perm[n // 4 : n // 4 + n // 16]] = 2**32 + num_kv - 1 + scale = D**-0.5 + ref = _sparse_prefill_attn_torch( + q, kv, indices, indptr, None, None, None, sink, scale + ) + out = pa_prefill_sparse(q, kv, indices, indptr, None, None, None, sink, scale) + torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2) + + +# The validity check must not rely on num_kv fitting int32. A real pool that +# large cannot be allocated in a test, so expand one KV row to the full row +# count (stride 0, no memory): every slot then reads the same row, and with a +# sink the output still depends on how many slots were accepted, +# out = v * n e^s / (n e^s + e^sink), +# so a negative slot wrongly accepted as valid changes it. +def test_pa_prefill_sparse_huge_pool_rejects_negative_slots(): + _triton_branch_only() + num_rows = 2**31 + 5 + torch.manual_seed(4) + T, H, D, L, dev = 256, 16, 512, 64, "cuda" + q = torch.randn(T, H, D, dtype=torch.bfloat16, device=dev) + row = torch.randn(1, D, dtype=torch.bfloat16, device=dev) + kv = row.expand(num_rows, D) + sink = torch.randn(H, dtype=torch.float32, device=dev) + indptr = torch.arange(0, (T + 1) * L, L, dtype=torch.int32, device=dev) + indices = torch.randint(0, 2**31 - 1, (T * L,), dtype=torch.int32, device=dev) + indices[0::3] = -1 + indices[1::7] = -(2**31) + scale = D**-0.5 + out = pa_prefill_sparse(q, kv, indices, indptr, None, None, None, sink, scale) + + n = (indices.view(T, L) >= 0).sum(1).float()[:, None] # accepted slots per query + s_ = (q.float() @ row.float()[0]) * scale # [T, H]: every key scores the same + m = torch.maximum(s_, sink[None, :]) + w = n * torch.exp(s_ - m) / (n * torch.exp(s_ - m) + torch.exp(sink[None, :] - m)) + ref = (w[..., None] * row.float()[0]).to(q.dtype) + torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2) + + +# A one-row pool: Triton specializes the integer 1 as a constant. +def test_pa_prefill_sparse_one_row_pool(): + _triton_branch_only() + T, H, D, dev = 64, 16, 512, "cuda" + q = torch.randn(T, H, D, dtype=torch.bfloat16, device=dev) + kv = torch.randn(1, D, dtype=torch.bfloat16, device=dev) + sink = torch.randn(H, dtype=torch.float32, device=dev) + indptr = torch.arange(0, (T + 1) * 4, 4, dtype=torch.int32, device=dev) + indices = torch.tensor([0, -1, 1, 0] * T, dtype=torch.int32, device=dev) + ref = _sparse_prefill_attn_torch( + q, kv, indices, indptr, None, None, None, sink, D**-0.5 + ) + out = pa_prefill_sparse(q, kv, indices, indptr, None, None, None, sink, D**-0.5) + torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2) + + +# Invalid slots are masked out of the gather, so a non-finite row the indices +# never select validly (row 0 here) must not reach the output. Loading it and +# giving it zero weight would not be enough: 0 * NaN is NaN in the PV dot. +def test_pa_prefill_sparse_nonfinite_row0_ignored_by_invalid_slots(): + _triton_branch_only() + torch.manual_seed(5) + H, T, D, num_kv, dev = 16, 256, 512, 1024, "cuda" + q = torch.randn(T, H, D, dtype=torch.bfloat16, device=dev) + kv = torch.randn(num_kv, D, dtype=torch.bfloat16, device=dev) + kv[0, ::2] = float("nan") + kv[0, 1::2] = float("inf") + sink = torch.randn(H, dtype=torch.float32, device=dev) + indptr = torch.arange(0, (T + 1) * 48, 48, dtype=torch.int32, device=dev) + indices = torch.randint(1, num_kv, (T * 48,), dtype=torch.int32, device=dev) + indices[0::4] = -1 + indices[1::9] = num_kv + 3 + ref = _sparse_prefill_attn_torch( + q, kv, indices, indptr, None, None, None, sink, D**-0.5 + ) + out = pa_prefill_sparse(q, kv, indices, indptr, None, None, None, sink, D**-0.5) + assert torch.isfinite(out).all() + torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)