[executorch][cuda] Lower off-graph KV to the batched-sequence layout

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

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

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 = [