[executorch][cuda] Test CudaExecutor on every step width up to 65

backends/cuda/batching/test/export_toy_decoder.py

@@ -18,6 +18,9 @@
 - ``device_sampling/``: dense's methods with their logits kept on the device,
   plus the ``sample`` and ``sample_argmax`` device samplers
   (extension/llm/batching/sampler.py).
+- ``mg_widths/``: Muse Glimmer's method set: forward_{1,2,4,...,64}, and
+  forward_others from 64 tokens selecting at least five rows, over longer
+  prompts, so steps of every width up to 65 can be formed.
 
 Each has ``expected.txt``: one ``prompt;continuation`` line per prompt, the
 greedy continuation computed eagerly with the neutral reference cache.
@@ -94,14 +97,14 @@ def forward(self, tokens, input_pos, logits_to_keep):
         return self.head(x[:, logits_to_keep, :])[0].float()
 
 
-def _manifest() -> str:
+def _manifest(max_step: int, max_cells: int) -> str:
     return json.dumps(
         {
             "version": 1,
             "layout": "cell",
-            "maximum_capacity": MAX_CELLS,
-            "max_cells": MAX_CELLS,
-            "max_write": MAX_STEP,
+            "maximum_capacity": max_cells,
+            "max_cells": max_cells,
+            "max_write": max_step,
             "layers": [
                 {"layer_id": i, "policy": "ring" if w else "flat", "window": w}
                 for i, w in enumerate(WINDOWS)
@@ -110,7 +113,7 @@ def _manifest() -> str:
     )
 
 
-def _greedy(model: ToyDecoder, prompt) -> list:
+def _greedy(model: ToyDecoder, prompt, new_tokens: int = NEW_TOKENS) -> list:
     from executorch.extension.llm.cache.reference_cache import (
         CacheConfig,
         LayerPolicy,
@@ -136,7 +139,7 @@ def _greedy(model: ToyDecoder, prompt) -> list:
             generated = []
             tokens = list(prompt)
             start = 0
-            for _ in range(NEW_TOKENS):
+            for _ in range(new_tokens):
                 step = torch.tensor([tokens[start:]])
                 logits = model(
                     step,
@@ -156,11 +159,31 @@ def _greedy(model: ToyDecoder, prompt) -> list:
     "dense": {"static": (1, 2, 4, 8), "min_rows": 1},
     "sparse": {"static": (1, 4), "min_rows": 5},
     "device_sampling": {"static": (1, 2, 4, 8), "min_rows": 1, "device_sampling": True},
+    # Each prompt plus its continuation fits MAX_CONTEXT; together they hold
+    # far more than the widest static step.
+    "mg_widths": {
+        "static": (1, 2, 4, 8, 16, 32, 64),
+        "min_rows": 5,
+        "max_step": 96,
+        "max_cells": 1024,
+        "new_tokens": 12,
+        "prompts": tuple(
+            [(7 * i + 3 * j) % VOCAB for j in range(n)]
+            for i, n in enumerate((1, 3, 7, 12, 20, 28, 40, 50))
+        ),
+    },
 }
 
 
 def export(
-    output_dir: str, static_widths, min_rows: int, device_sampling: bool = False
+    output_dir: str,
+    static_widths,
+    min_rows: int,
+    device_sampling: bool = False,
+    max_step: int = MAX_STEP,
+    max_cells: int = MAX_CELLS,
+    prompts=PROMPTS,
+    new_tokens: int = NEW_TOKENS,
 ) -> None:
     import torch._inductor.config as inductor_config
     from executorch.backends.cuda.cuda_backend import CudaBackend
@@ -194,7 +217,7 @@ def export(
     inductor_config.aot_inductor.precompile_headers = False
     torch.manual_seed(0)
     model = ToyDecoder().eval()
-    expected = [_greedy(model, prompt) for prompt in PROMPTS]
+    expected = [_greedy(model, prompt, new_tokens) for prompt in prompts]
 
     model = model.to(dtype=torch.bfloat16)
     long = {"dtype": torch.long}
@@ -207,8 +230,8 @@ def inputs(width: int):
         )
 
     widest_static = max(static_widths)
-    width = Dim("width", min=widest_static, max=MAX_STEP)
-    rows = Dim("rows", min=min_rows, max=MAX_STEP)
+    width = Dim("width", min=widest_static, max=max_step)
+    rows = Dim("rows", min=min_rows, max=max_step)
     programs = {}
     with torch.no_grad():
         for n in static_widths:
@@ -217,7 +240,7 @@ def inputs(width: int):
             )
         programs["forward_others"] = torch.export.export(
             model,
-            inputs(MAX_STEP),
+            inputs(max_step),
             dynamic_shapes=({1: width}, {0: width}, {0: rows}),
             strict=True,
         )
@@ -244,20 +267,22 @@ def partitioner(name: str) -> CudaPartitioner:
             [
                 CudaBackend.generate_method_name_compile_spec(name),
                 CompileSpec("low_memory_mode", b"ON"),
-                CompileSpec(OFFGRAPH_KV_COMPILE_SPEC, _manifest().encode()),
+                CompileSpec(
+                    OFFGRAPH_KV_COMPILE_SPEC, _manifest(max_step, max_cells).encode()
+                ),
                 CompileSpec(OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC, b"1:0"),
             ]
         )
 
     constant_methods = {
         **write_max_context_len(MAX_CONTEXT),
-        **write_max_seq_len(MAX_STEP),
+        **write_max_seq_len(max_step),
         **write_vocab_size(VOCAB),
         **write_logits_to_keep_mode("selected"),
         **write_cache_geometry(
             [N_KV_HEADS] * len(WINDOWS), [HEAD_DIM] * len(WINDOWS), list(WINDOWS)
         ),
-        "get_offgraph_kv_max_cells": MAX_CELLS,
+        "get_offgraph_kv_max_cells": max_cells,
     }
     if min_rows > 1:
         constant_methods["get_min_prefill_chunk"] = min_rows
@@ -290,7 +315,7 @@ def partitioner(name: str) -> CudaPartitioner:
         program.write_to_file(f)
     program.write_tensor_data_to_file(output_dir)
     with open(os.path.join(output_dir, "expected.txt"), "w") as f:
-        for prompt, continuation in zip(PROMPTS, expected):
+        for prompt, continuation in zip(prompts, expected):
             f.write(
                 " ".join(map(str, prompt))
                 + ";"
@@ -310,6 +335,10 @@ def main() -> None:
             artifact["static"],
             artifact["min_rows"],
             artifact.get("device_sampling", False),
+            artifact.get("max_step", MAX_STEP),
+            artifact.get("max_cells", MAX_CELLS),
+            artifact.get("prompts", PROMPTS),
+            artifact.get("new_tokens", NEW_TOKENS),
         )
 
 

backends/cuda/batching/test/test_cuda_executor.cpp

@@ -10,7 +10,9 @@
 // under $ET_CUDA_BATCHING_TOY_DIR, and checks every generation against the
 // eager greedy continuation it recorded. Runs once per artifact: forward_{1,2,
 // 4,8} + forward_others, and the sparser forward_{1,4} + forward_others with a
-// five-row selector minimum, where steps pad across a gap.
+// five-row selector minimum, where steps pad across a gap. mg_widths carries
+// Muse Glimmer's forward_{1,...,64} + forward_others and drives steps of every
+// width up to 65.
 
 #include <executorch/backends/cuda/batching/cuda_executor.h>
 #include <executorch/extension/llm/batching/decode_first_scheduler.h>
@@ -23,8 +25,11 @@
 #include <algorithm>
 #include <cstdlib>
 #include <fstream>
+#include <map>
 #include <memory>
 #include <mutex>
+#include <numeric>
+#include <random>
 #include <sstream>
 #include <string>
 #include <vector>
@@ -597,3 +602,186 @@ INSTANTIATE_TEST_SUITE_P(
     Toy,
     CudaExecutorDeviceSamplingTest,
     ::testing::Values("device_sampling"));
+
+namespace {
+
+// Must match export_toy_decoder.py's mg_widths artifact.
+constexpr int kWidthsMaxStep = 96;
+constexpr int kWidthsMaxCells = 1024;
+
+int next_pow2(int n) {
+  int p = 1;
+  while (p < n) {
+    p *= 2;
+  }
+  return p;
+}
+
+// Muse Glimmer's method set: forward_{1,2,...,64} and forward_others from 64.
+class CudaExecutorWidthsTest : public CudaExecutorTest {};
+
+} // namespace
+
+// Steps of every width 1..65 (and 96, 97), each a random mix of sequences and
+// chunk lengths: a static method's step pads up to the next power of two and
+// drops the padding's rows, 65..96 runs forward_others unpadded, and 97 slices.
+// Every output a step asks for is the eager greedy successor, so padding
+// neither reaches a real row nor disturbs a later step through the cache.
+TEST_P(CudaExecutorWidthsTest, EveryStepWidthUpTo65MatchesEagerGreedy) {
+  const int sessions = static_cast<int>(cases_.size());
+  ASSERT_LT(sessions * kMaxContext, kWidthsMaxCells);
+  for (const bool graphs : {true, false}) {
+    SCOPED_TRACE(graphs ? "cuda graphs" : "eager");
+    cb::CudaExecutorOptions options;
+    options.cuda_graph_for_static_methods = graphs;
+    auto created = cb::CudaExecutor::create(
+        module(),
+        sessions,
+        kMaxContext,
+        kBFloat16,
+        /*initial_capacity=*/16,
+        options);
+    ASSERT_EQ(created.error(), Error::Ok);
+    auto& exec = *created.get();
+    ASSERT_TRUE(exec.initialize());
+    ASSERT_EQ(
+        exec.preferred_batch_tokens(), static_cast<size_t>(kWidthsMaxStep));
+
+    // Teacher forcing: each sequence is its prompt and eager continuation, so
+    // a step may end any chunk at a position whose successor is known.
+    std::vector<std::shared_ptr<const std::vector<batching::Token>>> seqs;
+    for (const Case& c : cases_) {
+      auto full = c.prompt;
+      full.insert(full.end(), c.expected.begin(), c.expected.end());
+      ASSERT_LE(full.size(), static_cast<size_t>(kMaxContext));
+      seqs.push_back(std::make_shared<const std::vector<batching::Token>>(
+          std::move(full)));
+    }
+    std::vector<batching::SessionId> sids(sessions, -1);
+    std::vector<size_t> fed(sessions, 0);
+    auto reset = [&] {
+      for (int i = 0; i < sessions; ++i) {
+        if (sids[i] >= 0) {
+          exec.close_session(sids[i]);
+        }
+        auto sid = exec.open_session();
+        ASSERT_TRUE(sid.has_value());
+        sids[i] = *sid;
+        exec.set_sampling(sids[i], {}, 0);
+        fed[i] = 0;
+      }
+    };
+    // Each sequence's last token has no known successor, so it is never fed.
+    auto remaining = [&](int i) { return seqs[i]->size() - 1 - fed[i]; };
+    reset();
+
+    std::vector<int> widths;
+    for (int round = 0; round < 6; ++round) {
+      for (int n = 1; n <= 65; ++n) {
+        widths.push_back(n);
+      }
+      widths.push_back(kWidthsMaxStep);
+      widths.push_back(kWidthsMaxStep + 1);
+    }
+    std::mt19937 rng(1234);
+    std::shuffle(widths.begin(), widths.end(), rng);
+
+    size_t checked = 0;
+    for (const int n : widths) {
+      SCOPED_TRACE("step of " + std::to_string(n));
+      size_t total = 0;
+      for (int i = 0; i < sessions; ++i) {
+        total += remaining(i);
+      }
+      if (total < static_cast<size_t>(n)) {
+        reset();
+      }
+      // A random subset and split of the sequences that sums to n, each
+      // sequence at most once.
+      std::vector<int> order(sessions);
+      std::iota(order.begin(), order.end(), 0);
+      std::shuffle(order.begin(), order.end(), rng);
+      size_t after = 0;
+      for (const int i : order) {
+        after += remaining(i);
+      }
+      batching::BatchInput batch;
+      std::vector<int> owner;
+      size_t left = static_cast<size_t>(n);
+      for (const int i : order) {
+        after -= remaining(i);
+        if (left == 0 || remaining(i) == 0) {
+          continue;
+        }
+        const size_t most = std::min(remaining(i), left);
+        // Enough that the sequences after this one can cover the rest.
+        const size_t least = left > after ? left - after : 0;
+        if (most < least || (least == 0 && rng() % 3 == 0)) {
+          continue;
+        }
+        const size_t take = std::uniform_int_distribution<size_t>(
+            std::max<size_t>(least, 1), most)(rng);
+        const size_t end = fed[i] + take;
+        // Ask for the chunk's last row whenever its successor is known, but
+        // not always, so some steps want none.
+        const bool known = end >= cases_[i].prompt.size();
+        batch.inputs.push_back(batching::Input{
+            sids[i], known && rng() % 4 != 0, fed[i], take, seqs[i], 0});
+        owner.push_back(i);
+        left -= take;
+      }
+      ASSERT_EQ(left, 0u);
+      ASSERT_EQ(batch.size(), static_cast<size_t>(n));
+
+      const auto before = exec.method_calls();
+      batching::BatchOutput out;
+      ASSERT_TRUE(exec.execute(batch, out));
+      const auto calls = exec.method_calls();
+
+      // The method(s) the width runs through.
+      std::map<std::string, std::uint64_t> ran;
+      for (size_t m = 0; m < calls.size(); ++m) {
+        if (calls[m].calls != before[m].calls) {
+          ran[calls[m].name] = calls[m].calls - before[m].calls;
+        }
+      }
+      std::map<std::string, std::uint64_t> want;
+      if (n <= 64) {
+        want["forward_" + std::to_string(next_pow2(n))] = 1;
+      } else if (n <= kWidthsMaxStep) {
+        want["forward_others"] = 1;
+      } else {
+        want["forward_others"] = 1;
+        want["forward_" + std::to_string(next_pow2(n - kWidthsMaxStep))] = 1;
+      }
+      EXPECT_EQ(ran, want);
+
+      ASSERT_EQ(out.outputs.size(), batch.inputs.size());
+      for (size_t j = 0; j < batch.inputs.size(); ++j) {
+        const int i = owner[j];
+        const auto& input = batch.inputs[j];
+        const size_t end = input.offset + input.size;
+        if (input.produce_output) {
+          ASSERT_TRUE(out.outputs[j].has_value()) << "sequence " << i;
+          EXPECT_EQ(out.outputs[j]->sid, sids[i]);
+          ASSERT_EQ(out.outputs[j]->tokens.size(), 1u);
+          EXPECT_EQ(out.outputs[j]->tokens[0], (*seqs[i])[end])
+              << "sequence " << i << " after " << end << " tokens";
+          ++checked;
+        } else {
+          EXPECT_FALSE(out.outputs[j].has_value()) << "sequence " << i;
+        }
+        fed[i] = end;
+      }
+    }
+    EXPECT_GT(checked, widths.size());
+    for (const batching::SessionId sid : sids) {
+      exec.close_session(sid);
+    }
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    Toy,
+    CudaExecutorWidthsTest,
+    ::testing::Values("mg_widths"));