Skip to content

Commit cc05bff

Browse files
[executorch][cuda] Lower off-graph KV to the batched-sequence layout (#24881)
This PR was created by the merge bot to help merge the original PR into the main branch. ghstack PR number: #24074 by @Gasoonjia ^ Please use this as the source of truth for the PR details, comments, and reviews ghstack PR base: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/92/base ghstack PR head: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/92/head Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/91/orig Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/92/orig Differential Revision: [D125987298](https://our.internmc.facebook.com/intern/diff/D125987298/) @diff-train-skip-merge Co-authored-by: gasoonjia <gasoonjia@icloud.com>
1 parent 63c24b9 commit cc05bff

4 files changed

Lines changed: 635 additions & 19 deletions

File tree

‎backends/cuda/cuda_backend.py‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
LowerOffGraphKVPass,
3939
OFFGRAPH_KV_COMPILE_SPEC,
4040
OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC,
41+
PACKED_LAYOUTS,
4142
parse_offgraph_kv_manifest,
4243
)
4344
from executorch.backends.cuda.passes.move_cond_predicate_to_cpu import (
@@ -1277,18 +1278,17 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any]
12771278
ReplaceFusedSdpaWithSdpaPass(),
12781279
ReplaceMaskedCopyWithWherePass(),
12791280
]
1280-
cell_layout = any(
1281-
spec.key == OFFGRAPH_KV_COMPILE_SPEC
1282-
and parse_offgraph_kv_manifest(spec.value)["layout"] == "cell"
1283-
for spec in compile_specs
1281+
layout = next(
1282+
(
1283+
parse_offgraph_kv_manifest(spec.value)["layout"]
1284+
for spec in compile_specs
1285+
if spec.key == OFFGRAPH_KV_COMPILE_SPEC
1286+
),
1287+
"sequence",
12841288
)
12851289
for spec in compile_specs:
12861290
if spec.key == OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC:
1287-
passes.append(
1288-
CheckOffGraphKVStepWidthPass(
1289-
spec.value, layout="cell" if cell_layout else "sequence"
1290-
)
1291-
)
1291+
passes.append(CheckOffGraphKVStepWidthPass(spec.value, layout=layout))
12921292
for spec in compile_specs:
12931293
if spec.key == OFFGRAPH_KV_COMPILE_SPEC:
12941294
passes.append(
@@ -1441,9 +1441,9 @@ def get_extra_aoti_compile_context_manager(
14411441
low_memory_mode = mode
14421442
elif spec.key == CUDA_GRAPH_AUTOTUNE_TIMING_COMPILE_SPEC:
14431443
cuda_graph_autotune_timing = _on_off_compile_spec_value(spec)
1444-
cell_layout = any(
1444+
packed_layout = any(
14451445
spec.key == OFFGRAPH_KV_COMPILE_SPEC
1446-
and parse_offgraph_kv_manifest(spec.value)["layout"] == "cell"
1446+
and parse_offgraph_kv_manifest(spec.value)["layout"] in PACKED_LAYOUTS
14471447
for spec in compile_specs or []
14481448
)
14491449

@@ -1497,8 +1497,8 @@ def _combined():
14971497
stack.enter_context(
14981498
cuda_graph_timing(offload=_moved_program_tensors)
14991499
)
1500-
if cell_layout:
1501-
# The cell layout's step buffers are constants the graph
1500+
if packed_layout:
1501+
# A packed layout's step buffers are constants the graph
15021502
# only reads, while the runtime rewrites them before every
15031503
# forward. Folding them, or inlining a small one as a
15041504
# literal, would evaluate them at compile time -- on

‎backends/cuda/passes/lower_offgraph_kv.py‎

Lines changed: 186 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,9 @@
1414
# Imported for its registration side effect: the decomposition calls
1515
# torch.ops.triton.sdpa, which only exists once this module has been loaded.
1616
from executorch.backends.cuda.triton.kernels.sdpa import sdpa # noqa: F401
17+
from executorch.backends.cuda.triton.kernels.sdpa_varlen import ( # noqa: F401
18+
sdpa_varlen,
19+
)
1720
from executorch.backends.transforms.utils import create_constant_placeholder
1821
from executorch.exir.dialects._ops import ops as exir_ops
1922
from torch.export import ExportedProgram
@@ -36,6 +39,26 @@
3639
OFFGRAPH_KV_READ_LEN_FQN = OFFGRAPH_KV_FQN_PREFIX + "read_len"
3740

3841

42+
# Batched-sequence layout: each sequence keeps its own column of every layer's
43+
# storage, and the runtime cache writes which column each token belongs to
44+
# (seq_cols) and each run of one sequence's query rows (seq_info) before each
45+
# forward. Declared once per program at a fixed shape, like the cell buffers.
46+
OFFGRAPH_KV_SEQ_COLS_FQN = OFFGRAPH_KV_FQN_PREFIX + "seq_cols"
47+
OFFGRAPH_KV_SEQ_INFO_FQN = OFFGRAPH_KV_FQN_PREFIX + "seq_info"
48+
49+
# Layouts whose step packs tokens of several sequences on one axis.
50+
PACKED_LAYOUTS = ("cell", "batched-sequence")
51+
52+
53+
def batched_sequence_max_segments(max_seqs: int) -> int:
54+
"""Rows of seq_info: one per run of a sequence's query rows in a step.
55+
56+
Same rule as batched_sequence_max_segments() in
57+
backends/cuda/runtime/cuda_kv_cache.h, which checks steps against it.
58+
"""
59+
return 2 * max_seqs
60+
61+
3962
def offgraph_kv_mask_fqn(window: int) -> str:
4063
"""The step mask shared by every layer of ``window`` (0 = full history)."""
4164
return f"{OFFGRAPH_KV_FQN_PREFIX}mask_w{window}"
@@ -70,8 +93,25 @@ def parse_offgraph_kv_manifest(value: bytes) -> dict[str, Any]: # noqa: C901
7093
if not isinstance(max_write, int) or max_write <= 0:
7194
raise ValueError("off-graph max_write must be positive")
7295
layout = manifest.setdefault("layout", "sequence")
73-
if layout not in ("sequence", "cell"):
96+
if layout not in ("sequence", *PACKED_LAYOUTS):
7497
raise ValueError(f"invalid off-graph KV layout {layout!r}")
98+
if layout == "batched-sequence":
99+
for key in ("max_seqs", "max_segments"):
100+
if not isinstance(manifest.get(key), int) or manifest[key] <= 0:
101+
raise ValueError(f"off-graph batched-sequence layout needs {key} > 0")
102+
if manifest["max_segments"] != batched_sequence_max_segments(
103+
manifest["max_seqs"]
104+
):
105+
raise ValueError(
106+
"off-graph batched-sequence max_segments must equal "
107+
"batched_sequence_max_segments(max_seqs)"
108+
)
109+
if not isinstance(manifest.setdefault("causal", True), bool):
110+
raise ValueError("off-graph causal must be a bool")
111+
if not isinstance(manifest.setdefault("context_window", False), bool):
112+
raise ValueError("off-graph context_window must be a bool")
113+
if manifest["context_window"] and manifest["causal"]:
114+
raise ValueError("off-graph context_window requires non-causal attention")
75115
if layout == "cell":
76116
# A cell step packs tokens of several sequences into the shared pool,
77117
# so the pool bounds its width, not one sequence's capacity.
@@ -240,6 +280,92 @@ def cell_step(
240280
return out
241281

242282

283+
def batched_sequence_step(
284+
q: torch.Tensor,
285+
k: torch.Tensor,
286+
v: torch.Tensor,
287+
position: torch.Tensor,
288+
k_storage: torch.Tensor,
289+
v_storage: torch.Tensor,
290+
seq_cols: torch.Tensor,
291+
seq_info: torch.Tensor,
292+
scale: float,
293+
rows: int,
294+
columns: int,
295+
window: int,
296+
ring: bool,
297+
causal: bool,
298+
out_dtype: torch.dtype | None = None,
299+
context_window: bool = False,
300+
) -> torch.Tensor:
301+
"""Write this step's K/V into each sequence's column, then attend.
302+
303+
The storage is declared ``[1, rows * columns, H, D]``: row ``r`` of
304+
column ``c`` is slot ``r * columns + c``, so it is ``[rows, columns, H, D]``
305+
and grows by whole rows, as the sequence layout's does. Token ``i`` lands
306+
in column ``seq_cols[i]`` at row ``position % rows`` (a ring layer's rows
307+
wrap; a flat layer's never do). ``seq_info`` names each segment's query
308+
rows and length, so every query reads only its own column.
309+
310+
``seq_cols`` is declared at the widest step; writes use the first
311+
K/V-width entries. ``k``/``v`` are BHSD; the storage is BSHD.
312+
``context_window`` gives every draft block query the same retained context
313+
before the block; ordinary attention uses each query's own window.
314+
"""
315+
position = position.reshape(-1)
316+
slots = (position % rows) * columns + seq_cols[: k.shape[2]]
317+
_write(k_storage, slots, k)
318+
_write(v_storage, slots, v)
319+
heads, dim = k_storage.shape[2], k_storage.shape[3]
320+
out = torch.ops.triton.sdpa_varlen(
321+
q,
322+
k_storage.view(rows, columns, heads, dim),
323+
v_storage.view(rows, columns, heads, dim),
324+
seq_info,
325+
scale,
326+
causal,
327+
window,
328+
rows if ring else 0,
329+
context_window=context_window,
330+
)
331+
if out_dtype is not None and out.dtype != out_dtype:
332+
out = out.to(out_dtype)
333+
return out
334+
335+
336+
def _batched_sequence_step_fn(
337+
scale: float,
338+
rows: int,
339+
columns: int,
340+
window: int,
341+
ring: bool,
342+
causal: bool,
343+
out_dtype,
344+
context_window: bool = False,
345+
):
346+
def fn(q, k, v, position, k_storage, v_storage, seq_cols, seq_info):
347+
return batched_sequence_step(
348+
q,
349+
k,
350+
v,
351+
position,
352+
k_storage,
353+
v_storage,
354+
seq_cols,
355+
seq_info,
356+
scale,
357+
rows,
358+
columns,
359+
window,
360+
ring,
361+
causal,
362+
out_dtype,
363+
context_window,
364+
)
365+
366+
return fn
367+
368+
243369
def _cell_step_fn(scale: float, out_dtype):
244370
def fn(q, k, v, k_pool, v_pool, cells, read_len, mask):
245371
return cell_step(
@@ -380,9 +506,9 @@ class CheckOffGraphKVStepWidthPass:
380506
requires_exported_program = True
381507

382508
def __init__(self, value: bytes, layout: str = "sequence") -> None:
383-
if layout not in ("sequence", "cell"):
509+
if layout not in ("sequence", *PACKED_LAYOUTS):
384510
raise ValueError(f"invalid off-graph KV layout {layout!r}")
385-
self._cell_layout = layout == "cell"
511+
self._packed = layout in PACKED_LAYOUTS
386512
text = value.decode("utf-8")
387513
fields = text.split(":")
388514
if len(fields) != 2 or not all(f.isdigit() for f in fields):
@@ -411,7 +537,7 @@ def __call__(self, exported_program: ExportedProgram) -> ExportedProgram:
411537
position = node.args[3]
412538
source = (
413539
_cell_position_source(position)
414-
if self._cell_layout
540+
if self._packed
415541
else _source_input(position)
416542
)
417543
if source is not declared:
@@ -449,6 +575,15 @@ def __init__(self, manifest: dict[str, Any]) -> None:
449575
def _cell_layout(self) -> bool:
450576
return self._manifest.get("layout", "sequence") == "cell"
451577

578+
@property
579+
def _batched_sequence_layout(self) -> bool:
580+
return self._manifest.get("layout", "sequence") == "batched-sequence"
581+
582+
@property
583+
def _columns(self) -> int:
584+
# One per sequence, and a scratch column static-width padding writes.
585+
return self._manifest["max_seqs"] + 1
586+
452587
@staticmethod
453588
def _compile_storage(shape, dtype: torch.dtype, device):
454589
storage = torch.empty(shape, dtype=dtype, device=device)
@@ -485,7 +620,12 @@ def _layer_capacity(self, layer: dict[str, Any]) -> int:
485620
# cell layout keeps every layer's history in the one pool.
486621
return self._manifest["max_cells"]
487622
if layer["policy"] == "ring":
488-
return ring_physical_capacity(layer["window"], self._manifest["max_write"])
623+
rows = ring_physical_capacity(layer["window"], self._manifest["max_write"])
624+
if self._batched_sequence_layout:
625+
# A sequence never reaches past its capacity, so a ring that
626+
# large never wraps.
627+
rows = min(rows, self._manifest["maximum_capacity"])
628+
return rows
489629
return self._manifest["maximum_capacity"]
490630

491631
def _storage_nodes(
@@ -508,6 +648,9 @@ def _storage_nodes(
508648
# back this with an allocation holding only the rows written so far
509649
# and grow it. A BHSD declaration would put head h at h * capacity,
510650
# past any smaller allocation.
651+
if self._batched_sequence_layout:
652+
# Rows outermost, so the runtime grows every column at once.
653+
capacity *= self._columns
511654
shape = (kv.shape[0], capacity, kv.shape[1], kv.shape[3])
512655
prefix = f"{OFFGRAPH_KV_FQN_PREFIX}layer_{layer_id}"
513656
names = (f"{prefix}_k", f"{prefix}_v")
@@ -534,7 +677,7 @@ def _runtime_buffer(exported_program: ExportedProgram, name: str, value):
534677
)
535678

536679
def _cell_buffer(self, exported_program, name: str, shape, dtype, device):
537-
"""A step buffer of the cell layout, declared once per program.
680+
"""A step buffer of a packed layout, declared once per program.
538681
539682
Unlike the pools these are small, so they carry real (zero) bytes at
540683
compile time: Inductor reads small constants while it compiles, and
@@ -580,6 +723,39 @@ def _lower_cell(
580723
example_args = tuple(n.meta["val"] for n in call_args)
581724
return _cell_step_fn(node.args[5], out_dtype), call_args, example_args
582725

726+
def _lower_batched_sequence(
727+
self, exported_program, node, layer, k_storage, v_storage, out_dtype
728+
):
729+
device = node.args[1].meta["val"].device
730+
seq_cols = self._cell_buffer(
731+
exported_program,
732+
OFFGRAPH_KV_SEQ_COLS_FQN,
733+
(self._manifest["max_write"],),
734+
torch.int64,
735+
device,
736+
)
737+
seq_info = self._cell_buffer(
738+
exported_program,
739+
OFFGRAPH_KV_SEQ_INFO_FQN,
740+
(self._manifest["max_segments"], 4),
741+
torch.int32,
742+
device,
743+
)
744+
call_args = (*node.args[0:4], k_storage, v_storage, seq_cols, seq_info)
745+
example_args = tuple(n.meta["val"] for n in call_args)
746+
ring = layer["policy"] == "ring"
747+
fn = _batched_sequence_step_fn(
748+
node.args[5],
749+
self._layer_capacity(layer),
750+
self._columns,
751+
layer["window"] if ring else 0,
752+
ring,
753+
self._manifest["causal"],
754+
out_dtype,
755+
self._manifest.get("context_window", False),
756+
)
757+
return fn, call_args, example_args
758+
583759
def _ring_mask(self, graph, position, buf_size: int, window: int, before):
584760
"""Emit the ring mask once and share it across layers that match.
585761
@@ -671,6 +847,10 @@ def __call__(self, exported_program: ExportedProgram) -> ExportedProgram:
671847
fn, call_args, example_args = self._lower_cell(
672848
exported_program, node, layer, k_storage, v_storage, out_dtype
673849
)
850+
elif self._batched_sequence_layout:
851+
fn, call_args, example_args = self._lower_batched_sequence(
852+
exported_program, node, layer, k_storage, v_storage, out_dtype
853+
)
674854
else:
675855
fn, call_args, example_args = self._lower_sequence(
676856
graph_module.graph, node, layer, k_storage, v_storage, out_dtype

‎backends/cuda/tests/targets.bzl‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -304,6 +304,32 @@ def define_common_targets(is_fbcode = False):
304304
),
305305
)
306306

307+
python_unittest_remote_gpu(
308+
name = "test_offgraph_kv_batched_sequence",
309+
srcs = [
310+
"test_offgraph_kv_batched_sequence.py",
311+
],
312+
visibility = [
313+
"//executorch/...",
314+
],
315+
deps = [
316+
"//caffe2:torch",
317+
"//executorch/backends/cuda:cuda_backend",
318+
"//executorch/backends/cuda:cuda_partitioner",
319+
"//executorch/backends/cuda:cuda_passes",
320+
"//executorch/exir:lib",
321+
"//executorch/exir/backend:compile_spec_schema",
322+
"//executorch/exir/dialects:lib",
323+
# The oracle is the neutral op itself, not a reference rebuilt here.
324+
"//executorch/extension/llm/cache:cache",
325+
],
326+
keep_gpu_sections = True,
327+
remote_execution = re_test_utils.remote_execution(
328+
platform = "gpu-remote-execution",
329+
subplatform = "A100-exclusive",
330+
),
331+
)
332+
307333
python_unittest(
308334
name = "test_quantized_gemm_family",
309335
srcs = [

0 commit comments

Comments
 (0)