From 020a10c49f3836421b1169057e89a1fc2d752b4a Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Mon, 9 Nov 2026 21:05:24 -0800 Subject: [PATCH] [executorch][cuda] Add a length-aware Triton SDPA over per-sequence KV columns Pull Request resolved: https://github.com/pytorch/executorch/pull/24072 Adds `triton::sdpa_varlen`, the attention kernel the CUDA batched-sequence KV cache needs. Each resident sequence owns one column of a `[rows, cols, H, D]` store, and a device segment table `seq_info Int32[S, 4] = (q_start, q_len, col, kv_end)` says which packed query rows belong to which column and how long that sequence is. Lengths are read on device, so one captured CUDA graph serves any mix of sequence lengths, as the MLX per-sequence SDPA does per group. Covers causal, non-causal (DFlash draft blocks), sliding window and ring storage, GQA, rows no segment covers (zeroed), and split-K for short steps. A symbolic query length takes the generic path without host-side length branches, so exporting it with a dynamic step width does not pin that width. It is a new op in its own file; `triton::sdpa` and its dispatch are untouched. ghstack-source-id: 469359730 @exported-using-ghexport Differential Revision: [D125987296](https://our.internmc.facebook.com/intern/diff/D125987296/) --- backends/cuda/BUCK | 1 + backends/cuda/tests/bench_sdpa_varlen.py | 127 +++++ backends/cuda/tests/targets.bzl | 20 + backends/cuda/tests/test_sdpa_varlen.py | 429 +++++++++++++++ backends/cuda/triton/kernels/sdpa_varlen.py | 573 ++++++++++++++++++++ examples/models/muse-glimmer/model/BUCK | 25 + 6 files changed, 1175 insertions(+) create mode 100644 backends/cuda/tests/bench_sdpa_varlen.py create mode 100644 backends/cuda/tests/test_sdpa_varlen.py create mode 100644 backends/cuda/triton/kernels/sdpa_varlen.py diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index 0d7a2e59ced..5ab15c33016 100644 --- a/backends/cuda/BUCK +++ b/backends/cuda/BUCK @@ -178,6 +178,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..57276aced7c --- /dev/null +++ b/backends/cuda/tests/bench_sdpa_varlen.py @@ -0,0 +1,127 @@ +# 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) + step_args = (q, k, v, info, scale, True, window, 0) + cell_args = ( + q, + k_pool.transpose(1, 2), + v_pool.transpose(1, 2), + mask, + 0.0, + False, + scale, + True, + read_len, + ) + + def varlen(args=step_args): + return sdpa_varlen(*args) + + def cell(args=cell_args): + return sdpa(*args) + + 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 f60a0261557..090f0db16ed 100644 --- a/backends/cuda/tests/targets.bzl +++ b/backends/cuda/tests/targets.bzl @@ -134,6 +134,26 @@ 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", + "//executorch/examples/models/muse-glimmer/model:dflash_model", + ], + 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..297f9dca809 --- /dev/null +++ b/backends/cuda/tests/test_sdpa_varlen.py @@ -0,0 +1,429 @@ +# 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 +import torch.nn.functional as F + +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 + + +def _dflash_reference(q, k_store, v_store, segments, scale, window, ring): + """Use the original drafter's context-ring + whole-block mask.""" + from executorch.examples.models.muse_glimmer.model.dflash_model import ( + build_dflash_swa_mask, + ) + + out = torch.zeros_like(q, dtype=torch.float32) + groups = q.shape[1] // k_store.shape[2] + for start, block, col, end in segments: + if block == 0: + continue + context = end - block + slots = torch.arange(window, device=q.device) + context_positions = slots + ((context - 1 - slots) // window) * window + positions = torch.cat( + (context_positions, torch.arange(context, end, device=q.device)) + ).clamp_min(0) + rows = positions % ring if ring else positions + keys = k_store[rows, col].float().repeat_interleave(groups, dim=1) + values = v_store[rows, col].float().repeat_interleave(groups, dim=1) + mask = build_dflash_swa_mask(context, block, window, window).to(q.device) + out[:, :, start : start + block] = F.scaled_dot_product_attention( + q[:, :, start : start + block].float(), + keys.transpose(0, 1)[None], + values.transpose(0, 1)[None], + attn_mask=mask, + scale=scale, + ) + return out + + +def _dflash_inputs(contexts, block, window, ring, *, hq=8, hkv=2, d=64): + rows = ring or max(contexts) + block + q = torch.zeros( + 1, hq, len(contexts) * block, d, device="cuda", dtype=torch.bfloat16 + ) + q[..., 0] = 1 + k = torch.zeros(rows, 4, hkv, d, device="cuda", dtype=torch.bfloat16) + v = torch.zeros_like(k) + segments = [] + for col, context in enumerate(contexts): + oldest = max(0, context - window) + row = oldest % ring if ring else oldest + # Make dropping even the oldest retained context observable: ordinary + # per-query windows omit this key from every draft block query. + k[row, col, :, 0] = 100 + v[row, col] = 1 + segments.append((col * block, block, col, context + block)) + padded = segments + [(0, 0, 0, 0)] * (4 - len(segments)) + info = torch.tensor(padded, device="cuda", dtype=torch.int32) + return q, k, v, info, segments + + +@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_dflash_context_window_matches_original_mask(self): + window = 2048 + for block, ring in itertools.product((2, 16), (0, 2175)): + q, k, v, info, segments = _dflash_inputs((3000, 7100), block, window, ring) + got = sdpa_varlen( + q, k, v, info, 0.125, False, window, ring, context_window=True + ) + want = _dflash_reference(q, k, v, segments, 0.125, window, ring) + torch.testing.assert_close(got.float(), want, atol=2e-2, rtol=2e-2) + ordinary = sdpa_varlen(q, k, v, info, 0.125, False, window, ring) + ordinary_want = _reference(q, k, v, segments, 0.125, False, window, ring) + torch.testing.assert_close( + ordinary.float(), ordinary_want, atol=2e-2, rtol=2e-2 + ) + self.assertGreater((got.float() - ordinary.float()).abs().max().item(), 0.9) + + def test_dflash_context_window_muse_head_geometry(self): + window = 2048 + for block in (2, 16): + ring = window + block + q, k, v, info, segments = _dflash_inputs( + (3000, 7100), block, window, ring, hq=32, hkv=8, d=128 + ) + scale = 128**-0.5 + got = sdpa_varlen( + q, k, v, info, scale, False, window, ring, context_window=True + ) + want = _dflash_reference(q, k, v, segments, scale, window, ring) + torch.testing.assert_close(got.float(), want, atol=2e-2, rtol=2e-2) + + def test_dflash_context_window_graph_replay(self): + window, ring, block = 2048, 2175, 16 + q, k, v, info, _ = _dflash_inputs((4096, 6000), block, window, ring) + sdpa_varlen(q, k, v, info, 0.125, False, window, ring, context_window=True) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = sdpa_varlen( + q, k, v, info, 0.125, False, window, ring, context_window=True + ) + for contexts in ((4096, 6000), (8192, 2500)): + _, new_k, new_v, new_info, segments = _dflash_inputs( + contexts, block, window, ring + ) + k.copy_(new_k) + v.copy_(new_v) + info.copy_(new_info) + graph.replay() + want = _dflash_reference(q, k, v, segments, 0.125, window, ring) + torch.testing.assert_close(out.float(), want, atol=2e-2, rtol=2e-2) + + def test_dynamic_dflash_context_window(self): + window, ring, block = 2048, 2175, 16 + compiled = torch.compile( + lambda q, k, v, info: sdpa_varlen( + q, k, v, info, 0.125, False, window, ring, context_window=True + ), + fullgraph=True, + dynamic=True, + ) + torch._dynamo.reset() + frames = torch._dynamo.utils.counters["stats"] + before = frames["unique_graphs"] + for owners in (2, 3, 4): + q, k, v, info, segments = _dflash_inputs( + (3000, 5000, 7000, 9000)[:owners], block, window, ring + ) + got = compiled(q, k, v, info) + want = _dflash_reference(q, k, v, segments, 0.125, window, ring) + torch.testing.assert_close(got.float(), want, atol=2e-2, rtol=2e-2) + self.assertEqual(frames["unique_graphs"] - before, 1) + + 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"] + before = frames["unique_graphs"] + 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"] - before, 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) + + 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) + with self.assertRaisesRegex(ValueError, "context_window requires non-causal"): + sdpa_varlen(q, k, v, info, 0.125, context_window=True) + + +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..bc3ecea5698 --- /dev/null +++ b/backends/cuda/triton/kernels/sdpa_varlen.py @@ -0,0 +1,573 @@ +# 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, + CONTEXT_WINDOW: 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: + if CONTEXT_WINDOW: + lo = tl.maximum(kv_end - q_len - WINDOW, 0) + else: + 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: + if CONTEXT_WINDOW: + visible = visible & (kv_pos[None, :] >= kv_end - q_len - WINDOW) + else: + 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, + CONTEXT_WINDOW: 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: + if CONTEXT_WINDOW: + lo = tl.maximum(kv_end - q_len - WINDOW, 0) + else: + 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: + if CONTEXT_WINDOW: + visible &= kv_pos[None, :] >= kv_end - q_len - WINDOW + else: + 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, + context_window: bool = False, +) -> 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. + context_window: non-causal draft queries share the last ``window`` + context positions before ``kv_end - q_len`` and see the entire + query block. Ordinary attention keeps its per-query window. + 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") + if context_window and is_causal: + raise ValueError("context_window requires non-causal attention") + 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, + CONTEXT_WINDOW=context_window, + 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, + CONTEXT_WINDOW=context_window, + 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, + context_window=False, +): + return torch.empty_like(query) diff --git a/examples/models/muse-glimmer/model/BUCK b/examples/models/muse-glimmer/model/BUCK index d08324fd768..c59d50a96c0 100644 --- a/examples/models/muse-glimmer/model/BUCK +++ b/examples/models/muse-glimmer/model/BUCK @@ -14,3 +14,28 @@ fbcode_target( "//caffe2:torch", ], ) + +fbcode_target( + _kind = runtime.python_library, + name = "muse_glimmer_model", + srcs = ["model.py"], + _is_external_target = True, + base_module = "executorch.examples.models.muse_glimmer.model", + visibility = ["PUBLIC"], + deps = [ + "//caffe2:torch", + ], +) + +fbcode_target( + _kind = runtime.python_library, + name = "dflash_model", + srcs = ["dflash_model.py"], + _is_external_target = True, + base_module = "executorch.examples.models.muse_glimmer.model", + visibility = ["PUBLIC"], + deps = [ + "//caffe2:torch", + "//executorch/examples/models/muse-glimmer/model:muse_glimmer_model", + ], +)