[executorch][muse] Serve Muse Glimmer CUDA batching on a batched-sequence cache

backends/cuda/batching/cuda_executor.cpp

@@ -363,21 +363,40 @@ Result<std::unique_ptr<CudaExecutor>> CudaExecutor::create(
       NotSupported,
       "CudaExecutor: the program must select its logits rows");
   ET_ASSIGN_OR_RETURN(geometry, metadata::read_cache_geometry(*module));
+  // The program fixes its cache layout: a private history per sequence when it
+  // publishes its sequence count, else one shared pool of cells.
   ET_ASSIGN_OR_RETURN(
-      max_cells,
-      metadata::detail::read_required_positive_int(*module, kMaxCellsMethod));
-
-  ET_CHECK_OR_RETURN_ERROR(
-      max_cells > 1 && max_cells <= std::numeric_limits<int>::max(),
-      InvalidProgram,
-      "CudaExecutor: cell count must fit an int and leave a padding row");
-  if (max_session_tokens == 0) {
-    max_session_tokens = static_cast<int>(
-        std::min(max_context_length, (max_cells - 1) / max_sessions));
+      declared_seqs,
+      metadata::detail::read_int_method(*module, kMaxSeqsMethod));
+  const bool sequence_layout = declared_seqs.has_value();
+  std::int64_t max_cells = 0;
+  if (sequence_layout) {
+    ET_CHECK_OR_RETURN_ERROR(
+        *declared_seqs > 0 &&
+            *declared_seqs <= std::numeric_limits<int>::max() &&
+            max_context_length <= std::numeric_limits<int>::max(),
+        InvalidProgram,
+        "CudaExecutor: invalid sequence count or context length");
+    if (max_session_tokens == 0) {
+      max_session_tokens = static_cast<int>(max_context_length);
+    }
+  } else {
+    ET_ASSIGN_OR_RETURN(
+        cells,
+        metadata::detail::read_required_positive_int(*module, kMaxCellsMethod));
+    max_cells = cells;
     ET_CHECK_OR_RETURN_ERROR(
-        max_session_tokens > 0,
-        InvalidArgument,
-        "CudaExecutor: no cell budget remains for each resident session");
+        max_cells > 1 && max_cells <= std::numeric_limits<int>::max(),
+        InvalidProgram,
+        "CudaExecutor: cell count must fit an int and leave a padding row");
+    if (max_session_tokens == 0) {
+      max_session_tokens = static_cast<int>(
+          std::min(max_context_length, (max_cells - 1) / max_sessions));
+      ET_CHECK_OR_RETURN_ERROR(
+          max_session_tokens > 0,
+          InvalidArgument,
+          "CudaExecutor: no cell budget remains for each resident session");
+    }
   }
 
   // Every forward_{N} and forward_others the program exports, narrowest first.
@@ -444,11 +463,11 @@ Result<std::unique_ptr<CudaExecutor>> CudaExecutor::create(
       "CudaExecutor: more than one dynamic method");
   const int max_step_tokens = methods.back().max_tokens;
   ET_CHECK_OR_RETURN_ERROR(
-      max_step_tokens < max_cells,
+      sequence_layout ? max_step_tokens <= max_context_length
+                      : max_step_tokens < max_cells,
       InvalidProgram,
-      "CudaExecutor: a %d-token step does not fit %" PRId64 " cells",
-      max_step_tokens,
-      max_cells);
+      "CudaExecutor: a %d-token step does not fit the cache",
+      max_step_tokens);
   // A dynamic method whose kernels switch at a small width is exported from
   // more selected rows than one; its selector is padded up to that.
   int min_selected_rows = 1;
