From d0c851218442568e4b9533d3b4f8aa79d345d922 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Sun, 18 Oct 2026 16:46:40 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/cuda/cuda_backend.py | 26 +- backends/cuda/passes/lower_offgraph_kv.py | 174 ++++++++- backends/cuda/tests/targets.bzl | 26 ++ .../test_offgraph_kv_batched_sequence.py | 346 ++++++++++++++++++ 4 files changed, 553 insertions(+), 19 deletions(-) create mode 100644 backends/cuda/tests/test_offgraph_kv_batched_sequence.py diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index ec58eb6818c..7031880d201 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -33,6 +33,7 @@ LowerOffGraphKVPass, OFFGRAPH_KV_COMPILE_SPEC, OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC, + PACKED_LAYOUTS, parse_offgraph_kv_manifest, ) from executorch.backends.cuda.passes.move_cond_predicate_to_cpu import ( @@ -820,18 +821,17 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any] ) triton_kernel_mode = mode passes = [MoveCondPredicateToCpuPass(), ReplaceInt64FloorDivWithFloatPass()] - cell_layout = any( - spec.key == OFFGRAPH_KV_COMPILE_SPEC - and parse_offgraph_kv_manifest(spec.value)["layout"] == "cell" - for spec in compile_specs + layout = next( + ( + parse_offgraph_kv_manifest(spec.value)["layout"] + for spec in compile_specs + if spec.key == OFFGRAPH_KV_COMPILE_SPEC + ), + "sequence", ) for spec in compile_specs: if spec.key == OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC: - passes.append( - CheckOffGraphKVStepWidthPass( - spec.value, layout="cell" if cell_layout else "sequence" - ) - ) + passes.append(CheckOffGraphKVStepWidthPass(spec.value, layout=layout)) for spec in compile_specs: if spec.key == OFFGRAPH_KV_COMPILE_SPEC: passes.append( @@ -972,9 +972,9 @@ def get_extra_aoti_compile_context_manager( low_memory_mode = mode elif spec.key == CUDA_GRAPH_AUTOTUNE_TIMING_COMPILE_SPEC: cuda_graph_autotune_timing = _on_off_compile_spec_value(spec) - cell_layout = any( + packed_layout = any( spec.key == OFFGRAPH_KV_COMPILE_SPEC - and parse_offgraph_kv_manifest(spec.value)["layout"] == "cell" + and parse_offgraph_kv_manifest(spec.value)["layout"] in PACKED_LAYOUTS for spec in compile_specs or [] ) @@ -1024,8 +1024,8 @@ def _combined(): stack.enter_context( cuda_graph_timing(offload=_moved_program_tensors) ) - if cell_layout: - # The cell layout's step buffers are constants the graph + if packed_layout: + # A packed layout's step buffers are constants the graph # only reads, while the runtime rewrites them before every # forward. Folding them, or inlining a small one as a # literal, would evaluate them at compile time -- on diff --git a/backends/cuda/passes/lower_offgraph_kv.py b/backends/cuda/passes/lower_offgraph_kv.py index 0567e809139..a7f7ffe14ac 100644 --- a/backends/cuda/passes/lower_offgraph_kv.py +++ b/backends/cuda/passes/lower_offgraph_kv.py @@ -14,6 +14,9 @@ # Imported for its registration side effect: the decomposition calls # torch.ops.triton.sdpa, which only exists once this module has been loaded. from executorch.backends.cuda.triton.kernels.sdpa import sdpa # noqa: F401 +from executorch.backends.cuda.triton.kernels.sdpa_varlen import ( # noqa: F401 + sdpa_varlen, +) from executorch.backends.transforms.utils import create_constant_placeholder from executorch.exir.dialects._ops import ops as exir_ops from torch.export import ExportedProgram @@ -36,6 +39,26 @@ OFFGRAPH_KV_READ_LEN_FQN = OFFGRAPH_KV_FQN_PREFIX + "read_len" +# Batched-sequence layout: each sequence keeps its own column of every layer's +# storage, and the runtime cache writes which column each token belongs to +# (seq_cols) and each run of one sequence's query rows (seq_info) before each +# forward. Declared once per program at a fixed shape, like the cell buffers. +OFFGRAPH_KV_SEQ_COLS_FQN = OFFGRAPH_KV_FQN_PREFIX + "seq_cols" +OFFGRAPH_KV_SEQ_INFO_FQN = OFFGRAPH_KV_FQN_PREFIX + "seq_info" + +# Layouts whose step packs tokens of several sequences on one axis. +PACKED_LAYOUTS = ("cell", "batched-sequence") + + +def batched_sequence_max_segments(max_seqs: int) -> int: + """Rows of seq_info: one per run of a sequence's query rows in a step. + + Same rule as batched_sequence_max_segments() in + backends/cuda/runtime/cuda_kv_cache.h, which checks steps against it. + """ + return 2 * max_seqs + + def offgraph_kv_mask_fqn(window: int) -> str: """The step mask shared by every layer of ``window`` (0 = full history).""" return f"{OFFGRAPH_KV_FQN_PREFIX}mask_w{window}" @@ -70,8 +93,14 @@ def parse_offgraph_kv_manifest(value: bytes) -> dict[str, Any]: # noqa: C901 if not isinstance(max_write, int) or max_write <= 0: raise ValueError("off-graph max_write must be positive") layout = manifest.setdefault("layout", "sequence") - if layout not in ("sequence", "cell"): + if layout not in ("sequence", *PACKED_LAYOUTS): raise ValueError(f"invalid off-graph KV layout {layout!r}") + if layout == "batched-sequence": + for key in ("max_seqs", "max_segments"): + if not isinstance(manifest.get(key), int) or manifest[key] <= 0: + raise ValueError(f"off-graph batched-sequence layout needs {key} > 0") + if not isinstance(manifest.setdefault("causal", True), bool): + raise ValueError("off-graph causal must be a bool") if layout == "cell": # A cell step packs tokens of several sequences into the shared pool, # so the pool bounds its width, not one sequence's capacity. @@ -240,6 +269,86 @@ def cell_step( return out +def batched_sequence_step( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + position: torch.Tensor, + k_storage: torch.Tensor, + v_storage: torch.Tensor, + seq_cols: torch.Tensor, + seq_info: torch.Tensor, + scale: float, + rows: int, + columns: int, + window: int, + ring: bool, + causal: bool, + out_dtype: torch.dtype | None = None, +) -> torch.Tensor: + """Write this step's K/V into each sequence's column, then attend. + + The storage is declared ``[1, rows * columns, H, D]``: row ``r`` of + column ``c`` is slot ``r * columns + c``, so it is ``[rows, columns, H, D]`` + and grows by whole rows, as the sequence layout's does. Token ``i`` lands + in column ``seq_cols[i]`` at row ``position % rows`` (a ring layer's rows + wrap; a flat layer's never do). ``seq_info`` names each segment's query + rows and length, so every query reads only its own column. + + ``seq_cols`` is declared at the widest step; writes use the first + K/V-width entries. ``k``/``v`` are BHSD; the storage is BSHD. + """ + position = position.reshape(-1) + slots = (position % rows) * columns + seq_cols[: k.shape[2]] + _write(k_storage, slots, k) + _write(v_storage, slots, v) + heads, dim = k_storage.shape[2], k_storage.shape[3] + out = torch.ops.triton.sdpa_varlen( + q, + k_storage.view(rows, columns, heads, dim), + v_storage.view(rows, columns, heads, dim), + seq_info, + scale, + causal, + window, + rows if ring else 0, + ) + if out_dtype is not None and out.dtype != out_dtype: + out = out.to(out_dtype) + return out + + +def _batched_sequence_step_fn( + scale: float, + rows: int, + columns: int, + window: int, + ring: bool, + causal: bool, + out_dtype, +): + def fn(q, k, v, position, k_storage, v_storage, seq_cols, seq_info): + return batched_sequence_step( + q, + k, + v, + position, + k_storage, + v_storage, + seq_cols, + seq_info, + scale, + rows, + columns, + window, + ring, + causal, + out_dtype, + ) + + return fn + + def _cell_step_fn(scale: float, out_dtype): def fn(q, k, v, k_pool, v_pool, cells, read_len, mask): return cell_step( @@ -380,9 +489,9 @@ class CheckOffGraphKVStepWidthPass: requires_exported_program = True def __init__(self, value: bytes, layout: str = "sequence") -> None: - if layout not in ("sequence", "cell"): + if layout not in ("sequence", *PACKED_LAYOUTS): raise ValueError(f"invalid off-graph KV layout {layout!r}") - self._cell_layout = layout == "cell" + self._packed = layout in PACKED_LAYOUTS text = value.decode("utf-8") fields = text.split(":") if len(fields) != 2 or not all(f.isdigit() for f in fields): @@ -411,7 +520,7 @@ def __call__(self, exported_program: ExportedProgram) -> ExportedProgram: position = node.args[3] source = ( _cell_position_source(position) - if self._cell_layout + if self._packed else _source_input(position) ) if source is not declared: @@ -449,6 +558,15 @@ def __init__(self, manifest: dict[str, Any]) -> None: def _cell_layout(self) -> bool: return self._manifest.get("layout", "sequence") == "cell" + @property + def _batched_sequence_layout(self) -> bool: + return self._manifest.get("layout", "sequence") == "batched-sequence" + + @property + def _columns(self) -> int: + # One per sequence, and a scratch column static-width padding writes. + return self._manifest["max_seqs"] + 1 + @staticmethod def _compile_storage(shape, dtype: torch.dtype, device): storage = torch.empty(shape, dtype=dtype, device=device) @@ -485,7 +603,12 @@ def _layer_capacity(self, layer: dict[str, Any]) -> int: # cell layout keeps every layer's history in the one pool. return self._manifest["max_cells"] if layer["policy"] == "ring": - return ring_physical_capacity(layer["window"], self._manifest["max_write"]) + rows = ring_physical_capacity(layer["window"], self._manifest["max_write"]) + if self._batched_sequence_layout: + # A sequence never reaches past its capacity, so a ring that + # large never wraps. + rows = min(rows, self._manifest["maximum_capacity"]) + return rows return self._manifest["maximum_capacity"] def _storage_nodes( @@ -508,6 +631,9 @@ def _storage_nodes( # back this with an allocation holding only the rows written so far # and grow it. A BHSD declaration would put head h at h * capacity, # past any smaller allocation. + if self._batched_sequence_layout: + # Rows outermost, so the runtime grows every column at once. + capacity *= self._columns shape = (kv.shape[0], capacity, kv.shape[1], kv.shape[3]) prefix = f"{OFFGRAPH_KV_FQN_PREFIX}layer_{layer_id}" names = (f"{prefix}_k", f"{prefix}_v") @@ -534,7 +660,7 @@ def _runtime_buffer(exported_program: ExportedProgram, name: str, value): ) def _cell_buffer(self, exported_program, name: str, shape, dtype, device): - """A step buffer of the cell layout, declared once per program. + """A step buffer of a packed layout, declared once per program. Unlike the pools these are small, so they carry real (zero) bytes at compile time: Inductor reads small constants while it compiles, and @@ -580,6 +706,38 @@ def _lower_cell( example_args = tuple(n.meta["val"] for n in call_args) return _cell_step_fn(node.args[5], out_dtype), call_args, example_args + def _lower_batched_sequence( + self, exported_program, node, layer, k_storage, v_storage, out_dtype + ): + device = node.args[1].meta["val"].device + seq_cols = self._cell_buffer( + exported_program, + OFFGRAPH_KV_SEQ_COLS_FQN, + (self._manifest["max_write"],), + torch.int64, + device, + ) + seq_info = self._cell_buffer( + exported_program, + OFFGRAPH_KV_SEQ_INFO_FQN, + (self._manifest["max_segments"], 4), + torch.int32, + device, + ) + call_args = (*node.args[0:4], k_storage, v_storage, seq_cols, seq_info) + example_args = tuple(n.meta["val"] for n in call_args) + ring = layer["policy"] == "ring" + fn = _batched_sequence_step_fn( + node.args[5], + self._layer_capacity(layer), + self._columns, + layer["window"] if ring else 0, + ring, + self._manifest["causal"], + out_dtype, + ) + return fn, call_args, example_args + def _ring_mask(self, graph, position, buf_size: int, window: int, before): """Emit the ring mask once and share it across layers that match. @@ -671,6 +829,10 @@ def __call__(self, exported_program: ExportedProgram) -> ExportedProgram: fn, call_args, example_args = self._lower_cell( exported_program, node, layer, k_storage, v_storage, out_dtype ) + elif self._batched_sequence_layout: + fn, call_args, example_args = self._lower_batched_sequence( + exported_program, node, layer, k_storage, v_storage, out_dtype + ) else: fn, call_args, example_args = self._lower_sequence( graph_module.graph, node, layer, k_storage, v_storage, out_dtype diff --git a/backends/cuda/tests/targets.bzl b/backends/cuda/tests/targets.bzl index af1cf582af1..b5048aad415 100644 --- a/backends/cuda/tests/targets.bzl +++ b/backends/cuda/tests/targets.bzl @@ -302,6 +302,32 @@ def define_common_targets(is_fbcode = False): ), ) + python_unittest_remote_gpu( + name = "test_offgraph_kv_batched_sequence", + srcs = [ + "test_offgraph_kv_batched_sequence.py", + ], + visibility = [ + "//executorch/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/cuda:cuda_backend", + "//executorch/backends/cuda:cuda_partitioner", + "//executorch/backends/cuda:cuda_passes", + "//executorch/exir:lib", + "//executorch/exir/backend:compile_spec_schema", + "//executorch/exir/dialects:lib", + # The oracle is the neutral op itself, not a reference rebuilt here. + "//executorch/extension/llm/cache:cache", + ], + keep_gpu_sections = True, + remote_execution = re_test_utils.remote_execution( + platform = "gpu-remote-execution", + subplatform = "A100-exclusive", + ), + ) + python_unittest( name = "test_quantized_gemm_family", srcs = [ diff --git a/backends/cuda/tests/test_offgraph_kv_batched_sequence.py b/backends/cuda/tests/test_offgraph_kv_batched_sequence.py new file mode 100644 index 00000000000..9b6ed2df30e --- /dev/null +++ b/backends/cuda/tests/test_offgraph_kv_batched_sequence.py @@ -0,0 +1,346 @@ +# 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. + +"""The batched-sequence off-graph KV layout on CUDA, against the neutral op. + +The oracle is ``BatchedSequenceReferenceCache``: every sequence keeps a private +history, and a packed step's tokens attend only their own sequence -- the same +contract MLX's batched-sequence cache implements. The lowered decomposition +reads the step from the two buffers the runtime cache writes: each token's +storage column (``seq_cols``) and each run of one sequence's rows +(``seq_info``). +""" + +import json +import unittest + +import torch + +from executorch.backends.cuda.passes.lower_offgraph_kv import ( + batched_sequence_step, + CheckOffGraphKVStepWidthPass, + LowerOffGraphKVPass, + OFFGRAPH_KV_FQN_PREFIX, + OFFGRAPH_KV_SEQ_COLS_FQN, + OFFGRAPH_KV_SEQ_INFO_FQN, + parse_offgraph_kv_manifest, + ring_physical_capacity, +) +from executorch.exir import EdgeCompileConfig, to_edge +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.extension.llm.cache.reference_cache import ( + BatchedSequenceReferenceCache, + CacheConfig, + LayerPolicy, +) +from executorch.extension.llm.cache.update_and_attend import REGISTRY + +N_KV_HEADS = 2 +N_Q_HEADS = 4 +HEAD_DIM = 64 +SCALE = 0.125 + + +def _inputs(start: int, length: int): + q = torch.randn(1, N_Q_HEADS, length, HEAD_DIM, device="cuda", dtype=torch.bfloat16) + k = torch.randn( + 1, N_KV_HEADS, length, HEAD_DIM, device="cuda", dtype=torch.bfloat16 + ) + v = torch.randn_like(k) + return q, k, v, torch.arange(start, start + length, device="cuda") + + +def _batch(groups): + """Packs (seq_id, start, length) groups onto one token axis.""" + parts = [_inputs(start, length) for _, start, length in groups] + seq_ids = [seq for seq, _, length in groups for _ in range(length)] + q, k, v = (torch.cat([p[i] for p in parts], dim=2) for i in range(3)) + position = torch.cat([p[3] for p in parts]) + return seq_ids, q, k, v, position + + +def _plan(seq_ids, position, columns, width, max_segments, scratch): + """What the runtime cache writes before a forward: each token's column, + and one (q_start, q_len, col, kv_end) row per run of one sequence.""" + seq_cols = [columns[s] for s in seq_ids] + [scratch] * (width - len(seq_ids)) + info, start = [], 0 + positions = position.tolist() + while start < len(seq_ids): + end = start + while end < len(seq_ids) and seq_ids[end] == seq_ids[start]: + end += 1 + info.append( + (start, end - start, columns[seq_ids[start]], positions[end - 1] + 1) + ) + start = end + info += [(0, 0, 0, 0)] * (max_segments - len(info)) + return ( + torch.tensor(seq_cols, dtype=torch.long, device="cuda"), + torch.tensor(info, dtype=torch.int32, device="cuda"), + ) + + +class _Oracle: + def __init__(self, capacity: int, windows) -> None: + self.windows = tuple(windows) + self._cache = BatchedSequenceReferenceCache( + CacheConfig( + n_layers=len(self.windows), + n_kv_heads=N_KV_HEADS, + head_dim=HEAD_DIM, + capacity=capacity, + dtype=torch.float32, + layers=tuple( + LayerPolicy.ring(w) if w else LayerPolicy.flat() + for w in self.windows + ), + ) + ) + self._key = f"cuda-offgraph-batched-sequence-{id(self)}" + REGISTRY.install(self._key, self._cache) + + def step(self, seq_ids, q, k, v, position): + self._cache.declare_step(seq_ids) + outputs = [] + with REGISTRY.active(self._key): + for layer_id in range(len(self.windows)): + outputs.append( + torch.ops.kvcache.update_and_attend( + q.float().cpu(), + k.float().cpu(), + v.float().cpu(), + position.reshape(-1, 1).cpu(), + layer_id, + SCALE, + torch.float32, + ) + ) + return outputs + + def close(self) -> None: + REGISTRY.uninstall(self._key) + + +def _max_abs_diff(actual, expected) -> float: + return (actual.float().cpu() - expected).abs().max().item() + + +# Two sequences prefilling in one step, decoding together, then one decoding +# beside the other's chunk, then a third joining late. +_STEPS = ( + ((0, 0, 5), (1, 0, 3)), + ((0, 5, 1), (1, 3, 1)), + ((1, 4, 4), (0, 6, 1)), + ((2, 0, 6), (0, 7, 1), (1, 8, 1)), + ((0, 8, 1), (2, 6, 1)), +) + + +@unittest.skipUnless(torch.cuda.is_available(), "CUDA required") +class BatchedSequenceStepTest(unittest.TestCase): + """batched_sequence_step, fed the runtime's plan, matches the neutral op.""" + + MAX_SEQS = 3 + MAX_WRITE = 16 + + def _run(self, steps, capacity, window=0, max_segments=6): + torch.manual_seed(20) + oracle = _Oracle(4 * capacity, (window,)) + self.addCleanup(oracle.close) + columns = {seq: seq for seq in range(self.MAX_SEQS)} + cols = self.MAX_SEQS + 1 + ring = window > 0 + rows = ( + min(capacity, ring_physical_capacity(window, self.MAX_WRITE)) + if ring + else capacity + ) + k_store = torch.zeros( + 1, rows * cols, N_KV_HEADS, HEAD_DIM, device="cuda", dtype=torch.bfloat16 + ) + v_store = torch.zeros_like(k_store) + for groups in steps: + seq_ids, q, k, v, position = _batch(groups) + (expected,) = oracle.step(seq_ids, q, k, v, position) + # A static method one wider than the step: one padding row. + width = len(seq_ids) + 1 + pad = torch.zeros( + 1, q.shape[1], 1, HEAD_DIM, device="cuda", dtype=torch.bfloat16 + ) + kpad = torch.zeros( + 1, N_KV_HEADS, 1, HEAD_DIM, device="cuda", dtype=torch.bfloat16 + ) + seq_cols, seq_info = _plan( + seq_ids, position, columns, self.MAX_WRITE, max_segments, cols - 1 + ) + out = batched_sequence_step( + torch.cat([q, pad], 2), + torch.cat([k, kpad], 2), + torch.cat([v, kpad], 2), + torch.cat([position, position.new_zeros(1)]), + k_store, + v_store, + seq_cols, + seq_info, + SCALE, + rows, + cols, + window, + ring, + True, + ) + self.assertLess(_max_abs_diff(out[:, :, : width - 1], expected), 2e-2) + self.assertTrue((out[:, :, width - 1 :] == 0).all()) + return k_store, rows, cols + + def test_sequences_see_only_their_own_history(self) -> None: + self._run(_STEPS, capacity=32) + + def test_windowed_ring_layer(self) -> None: + self._run(_STEPS, capacity=32, window=4) + + def test_padding_writes_only_the_scratch_column(self) -> None: + k_store, rows, cols = self._run(_STEPS[:1], capacity=16) + grid = k_store.view(rows, cols, N_KV_HEADS, HEAD_DIM) + # Sequence 2 never ran; only row 0 of the scratch column took padding. + self.assertFalse(grid[:, 2].any()) + self.assertFalse(grid[1:, cols - 1].any()) + + +class _FlatAndRing(torch.nn.Module): + def forward(self, q, k, v, position): + flat = torch.ops.kvcache.update_and_attend( + q, k, v, position, 0, SCALE, torch.bfloat16 + ) + ring = torch.ops.kvcache.update_and_attend( + q, k, v, position, 1, SCALE, torch.bfloat16 + ) + return flat + ring + + +@unittest.skipUnless(torch.cuda.is_available(), "CUDA required") +class LowerBatchedSequencePassTest(unittest.TestCase): + CAPACITY = 32 + MAX_WRITE = 8 + MAX_SEQS = 3 + MAX_SEGMENTS = 6 + WINDOW = 4 + + @classmethod + def manifest(cls, **overrides) -> bytes: + manifest = { + "version": 1, + "layout": "batched-sequence", + "maximum_capacity": cls.CAPACITY, + "max_write": cls.MAX_WRITE, + "max_seqs": cls.MAX_SEQS, + "max_segments": cls.MAX_SEGMENTS, + "layers": [ + {"layer_id": 0, "policy": "flat"}, + {"layer_id": 1, "policy": "ring", "window": cls.WINDOW}, + ], + } + manifest.update(overrides) + return json.dumps(manifest).encode() + + def _lowered(self): + q, k, v, position = _inputs(0, self.MAX_WRITE) + t = torch.export.Dim("t", min=1, max=self.MAX_WRITE) + program = torch.export.export( + _FlatAndRing(), + (q, k, v, position.reshape(-1, 1)), + dynamic_shapes=({2: t}, {2: t}, {2: t}, {0: t}), + strict=False, + ) + edge = to_edge( + program, compile_config=EdgeCompileConfig(_check_ir_validity=False) + ).exported_program() + return LowerOffGraphKVPass(parse_offgraph_kv_manifest(self.manifest()))(edge) + + def test_manifest_is_validated(self) -> None: + parsed = parse_offgraph_kv_manifest(self.manifest()) + self.assertEqual(parsed["layout"], "batched-sequence") + self.assertTrue(parsed["causal"]) + self.assertFalse( + parse_offgraph_kv_manifest(self.manifest(causal=False))["causal"] + ) + for bad in ( + {"max_seqs": 0}, + {"max_segments": None}, + {"causal": 1}, + {"max_write": self.CAPACITY + 1}, + ): + with self.assertRaises(ValueError): + parse_offgraph_kv_manifest(self.manifest(**bad)) + + def test_declares_column_storage_and_step_buffers(self) -> None: + lowered = self._lowered() + graph = lowered.graph_module.graph + target = exir_ops.edge.kvcache.update_and_attend.default + self.assertFalse(any(n.target == target for n in graph.nodes)) + buffers = { + n.name: tuple(n.meta["val"].shape) + for n in graph.nodes + if n.op == "placeholder" and n.name.startswith(OFFGRAPH_KV_FQN_PREFIX) + } + cols = self.MAX_SEQS + 1 + ring = min(self.CAPACITY, ring_physical_capacity(self.WINDOW, self.MAX_WRITE)) + flat = (1, self.CAPACITY * cols, N_KV_HEADS, HEAD_DIM) + windowed = (1, ring * cols, N_KV_HEADS, HEAD_DIM) + self.assertEqual( + buffers, + { + f"{OFFGRAPH_KV_FQN_PREFIX}layer_0_k": flat, + f"{OFFGRAPH_KV_FQN_PREFIX}layer_0_v": flat, + f"{OFFGRAPH_KV_FQN_PREFIX}layer_1_k": windowed, + f"{OFFGRAPH_KV_FQN_PREFIX}layer_1_v": windowed, + OFFGRAPH_KV_SEQ_COLS_FQN: (self.MAX_WRITE,), + OFFGRAPH_KV_SEQ_INFO_FQN: (self.MAX_SEGMENTS, 4), + }, + ) + calls = [ + n for n in graph.nodes if n.target == torch.ops.triton.sdpa_varlen.default + ] + self.assertEqual(len(calls), 2) + + def test_lowered_program_matches_the_neutral_op(self) -> None: + torch.manual_seed(21) + lowered = self._lowered() + runtime = {} + for name, value in list(lowered.constants.items()): + if name.startswith(OFFGRAPH_KV_FQN_PREFIX): + runtime[name] = torch.zeros( + value.shape, device="cuda", dtype=value.dtype + ) + lowered.constants[name] = runtime[name] + module = lowered.module() + oracle = _Oracle(4 * self.CAPACITY, (0, self.WINDOW)) + self.addCleanup(oracle.close) + columns = {seq: seq for seq in range(self.MAX_SEQS)} + for groups in _STEPS: + seq_ids, q, k, v, position = _batch(groups) + flat, ring = oracle.step(seq_ids, q, k, v, position) + seq_cols, seq_info = _plan( + seq_ids, + position, + columns, + self.MAX_WRITE, + self.MAX_SEGMENTS, + self.MAX_SEQS, + ) + runtime[OFFGRAPH_KV_SEQ_COLS_FQN].copy_(seq_cols) + runtime[OFFGRAPH_KV_SEQ_INFO_FQN].copy_(seq_info) + out = module(q, k, v, position.reshape(-1, 1)) + self.assertLess(_max_abs_diff(out, flat + ring), 3e-2) + + def test_step_width_pass_accepts_the_packed_layout(self) -> None: + with self.assertRaises(ValueError): + CheckOffGraphKVStepWidthPass(b"0:0", layout="paged") + CheckOffGraphKVStepWidthPass(b"0:0", layout="batched-sequence") + + +if __name__ == "__main__": + unittest.main()