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"));