@@ -483,7 +502,9 @@ Result<std::unique_ptr<CudaExecutor>> CudaExecutor::create(
   // them all, or a step could find no free cell mid-generation. One of the
   // program's cells is the padding scratch row, never handed out.
   ET_CHECK_OR_RETURN_ERROR(
-      static_cast<std::int64_t>(max_sessions) * max_session_tokens < max_cells,
+      sequence_layout ||
+          static_cast<std::int64_t>(max_sessions) * max_session_tokens <
+              max_cells,
       InvalidArgument,
       "CudaExecutor: %d sessions of %d tokens exceed the program's %" PRId64
       " cells, one of which is reserved",
@@ -493,10 +514,16 @@ Result<std::unique_ptr<CudaExecutor>> CudaExecutor::create(
 
   // The pool is the program's: its size and widest step fix the shapes the
   // program declared for the step buffers.
+  // A sequence's columns hold the program's context; every sequence has its
+  // own.
   llm_cache::CacheConfig cfg{};
-  cfg.capacity = static_cast<int>(max_cells);
+  cfg.capacity =
+      static_cast<int>(sequence_layout ? max_context_length : max_cells);
   cfg.kv_dtype = kv_dtype;
   cfg.max_write = max_step_tokens;
+  if (sequence_layout) {
+    cfg.max_seqs = static_cast<int>(*declared_seqs);
+  }
   if (initial_capacity >= 0) {
     cfg.initial_capacity = std::min(initial_capacity, cfg.capacity);
   }
@@ -505,7 +532,11 @@ Result<std::unique_ptr<CudaExecutor>> CudaExecutor::create(
       InvalidProgram,
       "CudaExecutor: the program's layout is unusable");
   auto built = llm_cache::CacheFactory::global().build(
-      backend_id, llm_cache::kind::kBatchedCell, geometry, cfg);
+      backend_id,
+      sequence_layout ? llm_cache::kind::kBatchedSequence
+                      : llm_cache::kind::kBatchedCell,
+      geometry,
+      cfg);
   ET_CHECK_OK_OR_RETURN_ERROR(built.error());
   std::shared_ptr<llm_cache::Cache> cache = built.get();
   auto* const ctl = cache->as<llm_cache::BatchControl>();
@@ -700,9 +731,8 @@ bool CudaExecutor::execute_raw(
         ? slice.width
         : std::max<std::size_t>(selected.size(), min_selected_rows_);
     selected.resize(rows, selected.empty() ? 0 : selected.back());
-    if (!ctl_->declare_step(
-            std::vector<std::int32_t>(
-                step.seq_ids.begin() + off, step.seq_ids.begin() + off + n))) {
+    if (!ctl_->declare_step(std::vector<std::int32_t>(
+            step.seq_ids.begin() + off, step.seq_ids.begin() + off + n))) {
       return false;
     }
     std::vector<std::int64_t> token_values(
@@ -1083,9 +1113,8 @@ bool CudaExecutor::execute_embeddings(
     const std::size_t rows =
         std::max<std::size_t>(selected.size(), min_selected_rows_);
     selected.resize(rows, selected.empty() ? 0 : selected.back());
-    if (!ctl_->declare_step(
-            std::vector<std::int32_t>(
-                step.seq_ids.begin() + off, step.seq_ids.begin() + off + n))) {
+    if (!ctl_->declare_step(std::vector<std::int32_t>(
+            step.seq_ids.begin() + off, step.seq_ids.begin() + off + n))) {
       return false;
     }
     // Rows past n are padding the cache writes to its scratch cell.

backends/cuda/batching/cuda_executor.h

@@ -48,6 +48,9 @@ inline constexpr char kStaticMethodPrefix[] = "forward_";
 inline constexpr char kDynamicMethod[] = "forward_others";
 // The cell layout's pool size, which fixes the shape of its step buffers.
 inline constexpr char kMaxCellsMethod[] = "get_offgraph_kv_max_cells";
+// Published instead of the cell count by a program whose KV cache keeps a
+// private history per sequence.
+inline constexpr char kMaxSeqsMethod[] = "get_offgraph_kv_max_seqs";
 // Optional: the fewest selected rows forward_others was exported for. Its
 // selector is padded up to it. Defaults to 1.
 inline constexpr char kMinSelectedRowsMethod[] = "get_min_prefill_chunk";

examples/models/muse-glimmer/export/export_dflash.py

@@ -14,6 +14,7 @@
 
 import argparse
 import gc
+import json
 from typing import TYPE_CHECKING
 
 import torch
@@ -173,6 +174,8 @@ def export_dflash(
     draft_max_cells: int | None = None,
     max_draft_owners: int = 4,
     forward_sizes=None,
+    kv_layout: str = "cell",
+    max_kv_seqs: int = 4,
 ) -> None:
     """Export DFlash target + draft to one CUDA or MLX .pte.
 
@@ -278,6 +281,8 @@ def export_dflash(
             vision_model=vision_model,
             pos_embed_table=pos_embed_table,
             max_vision_patches=max_vision_patches,
+            kv_layout=kv_layout,
+            max_kv_seqs=max_kv_seqs,
         )
         return
 
@@ -305,9 +310,19 @@ def prepare_cuda_packed_draft(
     max_step_tokens: int,
     max_cells: int,
     max_owners: int,
+    kv_layout: str = "cell",
+    max_kv_seqs: int = 4,
 ):
-    """Trace the packed draft and publish its independent cell geometry."""
-    from executorch.examples.models.muse_glimmer.export.export_solo import cell_manifest
+    """Trace the packed draft and publish its independent cache geometry.
+
+    On the sequence layout the draft keeps one private history per resident
+    session, and each owner's block attends all of it and the whole block.
+    """
+    from executorch.examples.models.muse_glimmer.export.export_solo import (
+        batched_sequence_manifest,
+        cell_manifest,
+        MAX_KV_SEQS_METHOD,
+    )
     from executorch.examples.models.muse_glimmer.model.dflash_model import (
         DFlashPackedExportModel,
     )
@@ -330,7 +345,8 @@ def prepare_cuda_packed_draft(
     if type(max_step_tokens) is not int or max_step_tokens < 2:
         raise ValueError("packed DFlash max_step_tokens must be an integer >= 2")
     max_write = max_step_tokens + max_owners * block
-    if type(max_cells) is not int or max_cells <= max_write:
+    sequence_layout = kv_layout == "sequence"
+    if not sequence_layout and (type(max_cells) is not int or max_cells <= max_write):
         raise ValueError("draft max_cells must exceed context plus all owner blocks")
     if config.dim != target_config.dim:
         raise ValueError("draft and target hidden dimensions must agree")
@@ -345,8 +361,11 @@ def prepare_cuda_packed_draft(
 
     convert_quantized_tensors_for_cuda(draft_model)
     draft_model.to(torch.bfloat16)
-    manifest = cell_manifest(
-        enable_dflash_offgraph_kv_cache(draft_model, max_write), max_cells
+    single_manifest = enable_dflash_offgraph_kv_cache(draft_model, max_write)
+    manifest = (
+        batched_sequence_manifest(single_manifest, max_kv_seqs, causal=False)
+        if sequence_layout
+        else cell_manifest(single_manifest, max_cells)
     )
     combined = DFlashPackedExportModel(target_model, draft_model)
     owners = (
@@ -414,14 +433,21 @@ def prepare_cuda_packed_draft(
             "get_dflash_target_layers": torch.tensor(
                 config.target_layers, dtype=torch.int32
             ),
-            "get_draft_offgraph_kv_max_cells": max_cells,
             "get_draft_max_step_tokens": max_write,
             "get_draft_max_owners": max_owners,
             "get_draft_dflash_block_size": block,
             "get_dflash_target_cache_key": "muse_glimmer_target",
             "get_dflash_draft_cache_key": "muse_glimmer_draft",
         }
     )
+    if sequence_layout:
+        metadata[MAX_KV_SEQS_METHOD.replace("get_", "get_draft_", 1)] = max_kv_seqs
+        # Fixes each sequence column's rows, which the runtime must match.
+        metadata["get_draft_offgraph_kv_seq_capacity"] = json.loads(single_manifest)[
+            "maximum_capacity"
+        ]
+    else:
+        metadata["get_draft_offgraph_kv_max_cells"] = max_cells
     return programs, manifest, metadata
 
 
@@ -1126,6 +1152,11 @@ def main() -> None:
     parser.add_argument("--max-cells", type=int, default=65536)
     parser.add_argument("--draft-max-cells", type=int, default=None)
     parser.add_argument("--max-draft-owners", type=int, default=4)
+    from executorch.examples.models.muse_glimmer.export.export_solo import (
+        add_kv_layout_arguments,
+    )
+
+    add_kv_layout_arguments(parser)
     parser.add_argument("--forward-sizes", default="1,2,4,8,16,32,64")
     args = parser.parse_args()
 
@@ -1161,6 +1192,8 @@ def main() -> None:
         draft_max_cells=args.draft_max_cells,
         max_draft_owners=args.max_draft_owners,
         forward_sizes=[int(n) for n in args.forward_sizes.split(",") if n],
+        kv_layout=args.kv_layout,
+        max_kv_seqs=args.max_kv_seqs,
     )