Skip to content

Commit f068fa8

Browse files
authored
[executorch][cuda] Test CudaExecutor on every step width up to 65
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #23735 * #23734 * #23974 * #23973 * #23972 * #23971 * #23970 * #23969 * #23968 * #23967 * #23966 * #23965 * #23964 * #23963 * #23962 * #23961 * #23960 * #23959 * #23958 * #23985 * #23984 * #23983 * #23982 * #23981 * #23980 * #23979 * #23978 * __->__ #23819 * #23818 * #23733 * #23732 * #23731 * #23730 * #23729 * #23727 * #23728 CudaExecutor runs a step of T tokens through the narrowest static `forward_{N}` that holds it, padding T up to N (for power-of-two widths, the next power of two), and sends the padding rows' K/V to the pool's scratch row and their logits rows nowhere; wider steps run `forward_others` unpadded, and steps wider than it slice. The GPU tests drove only the toy's forward_{1,2,4,8} through the runner, so the widths Muse Glimmer's forward_{16,32,64} pad to, and the step just past the widest static method, were never run. - `export_toy_decoder.py` adds `mg_widths/`: Muse Glimmer's method set, forward_{1,2,4,...,64} and forward_others from 64 tokens selecting at least five rows, over eight prompts of 1..50 tokens with 12-token continuations. - `EveryStepWidthUpTo65MatchesEagerGreedy`: every width 1..65 (and 96, 97) six times, shuffled. Each step is a random mix of the sequences (each at most once) with random chunk lengths summing to the width, some chunks asking for their next token and some not; sequences are teacher-forced along their prompt and eager continuation. Each step must run the expected method(s) (`forward_{next pow2}`; forward_others for 65..96; forward_others + forward_1 for 97), return exactly the outputs asked for, and every output must be the eager greedy token at that position. Runs with CUDA graphs on and off. Differential Revision: D125110873 Pull Request resolved: #23819
1 parent 549114b commit f068fa8

2 files changed

Lines changed: 233 additions & 16 deletions

File tree

‎backends/cuda/batching/test/export_toy_decoder.py‎

Lines changed: 44 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,9 @@
1818
- ``device_sampling/``: dense's methods with their logits kept on the device,
1919
plus the ``sample`` and ``sample_argmax`` device samplers
2020
(extension/llm/batching/sampler.py).
21+
- ``mg_widths/``: Muse Glimmer's method set: forward_{1,2,4,...,64}, and
22+
forward_others from 64 tokens selecting at least five rows, over longer
23+
prompts, so steps of every width up to 65 can be formed.
2124
2225
Each has ``expected.txt``: one ``prompt;continuation`` line per prompt, the
2326
greedy continuation computed eagerly with the neutral reference cache.
@@ -94,14 +97,14 @@ def forward(self, tokens, input_pos, logits_to_keep):
9497
return self.head(x[:, logits_to_keep, :])[0].float()
9598

9699

