1414# Imported for its registration side effect: the decomposition calls
1515# torch.ops.triton.sdpa, which only exists once this module has been loaded.
1616from 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+ )
1720from executorch .backends .transforms .utils import create_constant_placeholder
1821from executorch .exir .dialects ._ops import ops as exir_ops
1922from torch .export import ExportedProgram
3639OFFGRAPH_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+
3962def 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+
243369def _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
0 commit comments