From 0be3c22784630633d3d8d798c4554ab5efe24280 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Mon, 19 Oct 2026 13:07:58 -0700 Subject: [PATCH] [executorch][cuda] Test CudaExecutor on every step width up to 65 Pull Request resolved: https://github.com/pytorch/executorch/pull/23819 CudaExecutor chooses the narrowest static forward_{N} that holds a step, pads unused rows into the KV scratch cell without returning their logits, and uses forward_others with slicing for wider steps. Extend the toy artifact to the Muse-sized method set forward_{1,2,4,...,64} plus forward_others, with eight independently referenced prompts and continuations. EveryStepWidthUpTo65MatchesEagerGreedy covers every width 1..65 plus 96 and 97, shuffled repeatedly, with CUDA graphs both on and off. Mix sequence owners and chunk lengths, teacher-force eager reference histories, and vary which chunks request output. Assert the selected method sequence, exact output ownership/count and eager greedy token at each requested position. This is correctness coverage of static padding, dynamic widths and slicing, without serving performance reporting. ghstack-source-id: 451855103 @exported-using-ghexport Differential Revision: [D125110873](https://our.internmc.facebook.com/intern/diff/D125110873/) --- .../cuda/batching/test/export_toy_decoder.py | 59 ++++-- .../cuda/batching/test/test_cuda_executor.cpp | 190 +++++++++++++++++- 2 files changed, 233 insertions(+), 16 deletions(-) diff --git a/backends/cuda/batching/test/export_toy_decoder.py b/backends/cuda/batching/test/export_toy_decoder.py index d591fe8d784..1ba2acdb50f 100644 --- a/backends/cuda/batching/test/export_toy_decoder.py +++ b/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), ) diff --git a/backends/cuda/batching/test/test_cuda_executor.cpp b/backends/cuda/batching/test/test_cuda_executor.cpp index 34c54930319..062861075a7 100644 --- a/backends/cuda/batching/test/test_cuda_executor.cpp +++ b/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 #include @@ -23,8 +25,11 @@ #include #include #include +#include #include #include +#include +#include #include #include #include @@ -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(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(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>> 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(kMaxContext)); + seqs.push_back(std::make_shared>( + std::move(full))); + } + std::vector sids(sessions, -1); + std::vector 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 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(n)) { + reset(); + } + // A random subset and split of the sequences that sums to n, each + // sequence at most once. + std::vector 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 owner; + size_t left = static_cast(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( + std::max(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(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 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 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"));