97-
def _manifest() -> str:
100+
def _manifest(max_step: int, max_cells: int) -> str:
98101
return json.dumps(
99102
{
100103
"version": 1,
101104
"layout": "cell",
102-
"maximum_capacity": MAX_CELLS,
103-
"max_cells": MAX_CELLS,
104-
"max_write": MAX_STEP,
105+
"maximum_capacity": max_cells,
106+
"max_cells": max_cells,
107+
"max_write": max_step,
105108
"layers": [
106109
{"layer_id": i, "policy": "ring" if w else "flat", "window": w}
107110
for i, w in enumerate(WINDOWS)
@@ -110,7 +113,7 @@ def _manifest() -> str:
110113
)
111114

112115

113-
def _greedy(model: ToyDecoder, prompt) -> list:
116+
def _greedy(model: ToyDecoder, prompt, new_tokens: int = NEW_TOKENS) -> list:
114117
from executorch.extension.llm.cache.reference_cache import (
115118
CacheConfig,
116119
LayerPolicy,
@@ -136,7 +139,7 @@ def _greedy(model: ToyDecoder, prompt) -> list:
136139
generated = []
137140
tokens = list(prompt)
138141
start = 0
139-
for _ in range(NEW_TOKENS):
142+
for _ in range(new_tokens):
140143
step = torch.tensor([tokens[start:]])
141144
logits = model(
142145
step,
@@ -156,11 +159,31 @@ def _greedy(model: ToyDecoder, prompt) -> list:
156159
"dense": {"static": (1, 2, 4, 8), "min_rows": 1},
157160
"sparse": {"static": (1, 4), "min_rows": 5},
158161
"device_sampling": {"static": (1, 2, 4, 8), "min_rows": 1, "device_sampling": True},
162+
# Each prompt plus its continuation fits MAX_CONTEXT; together they hold
163+
# far more than the widest static step.
164+
"mg_widths": {
165+
"static": (1, 2, 4, 8, 16, 32, 64),
166+
"min_rows": 5,
167+
"max_step": 96,
168+
"max_cells": 1024,
169+
"new_tokens": 12,
170+
"prompts": tuple(
171+
[(7 * i + 3 * j) % VOCAB for j in range(n)]
172+
for i, n in enumerate((1, 3, 7, 12, 20, 28, 40, 50))
173+
),
174+
},
159175
}
160176

161177

162178
def export(
163-
output_dir: str, static_widths, min_rows: int, device_sampling: bool = False
179+
output_dir: str,
180+
static_widths,
181+
min_rows: int,
182+
device_sampling: bool = False,
183+
max_step: int = MAX_STEP,
184+
max_cells: int = MAX_CELLS,
185+
prompts=PROMPTS,
186+
new_tokens: int = NEW_TOKENS,
164187
) -> None:
165188
import torch._inductor.config as inductor_config
166189
from executorch.backends.cuda.cuda_backend import CudaBackend
@@ -194,7 +217,7 @@ def export(
194217
inductor_config.aot_inductor.precompile_headers = False
195218
torch.manual_seed(0)
196219
model = ToyDecoder().eval()
197-
expected = [_greedy(model, prompt) for prompt in PROMPTS]
220+
expected = [_greedy(model, prompt, new_tokens) for prompt in prompts]
198221

199222
model = model.to(dtype=torch.bfloat16)
200223
long = {"dtype": torch.long}
@@ -207,8 +230,8 @@ def inputs(width: int):
207230
)
208231

209232
widest_static = max(static_widths)
210-
width = Dim("width", min=widest_static, max=MAX_STEP)
211-
rows = Dim("rows", min=min_rows, max=MAX_STEP)
233+
width = Dim("width", min=widest_static, max=max_step)
234+
rows = Dim("rows", min=min_rows, max=max_step)
212235
programs = {}
213236
with torch.no_grad():
214237
for n in static_widths:
@@ -217,7 +240,7 @@ def inputs(width: int):
217240
)
218241
programs["forward_others"] = torch.export.export(
219242
model,
220-
inputs(MAX_STEP),
243+
inputs(max_step),
221244
dynamic_shapes=({1: width}, {0: width}, {0: rows}),
222245
strict=True,
223246
)
@@ -244,20 +267,22 @@ def partitioner(name: str) -> CudaPartitioner:
244267
[
245268
CudaBackend.generate_method_name_compile_spec(name),
246269
CompileSpec("low_memory_mode", b"ON"),
247-
CompileSpec(OFFGRAPH_KV_COMPILE_SPEC, _manifest().encode()),
270+
CompileSpec(
271+
OFFGRAPH_KV_COMPILE_SPEC, _manifest(max_step, max_cells).encode()
272+
),
248273
CompileSpec(OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC, b"1:0"),
249274
]
250275
)
251276

252277
constant_methods = {
253278
**write_max_context_len(MAX_CONTEXT),
254-
**write_max_seq_len(MAX_STEP),
279+
**write_max_seq_len(max_step),
255280
**write_vocab_size(VOCAB),
256281
**write_logits_to_keep_mode("selected"),
257282
**write_cache_geometry(
258283
[N_KV_HEADS] * len(WINDOWS), [HEAD_DIM] * len(WINDOWS), list(WINDOWS)
259284
),
260-
"get_offgraph_kv_max_cells": MAX_CELLS,
285+
"get_offgraph_kv_max_cells": max_cells,
261286
}
262287
if min_rows > 1:
263288
constant_methods["get_min_prefill_chunk"] = min_rows
@@ -290,7 +315,7 @@ def partitioner(name: str) -> CudaPartitioner:
290315
program.write_to_file(f)
291316
program.write_tensor_data_to_file(output_dir)
292317
with open(os.path.join(output_dir, "expected.txt"), "w") as f:
293-
for prompt, continuation in zip(PROMPTS, expected):
318+
for prompt, continuation in zip(prompts, expected):
294319
f.write(
295320
" ".join(map(str, prompt))
296321
+ ";"
@@ -310,6 +335,10 @@ def main() -> None:
310335
artifact["static"],
311336
artifact["min_rows"],
312337
artifact.get("device_sampling", False),
338+
artifact.get("max_step", MAX_STEP),
339+
artifact.get("max_cells", MAX_CELLS),
340+
artifact.get("prompts", PROMPTS),
341+
artifact.get("new_tokens", NEW_TOKENS),
313342
)
314343

315344

‎backends/cuda/batching/test/test_cuda_executor.cpp‎

Lines changed: 189 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,9 @@
1010
// under $ET_CUDA_BATCHING_TOY_DIR, and checks every generation against the
1111
// eager greedy continuation it recorded. Runs once per artifact: forward_{1,2,
1212
// 4,8} + forward_others, and the sparser forward_{1,4} + forward_others with a
13-
// five-row selector minimum, where steps pad across a gap.
13+
// five-row selector minimum, where steps pad across a gap. mg_widths carries
14+
// Muse Glimmer's forward_{1,...,64} + forward_others and drives steps of every
15+
// width up to 65.
1416

1517
#include <executorch/backends/cuda/batching/cuda_executor.h>
1618
#include <executorch/extension/llm/batching/decode_first_scheduler.h>
@@ -23,8 +25,11 @@
2325
#include <algorithm>
2426
#include <cstdlib>
2527
#include <fstream>
28+
#include <map>
2629
#include <memory>
2730
#include <mutex>
31+
#include <numeric>
32+
#include <random>
2833
#include <sstream>
2934
#include <string>
3035
#include <vector>
@@ -597,3 +602,186 @@ INSTANTIATE_TEST_SUITE_P(
597602
Toy,
598603
CudaExecutorDeviceSamplingTest,
599604
::testing::Values("device_sampling"));
605+
606+
namespace {
607+
608+
// Must match export_toy_decoder.py's mg_widths artifact.
609+
constexpr int kWidthsMaxStep = 96;
610+
constexpr int kWidthsMaxCells = 1024;
611+
612+
int next_pow2(int n) {
613+
int p = 1;
614+
while (p < n) {
615+
p *= 2;
616+
}
617+
return p;
618+
}
619+
620+
// Muse Glimmer's method set: forward_{1,2,...,64} and forward_others from 64.
621+
class CudaExecutorWidthsTest : public CudaExecutorTest {};
622+
623+
} // namespace
624+
625+
// Steps of every width 1..65 (and 96, 97), each a random mix of sequences and
626+
// chunk lengths: a static method's step pads up to the next power of two and
627+
// drops the padding's rows, 65..96 runs forward_others unpadded, and 97 slices.
628+
// Every output a step asks for is the eager greedy successor, so padding
629+
// neither reaches a real row nor disturbs a later step through the cache.
630+
TEST_P(CudaExecutorWidthsTest, EveryStepWidthUpTo65MatchesEagerGreedy) {
631+
const int sessions = static_cast<int>(cases_.size());
632+
ASSERT_LT(sessions * kMaxContext, kWidthsMaxCells);
633+
for (const bool graphs : {true, false}) {
634+
SCOPED_TRACE(graphs ? "cuda graphs" : "eager");
635+
cb::CudaExecutorOptions options;
636+
options.cuda_graph_for_static_methods = graphs;
637+
auto created = cb::CudaExecutor::create(
638+
module(),
639+
sessions,
640+
kMaxContext,
641+
kBFloat16,
642+
/*initial_capacity=*/16,
643+
options);
644+
ASSERT_EQ(created.error(), Error::Ok);
645+
auto& exec = *created.get();
646+
ASSERT_TRUE(exec.initialize());
647+
ASSERT_EQ(
648+
exec.preferred_batch_tokens(), static_cast<size_t>(kWidthsMaxStep));
649+
650+
// Teacher forcing: each sequence is its prompt and eager continuation, so
651+
// a step may end any chunk at a position whose successor is known.
652+
std::vector<std::shared_ptr<const std::vector<batching::Token>>> seqs;
653+
for (const Case& c : cases_) {
654+
auto full = c.prompt;
655+
full.insert(full.end(), c.expected.begin(), c.expected.end());
656+
ASSERT_LE(full.size(), static_cast<size_t>(kMaxContext));
657+
seqs.push_back(std::make_shared<const std::vector<batching::Token>>(
658+
std::move(full)));
659+
}
660+
std::vector<batching::SessionId> sids(sessions, -1);
661+
std::vector<size_t> fed(sessions, 0);
662+
auto reset = [&] {
663+
for (int i = 0; i < sessions; ++i) {
664+
if (sids[i] >= 0) {
665+
exec.close_session(sids[i]);
666+
}
667+
auto sid = exec.open_session();
668+
ASSERT_TRUE(sid.has_value());
669+
sids[i] = *sid;
670+
exec.set_sampling(sids[i], {}, 0);
671+
fed[i] = 0;
672+
}
673+
};
674+
// Each sequence's last token has no known successor, so it is never fed.
675+
auto remaining = [&](int i) { return seqs[i]->size() - 1 - fed[i]; };
676+
reset();
677+
678+
std::vector<int> widths;
679+
for (int round = 0; round < 6; ++round) {
680+
for (int n = 1; n <= 65; ++n) {
681+
widths.push_back(n);
682+
}
683+
widths.push_back(kWidthsMaxStep);
684+
widths.push_back(kWidthsMaxStep + 1);
685+
}
686+
std::mt19937 rng(1234);
687+
std::shuffle(widths.begin(), widths.end(), rng);
688+
689+
size_t checked = 0;
690+
for (const int n : widths) {
691+
SCOPED_TRACE("step of " + std::to_string(n));
692+
size_t total = 0;
693+
for (int i = 0; i < sessions; ++i) {
694+
total += remaining(i);
695+
}
696+
if (total < static_cast<size_t>(n)) {
697+
reset();
698+
}
699+
// A random subset and split of the sequences that sums to n, each
700+
// sequence at most once.
701+
std::vector<int> order(sessions);
702+
std::iota(order.begin(), order.end(), 0);
703+
std::shuffle(order.begin(), order.end(), rng);
704+
size_t after = 0;
705+
for (const int i : order) {
706+
after += remaining(i);
707+
}
708+
batching::BatchInput batch;
709+
std::vector<int> owner;
710+
size_t left = static_cast<size_t>(n);
711+
for (const int i : order) {
712+
after -= remaining(i);
713+
if (left == 0 || remaining(i) == 0) {
714+
continue;
715+
}
716+
const size_t most = std::min(remaining(i), left);
717+
// Enough that the sequences after this one can cover the rest.
718+
const size_t least = left > after ? left - after : 0;
719+
if (most < least || (least == 0 && rng() % 3 == 0)) {
720+
continue;
721+
}
722+
const size_t take = std::uniform_int_distribution<size_t>(
723+
std::max<size_t>(least, 1), most)(rng);
724+
const size_t end = fed[i] + take;
725+
// Ask for the chunk's last row whenever its successor is known, but
726+
// not always, so some steps want none.
727+
const bool known = end >= cases_[i].prompt.size();
728+
batch.inputs.push_back(batching::Input{
729+
sids[i], known && rng() % 4 != 0, fed[i], take, seqs[i], 0});
730+
owner.push_back(i);
731+
left -= take;
732+
}
733+
ASSERT_EQ(left, 0u);
734+
ASSERT_EQ(batch.size(), static_cast<size_t>(n));
735+
736+
const auto before = exec.method_calls();
737+
batching::BatchOutput out;
738+
ASSERT_TRUE(exec.execute(batch, out));
739+
const auto calls = exec.method_calls();
740+
741+
// The method(s) the width runs through.
742+
std::map<std::string, std::uint64_t> ran;
743+
for (size_t m = 0; m < calls.size(); ++m) {
744+
if (calls[m].calls != before[m].calls) {
745+
ran[calls[m].name] = calls[m].calls - before[m].calls;
746+
}
747+
}
748+
std::map<std::string, std::uint64_t> want;
749+
if (n <= 64) {
750+
want["forward_" + std::to_string(next_pow2(n))] = 1;
751+
} else if (n <= kWidthsMaxStep) {
752+
want["forward_others"] = 1;
753+
} else {
754+
want["forward_others"] = 1;
755+
want["forward_" + std::to_string(next_pow2(n - kWidthsMaxStep))] = 1;
756+
}
757+
EXPECT_EQ(ran, want);
758+
759+
ASSERT_EQ(out.outputs.size(), batch.inputs.size());
760+
for (size_t j = 0; j < batch.inputs.size(); ++j) {
761+
const int i = owner[j];
762+
const auto& input = batch.inputs[j];
763+
const size_t end = input.offset + input.size;
764+
if (input.produce_output) {
765+
ASSERT_TRUE(out.outputs[j].has_value()) << "sequence " << i;
766+
EXPECT_EQ(out.outputs[j]->sid, sids[i]);
767+
ASSERT_EQ(out.outputs[j]->tokens.size(), 1u);
768+
EXPECT_EQ(out.outputs[j]->tokens[0], (*seqs[i])[end])
769+
<< "sequence " << i << " after " << end << " tokens";
770+
++checked;
771+
} else {
772+
EXPECT_FALSE(out.outputs[j].has_value()) << "sequence " << i;
773+
}
774+
fed[i] = end;
775+
}
776+
}
777+
EXPECT_GT(checked, widths.size());
778+
for (const batching::SessionId sid : sids) {
779+
exec.close_session(sid);
780+
}
781+
}
782+
}
783+
784+
INSTANTIATE_TEST_SUITE_P(
785+
Toy,
786+
CudaExecutorWidthsTest,
787+
::testing::Values("mg_widths"));

0 commit comments

Comments
 (0)