diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index ca8f9eee753..919a178e3fb 100644 --- a/backends/cuda/BUCK +++ b/backends/cuda/BUCK @@ -174,6 +174,7 @@ fbcode_target( "triton/kernels/quantized_gemm_family.py", "triton/kernels/quantized_gemm_utils.py", "triton/kernels/sdpa.py", + "triton/kernels/sdpa_varlen.py", "triton/kernels/topk.py", "triton/kernels/tq4_sdpa.py", "triton/kernels/triton_mutation_cache.py", diff --git a/backends/cuda/tests/bench_sdpa_varlen.py b/backends/cuda/tests/bench_sdpa_varlen.py new file mode 100644 index 00000000000..adbe239dae1 --- /dev/null +++ b/backends/cuda/tests/bench_sdpa_varlen.py @@ -0,0 +1,109 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Benchmark triton::sdpa_varlen against the cell layout's triton::sdpa. + +Muse Glimmer 30B attention (H_q=32, H_kv=2, D=128). Each case is a step of +the batching runtime: per sequence, the query rows it adds and the context it +then holds. The cell baseline sweeps one shared pool holding every sequence, +masked per row; sdpa_varlen reads each sequence's own column. + +usage: python bench_sdpa_varlen.py [--check] +""" + +import argparse + +import torch +import triton + +from executorch.backends.cuda.triton.kernels.sdpa import sdpa +from executorch.backends.cuda.triton.kernels.sdpa_varlen import sdpa_varlen + +HQ, HKV, D = 32, 2, 128 +WINDOW = 2048 + +# (name, [(query rows, context after the step)] per sequence) +CASES = [ + ("decode c1 ctx600", [(1, 600)]), + ("decode c4 ctx600", [(1, 600)] * 4), + ("decode c4 ctx4k", [(1, 4096)] * 4), + ("decode c4 mixed", [(1, 300), (1, 900), (1, 2500), (1, 4096)]), + ("verify c1 16x600", [(16, 600)]), + ("verify c4 16x600", [(16, 600)] * 4), + ("verify c4 16x4k", [(16, 4096)] * 4), + ("prefill 97", [(97, 97)]), + ("prefill 512 + 3 decode", [(512, 512), (1, 600), (1, 600), (1, 600)]), + ("prefill 4096", [(4096, 4096)]), +] + + +def _step(case, window, device="cuda"): + t = sum(q for q, _ in case) + cols = len(case) + 1 + rows = max(c for _, c in case) + q = torch.randn(1, HQ, t, D, device=device, dtype=torch.bfloat16) + k_seq = torch.randn(rows, cols, HKV, D, device=device, dtype=torch.bfloat16) + v_seq = torch.randn_like(k_seq) + info, start = [], 0 + for col, (q_len, ctx) in enumerate(case): + info.append((start, q_len, col, ctx)) + start += q_len + info = torch.tensor(info, dtype=torch.int32, device=device) + + # Cell layout: every sequence's context packed into one pool, each row's + # visibility in an explicit mask, the sweep bounded by the occupied extent. + pool = sum(c for _, c in case) + k_pool = torch.empty(1, pool, HKV, D, device=device, dtype=torch.bfloat16) + v_pool = torch.empty_like(k_pool) + mask = torch.zeros(1, 1, t, pool, device=device, dtype=torch.bool) + base, row = 0, 0 + for col, (q_len, ctx) in enumerate(case): + k_pool[0, base : base + ctx] = k_seq[:ctx, col] + v_pool[0, base : base + ctx] = v_seq[:ctx, col] + for r in range(q_len): + p = ctx - q_len + r + lo = max(0, p - window + 1) if window else 0 + mask[0, 0, row + r, base + lo : base + p + 1] = True + base += ctx + row += q_len + read_len = torch.tensor([pool], device=device, dtype=torch.int64) + return q, k_seq, v_seq, info, k_pool, v_pool, mask, read_len + + +def _time(fn): + fn() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fn() + return triton.testing.do_bench(graph.replay, warmup=25, rep=200) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--check", action="store_true") + args = parser.parse_args() + torch.manual_seed(0) + scale = D**-0.5 + print(f"{'case':28s} {'window':>6s} {'cell ms':>9s} {'varlen ms':>10s} {'speedup':>8s}") + for name, case in CASES: + for window in (0, WINDOW): + q, k, v, info, k_pool, v_pool, mask, read_len = _step(case, window) + varlen = lambda: sdpa_varlen(q, k, v, info, scale, True, window, 0) + cell = lambda: sdpa( + q, k_pool.transpose(1, 2), v_pool.transpose(1, 2), mask, 0.0, + False, scale, True, read_len, + ) + if args.check: + torch.testing.assert_close( + varlen().float(), cell().float(), atol=2e-2, rtol=2e-2 + ) + t_cell, t_var = _time(cell), _time(varlen) + print(f"{name:28s} {window:6d} {t_cell:9.4f} {t_var:10.4f} {t_cell / t_var:7.2f}x") + + +if __name__ == "__main__": + main() diff --git a/backends/cuda/tests/targets.bzl b/backends/cuda/tests/targets.bzl index 576943f147b..af1cf582af1 100644 --- a/backends/cuda/tests/targets.bzl +++ b/backends/cuda/tests/targets.bzl @@ -133,6 +133,25 @@ def define_common_targets(is_fbcode = False): ), ) + python_unittest_remote_gpu( + name = "test_sdpa_varlen", + srcs = [ + "test_sdpa_varlen.py", + ], + visibility = [ + "//executorch/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/cuda:triton_kernels", + ], + keep_gpu_sections = True, + remote_execution = re_test_utils.remote_execution( + platform = "gpu-remote-execution", + subplatform = "A100-exclusive", + ), + ) + python_unittest_remote_gpu( name = "test_int4_large_quantized_gemm", srcs = [ diff --git a/backends/cuda/tests/test_sdpa_varlen.py b/backends/cuda/tests/test_sdpa_varlen.py new file mode 100644 index 00000000000..dfd4944753c --- /dev/null +++ b/backends/cuda/tests/test_sdpa_varlen.py @@ -0,0 +1,280 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import itertools +import unittest + +import torch + +from executorch.backends.cuda.triton.kernels.sdpa_varlen import sdpa_varlen + + +def _reference(q, k_store, v_store, segments, scale, causal, window, ring): + """Per-segment attention over logical positions, in float32.""" + _, hq, t_total, d = q.shape + hkv = k_store.shape[2] + out = torch.zeros_like(q, dtype=torch.float32) + for q_start, q_len, col, kv_end in segments: + if q_len == 0: + continue + positions = torch.arange(kv_end, device=q.device) + rows = positions % ring if ring else positions + keep = positions > kv_end - 1 - ring if ring else positions >= 0 + k = k_store[rows, col].float() # [kv, H, D] + v = v_store[rows, col].float() + k = k.repeat_interleave(hq // hkv, dim=1).transpose(0, 1) + v = v.repeat_interleave(hq // hkv, dim=1).transpose(0, 1) + qs = q[0, :, q_start : q_start + q_len].float() + scores = qs @ k.transpose(1, 2) * scale # [H, q, kv] + query_pos = torch.arange(kv_end - q_len, kv_end, device=q.device) + visible = keep[None, :].expand(q_len, -1).clone() + if causal: + visible &= positions[None, :] <= query_pos[:, None] + if window: + visible &= positions[None, :] > query_pos[:, None] - window + scores = scores.masked_fill(~visible, float("-inf")) + out[0, :, q_start : q_start + q_len] = torch.softmax(scores, -1) @ v + return out + + +@unittest.skipUnless(torch.cuda.is_available(), "CUDA required") +class SdpaVarlenTest(unittest.TestCase): + def _run(self, segments, *, hq=8, hkv=2, d=64, rows=512, cols=5, t=None, + causal=True, window=0, ring=0, seed=0): + torch.manual_seed(seed) + dev = "cuda" + t = t or max(s + n for s, n, _, _ in segments) + q = torch.randn(1, hq, t, d, device=dev, dtype=torch.bfloat16) + k = torch.randn(rows, cols, hkv, d, device=dev, dtype=torch.bfloat16) + v = torch.randn(rows, cols, hkv, d, device=dev, dtype=torch.bfloat16) + info = torch.tensor(segments, dtype=torch.int32, device=dev) + scale = d**-0.5 + got = sdpa_varlen(q, k, v, info, scale, causal, window, ring) + want = _reference(q, k, v, segments, scale, causal, window, ring) + torch.testing.assert_close(got.float(), want, atol=2e-2, rtol=2e-2) + covered = torch.zeros(t, dtype=torch.bool) + for s, n, _, _ in segments: + covered[s : s + n] = True + self.assertTrue((got[0, :, ~covered] == 0).all()) + + def test_decode_rows_see_only_their_sequence(self): + for kv in (1, 63, 64, 200, 511): + self._run([(0, 1, 3, kv), (1, 1, 0, 37), (2, 1, 1, 512), (3, 1, 4, 5)]) + + def test_prefill_and_decode_mix(self): + self._run([(0, 100, 2, 150), (100, 1, 0, 300), (101, 30, 4, 30)]) + + def test_padding_segments_and_rows(self): + self._run([(0, 3, 1, 90), (0, 0, 0, 0), (3, 2, 2, 7)], t=16) + + def test_one_owner_split_across_segments(self): + # The second segment's rows see the first's, written in the same step. + self._run([(0, 4, 2, 44), (4, 3, 1, 10), (7, 5, 2, 49)]) + + def test_sliding_window(self): + for window in (1, 16, 100): + self._run([(0, 64, 0, 200), (64, 1, 3, 300)], window=window) + + def test_ring(self): + ring = 96 + for kv_end in (40, 96, 97, 300): + self._run( + [(0, 8, 1, kv_end), (8, 1, 2, max(kv_end - 3, 1))], + window=64, ring=ring, rows=ring, + ) + + def test_non_causal_draft_block(self): + self._run([(0, 16, 0, 80), (16, 16, 3, 16)], causal=False) + + def test_head_layouts(self): + for hq, hkv, d in itertools.product((4, 16), (1, 4), (64, 128)): + if hq % hkv: + continue + self._run([(0, 2, 0, 130), (2, 9, 1, 9)], hq=hq, hkv=hkv, d=d) + + def test_graph_replay_reads_new_lengths(self): + torch.manual_seed(1) + dev = "cuda" + q = torch.randn(1, 8, 4, 64, device=dev, dtype=torch.bfloat16) + k = torch.randn(256, 3, 2, 64, device=dev, dtype=torch.bfloat16) + v = torch.randn_like(k) + info = torch.tensor([[0, 1, 0, 10], [1, 1, 1, 20], [2, 2, 2, 30]], + dtype=torch.int32, device=dev) + sdpa_varlen(q, k, v, info, 0.125) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = sdpa_varlen(q, k, v, info, 0.125) + for segments in ([[0, 1, 0, 200], [1, 1, 1, 3], [2, 2, 2, 256]], + [[0, 4, 1, 100], [0, 0, 0, 0], [0, 0, 0, 0]]): + info.copy_(torch.tensor(segments, dtype=torch.int32)) + graph.replay() + want = _reference(q, k, v, segments, 0.125, True, 0, 0) + torch.testing.assert_close(out.float(), want, atol=2e-2, rtol=2e-2) + + def test_small_t_split_path(self): + production = [(0, 1, 0, 600), (1, 1, 1, 900), + (2, 1, 2, 2500), (3, 1, 3, 4096)] + for window in (0, 2048): + self._run( + production, + hq=32, + hkv=2, + d=128, + rows=4096, + cols=4, + t=4, + window=window, + ) + self._run( + [(0, 3, 1, 130), (3, 1, 0, 4096), (0, 0, 2, 0)], + hq=32, + hkv=2, + d=128, + rows=4096, + t=4, + causal=False, + ) + self._run( + [(0, 1, 1, 600), (2, 2, 0, 4096), (0, 0, 2, 0)], + hq=32, + hkv=2, + d=128, + rows=4096, + t=4, + ) + + def test_small_t_ring_and_uncovered_rows(self): + self._run( + [(0, 1, 1, 300), (0, 0, 0, 0), (2, 2, 2, 97)], + hq=32, + hkv=2, + d=128, + rows=96, + t=4, + window=64, + ring=96, + ) + + def test_overprovisioned_store_graph_replay(self): + torch.manual_seed(2) + dev = "cuda" + q = torch.randn(1, 8, 4, 64, device=dev, dtype=torch.bfloat16) + k = torch.randn(8192, 4, 2, 64, device=dev, dtype=torch.bfloat16) + v = torch.randn_like(k) + info = torch.tensor( + [[0, 1, 0, 100], [1, 1, 1, 600], [2, 1, 2, 900], [3, 1, 3, 1000]], + dtype=torch.int32, + device=dev, + ) + sdpa_varlen(q, k, v, info, 0.125) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = sdpa_varlen(q, k, v, info, 0.125) + scenarios = ( + [[0, 1, 0, 100], [1, 1, 1, 600], [2, 1, 2, 900], [3, 1, 3, 1000]], + [[0, 1, 3, 4096], [1, 1, 2, 17], [2, 1, 1, 8000], [3, 1, 0, 700]], + ) + for segments in scenarios: + info.copy_(torch.tensor(segments, dtype=torch.int32)) + graph.replay() + want = _reference(q, k, v, segments, 0.125, True, 0, 0) + torch.testing.assert_close(out.float(), want, atol=2e-2, rtol=2e-2) + + def test_rows_available(self): + torch.manual_seed(3) + dev = "cuda" + hq, hkv, d, available = 8, 2, 64, 128 + q = torch.randn(1, hq, 1, d, device=dev, dtype=torch.bfloat16) + k = torch.randn(512, 3, hkv, d, device=dev, dtype=torch.bfloat16) + v = torch.randn_like(k) + info = torch.tensor([[0, 1, 2, 300]], dtype=torch.int32, device=dev) + got = sdpa_varlen(q, k, v, info, d**-0.5, rows_available=available) + + keys = k[:available, 2].float().repeat_interleave(hq // hkv, dim=1) + values = v[:available, 2].float().repeat_interleave(hq // hkv, dim=1) + keys = keys.transpose(0, 1) + values = values.transpose(0, 1) + scores = q[0, :, 0].float()[:, None, :] @ keys.transpose(1, 2) + want = torch.softmax(scores * (d**-0.5), -1) @ values + torch.testing.assert_close( + got[0, :, 0].float(), want[:, 0], atol=2e-2, rtol=2e-2 + ) + + def test_torch_compile(self): + torch.manual_seed(4) + dev = "cuda" + q = torch.randn(1, 8, 4, 64, device=dev, dtype=torch.bfloat16) + k = torch.randn(512, 2, 2, 64, device=dev, dtype=torch.bfloat16) + v = torch.randn_like(k) + segments = [(0, 1, 0, 300), (1, 3, 1, 500)] + info = torch.tensor(segments, dtype=torch.int32, device=dev) + + def fn(q, k, v, info): + return sdpa_varlen(q, k, v, info, 0.125, True, 256, 0) + + got = torch.compile(fn, fullgraph=True)(q, k, v, info) + want = _reference(q, k, v, segments, 0.125, True, 256, 0) + torch.testing.assert_close(got.float(), want, atol=2e-2, rtol=2e-2) + + def test_torch_compile_generic(self): + torch.manual_seed(5) + dev = "cuda" + q = torch.randn(1, 8, 16, 64, device=dev, dtype=torch.bfloat16) + k = torch.randn(512, 2, 2, 64, device=dev, dtype=torch.bfloat16) + v = torch.randn_like(k) + segments = [(0, 16, 0, 500), (0, 0, 1, 0)] + info = torch.tensor(segments, dtype=torch.int32, device=dev) + + def fn(q, k, v, info): + return sdpa_varlen(q, k, v, info, 0.125, True, 256, 0) + + got = torch.compile(fn, fullgraph=True)(q, k, v, info) + want = _reference(q, k, v, segments, 0.125, True, 256, 0) + torch.testing.assert_close(got.float(), want, atol=2e-2, rtol=2e-2) + + def test_dynamic_query_length_is_not_specialized(self): + torch.manual_seed(6) + dev = "cuda" + k = torch.randn(512, 3, 2, 64, device=dev, dtype=torch.bfloat16) + v = torch.randn_like(k) + compiled = torch.compile( + lambda q, k, v, info: sdpa_varlen(q, k, v, info, 0.125, False), + fullgraph=True, + dynamic=True, + ) + torch._dynamo.reset() + frames = torch._dynamo.utils.counters["stats"] + for owners in (2, 3, 4): + t = owners * 16 + q = torch.randn(1, 8, t, 64, device=dev, dtype=torch.bfloat16) + segments = [(16 * i, 16, i % 3, 40 + 16 * i) for i in range(owners)] + # The segment table is fixed-size in a lowered step; pad it. + padded = segments + [(0, 0, 0, 0)] * (4 - owners) + info = torch.tensor(padded, dtype=torch.int32, device=dev) + got = compiled(q, k, v, info) + want = _reference(q, k, v, segments, 0.125, False, 0, 0) + torch.testing.assert_close(got.float(), want, atol=2e-2, rtol=2e-2) + self.assertEqual(frames["unique_graphs"], 1) + + def test_rejects_unsupported_layouts(self): + dev = "cuda" + q = torch.randn(1, 8, 1, 64, device=dev, dtype=torch.bfloat16) + k = torch.randn(128, 2, 2, 64, device=dev, dtype=torch.bfloat16) + v = torch.randn_like(k) + info = torch.zeros((1, 8), device=dev, dtype=torch.int32)[:, ::2] + with self.assertRaisesRegex(ValueError, "seq_info must be contiguous"): + sdpa_varlen(q, k, v, info, 0.125) + + q = torch.randn(1, 8, 1, 96, device=dev, dtype=torch.bfloat16) + k = torch.randn(128, 2, 2, 96, device=dev, dtype=torch.bfloat16) + v = torch.randn_like(k) + info = torch.tensor([[0, 1, 0, 100]], dtype=torch.int32, device=dev) + with self.assertRaisesRegex(ValueError, "head dimension must be a power"): + sdpa_varlen(q, k, v, info, 96**-0.5) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/cuda/triton/kernels/sdpa_varlen.py b/backends/cuda/triton/kernels/sdpa_varlen.py new file mode 100644 index 00000000000..8e80a025864 --- /dev/null +++ b/backends/cuda/triton/kernels/sdpa_varlen.py @@ -0,0 +1,545 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +# +# Optimized by KernelAgent-Oink(https://github.com/meta-pytorch/KernelAgent) + +"""Length-aware SDPA over a batched-sequence off-graph KV cache. + +A step packs the tokens of several sequences on one axis; each sequence keeps +its K/V in its own column of a ``[rows, columns, H_kv, D]`` storage, row = its +logical position (or ``position % ring`` for a sliding-window ring). Each +sequence's query rows attend only that sequence's column, bounded by its own +length, so a step costs the sum of the contexts rather than a sweep of every +sequence's cells. + +The runtime describes the step in ``seq_info`` (``Int32[S, 4]``), one row per +segment of consecutive query rows of one sequence: + +* ``q_start``, ``q_len``: the segment's rows on the packed query axis + (``q_len == 0`` marks an unused row); +* ``col``: the sequence's storage column; +* ``kv_end``: one past the logical position of the segment's last row, i.e. the + sequence's length once this step's writes land. + +Rows are the segment's newest positions: row ``r`` sits at +``kv_end - q_len + r``. Active segments must be disjoint and in bounds; +``seq_info`` is read on the device, so a captured CUDA graph replays it with new +contents. Head dimension ``D`` must be a power of two in ``[16, 256]``. +""" + +import torch +import triton +import triton.language as tl +from torch.library import triton_op, wrap_triton + +_SPLITK_MAX_QUERY = 4 +_SPLITK_ROWS_PER_SPLIT = 256 +_SPLITK_MAX_SPLITS = 16 + + +@triton.jit +def _sdpa_varlen_kernel( # noqa: C901 + Q_ptr, + K_ptr, + V_ptr, + O_ptr, + O_part_ptr, + M_part_ptr, + L_part_ptr, + Info_ptr, + Hq, + T, + stride_qh, + stride_qt, + stride_kr, + stride_kc, + stride_kh, + stride_vr, + stride_vc, + stride_vh, + stride_oh, + stride_ot, + sm_scale, + rows_available, + num_qb, + NUM_SPLITS: tl.constexpr, + IS_CAUSAL: tl.constexpr, + WINDOW: tl.constexpr, + RING: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + HEAD_DIM: tl.constexpr, + NUM_GROUPS: tl.constexpr, +): + sm_scale = sm_scale.to(tl.float32) + pid0 = tl.program_id(0) + h_kv = tl.program_id(1) + split = tl.program_id(2) + seg = pid0 // num_qb + qb = pid0 % num_qb + + q_start = tl.load(Info_ptr + seg * 4 + 0) + q_len = tl.load(Info_ptr + seg * 4 + 1) + col = tl.load(Info_ptr + seg * 4 + 2) + kv_end = tl.load(Info_ptr + seg * 4 + 3) + if qb * BLOCK_M >= q_len * NUM_GROUPS: + return + + offs_p = qb * BLOCK_M + tl.arange(0, BLOCK_M) + r = offs_p // NUM_GROUPS + h_q = h_kv * NUM_GROUPS + offs_p % NUM_GROUPS + row_valid = r < q_len + t = q_start + r + pos = kv_end - q_len + r + offs_d = tl.arange(0, HEAD_DIM) + + q = tl.load( + Q_ptr + h_q[:, None] * stride_qh + t[:, None] * stride_qt + offs_d[None, :], + mask=row_valid[:, None], + other=0.0, + ) + + # The rows this tile may see, before the split divides them. + r_last = tl.minimum(q_len - 1, (qb * BLOCK_M + BLOCK_M - 1) // NUM_GROUPS) + r_first = (qb * BLOCK_M) // NUM_GROUPS + if RING > 0: + lo = 0 + hi = tl.minimum(kv_end, RING) + else: + hi = kv_end + if IS_CAUSAL: + hi = kv_end - q_len + r_last + 1 + lo = 0 + if WINDOW > 0: + lo = tl.maximum(kv_end - q_len + r_first - WINDOW + 1, 0) + hi = tl.minimum(hi, rows_available) + span = tl.maximum(hi - lo, 0) + chunk = tl.cdiv(tl.cdiv(span, NUM_SPLITS), BLOCK_N) * BLOCK_N + start = lo + split * chunk + stop = tl.minimum(start + chunk, hi) + + m_i = tl.full([BLOCK_M], -float("inf"), dtype=tl.float32) + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32) + k_base = K_ptr + col * stride_kc + h_kv * stride_kh + v_base = V_ptr + col * stride_vc + h_kv * stride_vh + + for n0 in tl.range(start, stop, BLOCK_N): + offs_n = n0 + tl.arange(0, BLOCK_N) + n_valid = offs_n < stop + k = tl.load( + k_base + offs_n[:, None] * stride_kr + offs_d[None, :], + mask=n_valid[:, None], + other=0.0, + ) + qk = tl.dot(q, tl.trans(k)).to(tl.float32) * sm_scale + if RING > 0: + # The newest logical position that landed in each ring row. + kv_pos = offs_n + RING * ((kv_end - 1 - offs_n) // RING) + else: + kv_pos = offs_n + visible = n_valid[None, :] & row_valid[:, None] & (kv_pos[None, :] >= 0) + if IS_CAUSAL: + visible = visible & (kv_pos[None, :] <= pos[:, None]) + else: + visible = visible & (kv_pos[None, :] < kv_end) + if WINDOW > 0: + visible = visible & (kv_pos[None, :] > pos[:, None] - WINDOW) + qk = tl.where(visible, qk, -float("inf")) + + m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) + p = tl.where(m_ij[:, None] > -float("inf"), tl.exp(qk - m_ij[:, None]), 0.0) + alpha = tl.where(m_ij > -float("inf"), tl.exp(m_i - m_ij), 0.0) + v = tl.load( + v_base + offs_n[:, None] * stride_vr + offs_d[None, :], + mask=n_valid[:, None], + other=0.0, + ) + acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v) + l_i = l_i * alpha + tl.sum(p, axis=1) + m_i = m_ij + + if NUM_SPLITS == 1: + out = acc * tl.where(l_i > 0, 1.0 / l_i, 0.0)[:, None] + tl.store( + O_ptr + h_q[:, None] * stride_oh + t[:, None] * stride_ot + offs_d[None, :], + out.to(O_ptr.dtype.element_ty), + mask=row_valid[:, None], + ) + else: + part = (split * Hq + h_q) * T + t + tl.store( + O_part_ptr + part[:, None] * HEAD_DIM + offs_d[None, :], + acc, + mask=row_valid[:, None], + ) + tl.store(M_part_ptr + part, m_i, mask=row_valid) + tl.store(L_part_ptr + part, l_i, mask=row_valid) + + +@triton.jit +def _sdpa_varlen_decode_kernel( + Q_ptr, + K_ptr, + V_ptr, + O_ptr, + O_part_ptr, + M_part_ptr, + L_part_ptr, + Info_ptr, + Hq, + T, + stride_qh, + stride_qt, + stride_kr, + stride_kc, + stride_kh, + stride_vr, + stride_vc, + stride_vh, + stride_oh, + stride_ot, + sm_scale, + rows_available, + NUM_SEGMENTS: tl.constexpr, + NUM_SPLITS: tl.constexpr, + IS_CAUSAL: tl.constexpr, + WINDOW: tl.constexpr, + RING: tl.constexpr, + BLOCK_N: tl.constexpr, + HEAD_DIM: tl.constexpr, + NUM_GROUPS: tl.constexpr, + BLOCK_G: tl.constexpr, +): + sm_scale = sm_scale.to(tl.float32) + t = tl.program_id(0) + h_kv = tl.program_id(1) + split = tl.program_id(2) + + owned = False + q_start = 0 + q_len = 0 + col = 0 + kv_end = 0 + for seg in range(NUM_SEGMENTS): + seg_q_start = tl.load(Info_ptr + seg * 4) + seg_q_len = tl.load(Info_ptr + seg * 4 + 1) + take = (t >= seg_q_start) & (t < seg_q_start + seg_q_len) + owned |= take + q_start = tl.where(take, seg_q_start, q_start) + q_len = tl.where(take, seg_q_len, q_len) + col = tl.where(take, tl.load(Info_ptr + seg * 4 + 2), col) + kv_end = tl.where(take, tl.load(Info_ptr + seg * 4 + 3), kv_end) + + offs_g = tl.arange(0, BLOCK_G) + offs_d = tl.arange(0, HEAD_DIM) + g_valid = offs_g < NUM_GROUPS + row_valid = owned & g_valid + h_q = h_kv * NUM_GROUPS + offs_g + q = tl.load( + Q_ptr + h_q[:, None] * stride_qh + t * stride_qt + offs_d[None, :], + mask=row_valid[:, None], + other=0.0, + ) + + pos = kv_end - q_len + t - q_start + if RING > 0: + lo = 0 + hi = tl.minimum(kv_end, RING) + else: + hi = pos + 1 if IS_CAUSAL else kv_end + lo = 0 + if WINDOW > 0: + lo = tl.maximum(pos - WINDOW + 1, 0) + hi = tl.minimum(hi, rows_available) + span = tl.maximum(hi - lo, 0) + chunk = tl.cdiv(tl.cdiv(span, NUM_SPLITS), BLOCK_N) * BLOCK_N + start = lo + split * chunk + stop = tl.minimum(start + chunk, hi) + + m_i = tl.full([BLOCK_G], -float("inf"), dtype=tl.float32) + l_i = tl.zeros([BLOCK_G], dtype=tl.float32) + acc = tl.zeros([BLOCK_G, HEAD_DIM], dtype=tl.float32) + k_base = K_ptr + col * stride_kc + h_kv * stride_kh + v_base = V_ptr + col * stride_vc + h_kv * stride_vh + for n0 in tl.range(start, stop, BLOCK_N): + offs_n = n0 + tl.arange(0, BLOCK_N) + n_valid = offs_n < stop + k = tl.load( + k_base + offs_n[:, None] * stride_kr + offs_d[None, :], + mask=n_valid[:, None] & owned, + other=0.0, + ) + qk = tl.dot(q, tl.trans(k)).to(tl.float32) * sm_scale + if RING > 0: + kv_pos = offs_n + RING * ((kv_end - 1 - offs_n) // RING) + else: + kv_pos = offs_n + visible = n_valid[None, :] & row_valid[:, None] & (kv_pos[None, :] >= 0) + if IS_CAUSAL: + visible &= kv_pos[None, :] <= pos + else: + visible &= kv_pos[None, :] < kv_end + if WINDOW > 0: + visible &= kv_pos[None, :] > pos - WINDOW + qk = tl.where(visible, qk, -float("inf")) + + m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) + p = tl.where(m_ij[:, None] > -float("inf"), tl.exp(qk - m_ij[:, None]), 0.0) + alpha = tl.where(m_ij > -float("inf"), tl.exp(m_i - m_ij), 0.0) + v = tl.load( + v_base + offs_n[:, None] * stride_vr + offs_d[None, :], + mask=n_valid[:, None] & owned, + other=0.0, + ) + acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v) + l_i = l_i * alpha + tl.sum(p, axis=1) + m_i = m_ij + + if NUM_SPLITS == 1: + out = acc * tl.where(l_i > 0, 1.0 / l_i, 0.0)[:, None] + tl.store( + O_ptr + h_q[:, None] * stride_oh + t * stride_ot + offs_d[None, :], + out.to(O_ptr.dtype.element_ty), + mask=g_valid[:, None], + ) + else: + part = (split * Hq + h_q) * T + t + tl.store( + O_part_ptr + part[:, None] * HEAD_DIM + offs_d[None, :], + acc, + mask=g_valid[:, None], + ) + tl.store(M_part_ptr + part, m_i, mask=g_valid) + tl.store(L_part_ptr + part, l_i, mask=g_valid) + + +@triton.jit +def _sdpa_varlen_reduce_kernel( + O_part_ptr, + M_part_ptr, + L_part_ptr, + O_ptr, + Hq, + T, + stride_oh, + stride_ot, + NUM_SPLITS: tl.constexpr, + BLOCK_S: tl.constexpr, + HEAD_DIM: tl.constexpr, +): + pid = tl.program_id(0) + h_q = pid // T + t = pid % T + offs_s = tl.arange(0, BLOCK_S) + s_valid = offs_s < NUM_SPLITS + part = (offs_s * Hq + h_q) * T + t + m = tl.load(M_part_ptr + part, mask=s_valid, other=-float("inf")) + m_max = tl.max(m, axis=0) + alpha = tl.where(s_valid & (m > -float("inf")), tl.exp(m - m_max), 0.0) + l_sum = tl.sum(tl.load(L_part_ptr + part, mask=s_valid, other=0.0) * alpha, 0) + offs_d = tl.arange(0, HEAD_DIM) + o = tl.load( + O_part_ptr + part[:, None] * HEAD_DIM + offs_d[None, :], + mask=s_valid[:, None], + other=0.0, + ) + out = tl.sum(o * alpha[:, None], axis=0) * tl.where(l_sum > 0, 1.0 / l_sum, 0.0) + tl.store( + O_ptr + h_q * stride_oh + t * stride_ot + offs_d, + out.to(O_ptr.dtype.element_ty), + ) + + +def _num_splits(T: int, rows: int) -> int: + if T > _SPLITK_MAX_QUERY: + return 1 + return max(1, min(_SPLITK_MAX_SPLITS, triton.cdiv(rows, _SPLITK_ROWS_PER_SPLIT))) + + +@triton_op("triton::sdpa_varlen", mutates_args={}) +def sdpa_varlen( # noqa: C901 + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + seq_info: torch.Tensor, + scale: float, + is_causal: bool = True, + window: int = 0, + ring: int = 0, + rows_available: int = -1, +) -> torch.Tensor: + """Attend each segment of ``query`` to its own sequence's K/V. + + Args: + query: ``[1, H_q, T, D]`` bfloat16, the step's queries in packed order. + key, value: ``[rows, columns, H_kv, D]`` bfloat16 storage; one column per + sequence. Only the head dimension must be contiguous. + seq_info: ``Int32[S, 4]`` on the query's device, ``(q_start, q_len, + col, kv_end)`` per segment; see the module docstring. + scale: softmax scale. + is_causal: row ``r`` sees positions up to its own; otherwise every + row sees all of ``[0, kv_end)``. + window: a sliding window: a row sees only its last ``window`` + positions. Zero for full attention. + ring: the storage holds a ring of this many rows, indexed by + ``position % ring``. Zero when the row is the position. + rows_available: rows actually backed by storage, when fewer than its + declared extent; negative means all of them. + Returns: + ``[1, H_q, T, D]``; rows no segment covers are zero. + """ + if query.dim() != 4 or key.dim() != 4 or value.dim() != 4: + raise ValueError("query, key, and value must be rank 4") + if query.shape[0] != 1: + raise ValueError("query must have shape [1, H_q, T, D]") + if key.shape != value.shape: + raise ValueError("key and value must have the same shape") + _, Hq, T, D = query.shape + rows, _, Hkv, Dkv = key.shape + if D != Dkv: + raise ValueError("query, key, and value must have the same head dimension") + if ( + query.dtype != torch.bfloat16 + or key.dtype != torch.bfloat16 + or value.dtype != torch.bfloat16 + ): + raise ValueError("query, key, and value must be bfloat16") + if ( + query.device.type != "cuda" + or key.device != query.device + or value.device != query.device + ): + raise ValueError("query, key, and value must be on the same CUDA device") + if Hkv == 0 or Hq % Hkv != 0: + raise ValueError(f"H_q={Hq} is not a multiple of H_kv={Hkv}") + if D < 16 or D > 256 or D & (D - 1): + raise ValueError("head dimension must be a power of two in [16, 256]") + if query.stride(-1) != 1 or key.stride(-1) != 1 or value.stride(-1) != 1: + raise ValueError("sdpa_varlen needs a contiguous head dimension") + if seq_info.dtype != torch.int32 or seq_info.dim() != 2 or seq_info.shape[1] != 4: + raise ValueError("seq_info must be Int32[S, 4]") + if seq_info.device != query.device or not seq_info.is_contiguous(): + raise ValueError("seq_info must be contiguous on the query device") + groups = Hq // Hkv + available = rows if rows_available < 0 else min(rows_available, rows) + # Branching on a symbolic T would guard (and AOTI then pin) the exported + # range, so a dynamic step takes the generic path, valid for any T. + static = isinstance(T, int) + use_decode = static and T <= _SPLITK_MAX_QUERY + block_m = 16 if static and T * groups <= 16 else 64 + num_qb = (T * groups + block_m - 1) // block_m + generic_num_stages = 3 if static and T <= 64 else 2 + splits = _num_splits(T, rows) if static else 1 + # The row-complete decode path writes uncovered rows and empty split slots. + out = torch.empty_like(query) if use_decode else torch.zeros_like(query) + if splits > 1: + o_part = torch.empty( + (splits, Hq, T, D), device=query.device, dtype=torch.float32 + ) + m_part = torch.empty( + (splits, Hq, T), device=query.device, dtype=torch.float32 + ) + l_part = torch.empty((splits, Hq, T), device=query.device, dtype=torch.float32) + else: + o_part = m_part = l_part = out + if use_decode: + grid = (T, Hkv, splits) + wrap_triton(_sdpa_varlen_decode_kernel)[grid]( + query, + key, + value, + out, + o_part, + m_part, + l_part, + seq_info, + Hq, + T, + query.stride(1), + query.stride(2), + key.stride(0), + key.stride(1), + key.stride(2), + value.stride(0), + value.stride(1), + value.stride(2), + out.stride(1), + out.stride(2), + scale, + available, + NUM_SEGMENTS=seq_info.shape[0], + NUM_SPLITS=splits, + IS_CAUSAL=is_causal, + WINDOW=window, + RING=ring, + BLOCK_N=128, + HEAD_DIM=D, + NUM_GROUPS=groups, + BLOCK_G=triton.next_power_of_2(groups), + num_warps=4, + num_stages=2, + ) + else: + grid = (seq_info.shape[0] * num_qb, Hkv, splits) + wrap_triton(_sdpa_varlen_kernel)[grid]( + query, + key, + value, + out, + o_part, + m_part, + l_part, + seq_info, + Hq, + T, + query.stride(1), + query.stride(2), + key.stride(0), + key.stride(1), + key.stride(2), + value.stride(0), + value.stride(1), + value.stride(2), + out.stride(1), + out.stride(2), + scale, + available, + num_qb, + NUM_SPLITS=splits, + IS_CAUSAL=is_causal, + WINDOW=window, + RING=ring, + BLOCK_M=block_m, + BLOCK_N=64, + HEAD_DIM=D, + NUM_GROUPS=groups, + num_warps=4, + num_stages=generic_num_stages, + ) + if splits > 1: + wrap_triton(_sdpa_varlen_reduce_kernel)[(Hq * T,)]( + o_part, + m_part, + l_part, + out, + Hq, + T, + out.stride(1), + out.stride(2), + NUM_SPLITS=splits, + BLOCK_S=triton.next_power_of_2(splits), + HEAD_DIM=D, + ) + return out + + +@sdpa_varlen.register_fake +def _sdpa_varlen_fake( + query, key, value, seq_info, scale, is_causal=True, window=0, ring=0, + rows_available=-1, +): + return torch.empty_like(query)