From 2548fb590720a27e48e42dff0a137b047474a87a Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Fri, 2 Oct 2026 12:06:52 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .github/workflows/cuda.yml | 12 +- backends/cuda/batching/CMakeLists.txt | 35 +- backends/cuda/batching/__init__.py | 0 backends/cuda/batching/cuda_executor.cpp | 622 ++++++++++++++++++ backends/cuda/batching/cuda_executor.h | 158 +++++ backends/cuda/batching/targets.bzl | 56 ++ backends/cuda/batching/test/__init__.py | 0 .../cuda/batching/test/export_toy_decoder.py | 253 +++++++ .../cuda/batching/test/test_cuda_executor.cpp | 250 +++++++ 9 files changed, 1381 insertions(+), 5 deletions(-) create mode 100644 backends/cuda/batching/__init__.py create mode 100644 backends/cuda/batching/cuda_executor.cpp create mode 100644 backends/cuda/batching/cuda_executor.h create mode 100644 backends/cuda/batching/test/__init__.py create mode 100644 backends/cuda/batching/test/export_toy_decoder.py create mode 100644 backends/cuda/batching/test/test_cuda_executor.cpp diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 65d755e5280..1472146dbee 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -515,12 +515,22 @@ jobs: -v -o "addopts=" cmake --preset llm-release-cuda -DEXECUTORCH_BUILD_TESTS=ON - cmake --build cmake-out --target test_cuda_allocator test_cuda_mutable_state test_cuda_weight_cache test_cuda_kv_cache test_cuda_guard test_cuda_stream_guard test_step_plan -j$(nproc) + cmake --build cmake-out --target test_cuda_allocator test_cuda_mutable_state test_cuda_weight_cache test_cuda_kv_cache test_cuda_guard test_cuda_stream_guard test_step_plan test_cuda_executor -j$(nproc) ctest --test-dir cmake-out -R test_cuda_allocator --output-on-failure -V ctest --test-dir cmake-out -R test_cuda_mutable_state --output-on-failure -V ctest --test-dir cmake-out -R test_cuda_weight_cache --output-on-failure -V ctest --test-dir cmake-out -R test_cuda_kv_cache --output-on-failure -V ctest --test-dir cmake-out -R test_step_plan --output-on-failure -V + # CudaExecutor end to end on a toy decoder, against its eager greedy + # output. The test skips without the model, so fail if it skipped. + export ET_CUDA_BATCHING_TOY_DIR="${RUNNER_TEMP:-/tmp}/cuda_batching_toy" + python -m executorch.backends.cuda.batching.test.export_toy_decoder --output-dir "${ET_CUDA_BATCHING_TOY_DIR}" + ctest --test-dir cmake-out -R test_cuda_executor --output-on-failure -V > /tmp/test_cuda_executor.log 2>&1 || { cat /tmp/test_cuda_executor.log; exit 1; } + cat /tmp/test_cuda_executor.log + if grep -q "SKIPPED" /tmp/test_cuda_executor.log; then + echo "::error::test_cuda_executor skipped" + exit 1 + fi ctest --test-dir cmake-out -R test_cuda_guard --output-on-failure -V ctest --test-dir cmake-out -R test_cuda_stream_guard --output-on-failure -V diff --git a/backends/cuda/batching/CMakeLists.txt b/backends/cuda/batching/CMakeLists.txt index a46dd2b8e2d..c665ab79485 100644 --- a/backends/cuda/batching/CMakeLists.txt +++ b/backends/cuda/batching/CMakeLists.txt @@ -4,15 +4,33 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -# Step slicing for the CUDA batching executor: cuts a batch of packed tokens -# into forwards the exported decode/prefill methods can run. Added from the -# root after extension/llm/batching. +# The batching executor for programs the CUDA backend compiled: the CUDA +# counterpart of extension_llm_batching_module. Added from the root after +# extension/llm/batching, whose executor seam it implements. if(NOT EXECUTORCH_ROOT) set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..) endif() -install(FILES step_plan.h +add_library(cuda_batching cuda_executor.cpp) +target_include_directories(cuda_batching PUBLIC ${_common_include_directories}) +target_compile_options(cuda_batching PUBLIC ${_common_compile_options}) +target_compile_definitions(cuda_batching PRIVATE CUDA_AVAILABLE=1) +target_link_libraries( + cuda_batching + PUBLIC aoti_cuda_backend extension_llm_batching extension_llm_cache + extension_module extension_tensor + PRIVATE extension_llm_sampler +) + +install( + TARGETS cuda_batching + EXPORT ExecuTorchTargets + DESTINATION ${CMAKE_INSTALL_LIBDIR} + INCLUDES + DESTINATION ${_common_include_directories} +) +install(FILES cuda_executor.h step_plan.h DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/executorch/backends/cuda/batching ) @@ -20,4 +38,13 @@ if(BUILD_TESTING) include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake) et_cxx_test(test_step_plan SOURCES test/test_step_plan.cpp) + + # Skips unless ET_CUDA_BATCHING_TOY_DIR names the output of + # test/export_toy_decoder.py. The program's host/device copies are portable + # et_copy kernels, registered by the CPU ops library. + et_cxx_test( + test_cuda_executor SOURCES test/test_cuda_executor.cpp EXTRA_LIBS + cuda_batching optimized_native_cpu_ops_lib + ) + executorch_target_link_options_shared_lib(optimized_native_cpu_ops_lib) endif() diff --git a/backends/cuda/batching/__init__.py b/backends/cuda/batching/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/backends/cuda/batching/cuda_executor.cpp b/backends/cuda/batching/cuda_executor.cpp new file mode 100644 index 00000000000..b6c815fa490 --- /dev/null +++ b/backends/cuda/batching/cuda_executor.cpp @@ -0,0 +1,622 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace executorch::backends::cuda::batching { + +using ::executorch::extension::make_tensor_ptr; +using ::executorch::extension::Module; +using ::executorch::extension::llm::LogitsToKeepMode; +using ::executorch::extension::llm::Sampler; +using ::executorch::runtime::Error; +using ::executorch::runtime::MethodMeta; +using ::executorch::runtime::Result; +using llm_batching::BatchInput; +using llm_batching::BatchOutput; +using llm_batching::Input; +using llm_batching::Output; +using llm_batching::Position; +using llm_batching::SamplingParams; +using llm_batching::SessionId; +using llm_batching::Token; + +namespace metadata = ::executorch::extension::llm; + +namespace { + +// A prefill forward carries at least this many tokens; one token is decode's. +constexpr int kMinPrefillTokens = 2; + +struct SequenceGuard { + llm_cache::BatchControl& control; + std::int32_t seq_id; + bool owned = true; + + ~SequenceGuard() { + if (owned) { + control.seq_rm(seq_id); + } + } +}; + +std::uint64_t nondeterministic_seed() { + std::random_device device; + return device(); +} + +bool is_supported_logits_type(::executorch::aten::ScalarType type) { + using ScalarType = ::executorch::aten::ScalarType; + return type == ScalarType::Float || type == ScalarType::Half || + type == ScalarType::BFloat16; +} + +// The widest step `method` takes: the extent its token input was exported +// with, which for a dynamic dimension is its upper bound. +Result step_width(const MethodMeta& meta, const char* method) { + ET_CHECK_OR_RETURN_ERROR( + meta.num_inputs() == 3, + InvalidProgram, + "CudaExecutor: %s must take tokens, positions and a logits selector", + method); + const auto tokens = meta.input_tensor_meta(0); + const auto positions = meta.input_tensor_meta(1); + const auto selector = meta.input_tensor_meta(2); + ET_CHECK_OR_RETURN_ERROR( + tokens.ok() && positions.ok() && selector.ok(), + InvalidProgram, + "CudaExecutor: %s inputs must be tensors", + method); + const auto token_sizes = tokens->sizes(); + const auto position_sizes = positions->sizes(); + ET_CHECK_OR_RETURN_ERROR( + tokens->scalar_type() == ::executorch::aten::ScalarType::Long && + token_sizes.size() == 2 && token_sizes[0] == 1 && + token_sizes[1] > 0 && + positions->scalar_type() == ::executorch::aten::ScalarType::Long && + position_sizes.size() == 1 && position_sizes[0] == token_sizes[1] && + selector->scalar_type() == ::executorch::aten::ScalarType::Long && + selector->sizes().size() == 1, + InvalidProgram, + "CudaExecutor: %s must take Long[1, T] tokens, Long[T] positions and a " + "rank-one Long selector", + method); + return static_cast(token_sizes[1]); +} + +// The logits width `method` produces. +Result logits_width(const MethodMeta& meta, const char* method) { + ET_CHECK_OR_RETURN_ERROR( + meta.num_outputs() > 0, + InvalidProgram, + "CudaExecutor: %s publishes no outputs", + method); + const auto logits = meta.output_tensor_meta(0); + ET_CHECK_OK_OR_RETURN_ERROR(logits.error()); + const auto sizes = logits->sizes(); + ET_CHECK_OR_RETURN_ERROR( + sizes.size() >= 2 && sizes[sizes.size() - 1] > 0 && + is_supported_logits_type(logits->scalar_type()), + InvalidProgram, + "CudaExecutor: %s logits must have shape [..., vocab] and a float dtype", + method); + return sizes[sizes.size() - 1]; +} + +// The one backend `method` delegates to. +Result backend_of(const MethodMeta& meta, const char* method) { + std::string backend_id; + for (std::size_t i = 0; i < meta.num_backends(); ++i) { + const auto name = meta.get_backend_name(i); + ET_CHECK_OK_OR_RETURN_ERROR(name.error()); + ET_CHECK_OR_RETURN_ERROR( + backend_id.empty() || backend_id == name.get(), + InvalidProgram, + "CudaExecutor: %s spans more than one backend", + method); + backend_id = name.get(); + } + ET_CHECK_OR_RETURN_ERROR( + backend_id == kCudaBackendId, + InvalidProgram, + "CudaExecutor: %s is not delegated to %s", + method, + kCudaBackendId); + return backend_id; +} + +Error set_backend_options(const CudaExecutorOptions& options) { + ::executorch::runtime::BackendOptions<2> backend_options; + ET_CHECK_OK_OR_RETURN_ERROR(backend_options.set_option( + "weight_sharing_across_methods", options.weight_sharing_across_methods)); + ET_CHECK_OK_OR_RETURN_ERROR(backend_options.set_option( + "enable_cuda_graph_for_method", + options.cuda_graph_for_decode ? kDecodeMethod : "")); + return ::executorch::runtime::set_option( + kCudaBackendId, backend_options.view()); +} + +} // namespace + +CudaExecutor::CudaExecutor( + std::unique_ptr module, + std::shared_ptr cache, + int max_sessions, + int max_session_tokens, + std::string backend_id, + std::int32_t vocab_size, + int max_step_tokens) + : install_guard_(cache), + module_(std::move(module)), + ctl_(cache->as()), + kv_(cache->as()), + max_sessions_(max_sessions), + max_session_tokens_(max_session_tokens), + backend_id_(std::move(backend_id)), + vocab_size_(vocab_size), + max_step_tokens_(max_step_tokens) {} + +CudaExecutor::~CudaExecutor() = default; + +Result> CudaExecutor::create( + std::unique_ptr module, + int max_sessions, + int max_session_tokens, + int kv_dtype, + int initial_capacity, + CudaExecutorOptions options) { + ET_CHECK_OR_RETURN_ERROR( + module != nullptr, InvalidArgument, "CudaExecutor: no program"); + ET_CHECK_OR_RETURN_ERROR( + max_sessions > 0 && max_session_tokens > 0 && + max_sessions <= std::numeric_limits::max() / max_session_tokens, + InvalidArgument, + "CudaExecutor: invalid session limits"); + ET_CHECK_OK_OR_RETURN_ERROR(module->load()); + + ET_ASSIGN_OR_RETURN( + max_context_length, + metadata::read_max_context_length(*module)); + ET_CHECK_OR_RETURN_ERROR( + max_session_tokens <= max_context_length, + InvalidArgument, + "CudaExecutor: max session tokens %d exceeds model context length %" PRId64, + max_session_tokens, + max_context_length); + ET_ASSIGN_OR_RETURN( + logits_mode, + metadata::read_logits_to_keep_mode(*module)); + // decode's selector is static at one row, so only a program whose logits + // are selected per input fits. + ET_CHECK_OR_RETURN_ERROR( + logits_mode == LogitsToKeepMode::Selected, + NotSupported, + "CudaExecutor: the program must select its logits rows"); + ET_ASSIGN_OR_RETURN( + geometry, metadata::read_cache_geometry(*module)); + ET_ASSIGN_OR_RETURN( + max_cells, + metadata::detail::read_required_positive_int(*module, kMaxCellsMethod)); + + ET_ASSIGN_OR_RETURN( + decode_meta, module->method_meta(kDecodeMethod)); + ET_ASSIGN_OR_RETURN( + prefill_meta, module->method_meta(kPrefillMethod)); + ET_ASSIGN_OR_RETURN( + decode_width, step_width(decode_meta, kDecodeMethod)); + ET_ASSIGN_OR_RETURN( + max_step_tokens, step_width(prefill_meta, kPrefillMethod)); + ET_CHECK_OR_RETURN_ERROR( + decode_width == 1 && max_step_tokens >= kMinPrefillTokens && + max_step_tokens <= max_cells, + InvalidProgram, + "CudaExecutor: decode must take one token and prefill [%d, max_cells]; " + "got %d and %d", + kMinPrefillTokens, + decode_width, + max_step_tokens); + ET_ASSIGN_OR_RETURN( + decode_vocab, logits_width(decode_meta, kDecodeMethod)); + ET_ASSIGN_OR_RETURN( + prefill_vocab, + logits_width(prefill_meta, kPrefillMethod)); + ET_CHECK_OR_RETURN_ERROR( + decode_vocab == prefill_vocab, + InvalidProgram, + "CudaExecutor: decode and prefill disagree on the vocabulary"); + ET_ASSIGN_OR_RETURN( + published_vocab, metadata::read_vocab_size(*module)); + ET_ASSIGN_OR_RETURN( + vocab_size, + metadata::check_vocab_size(published_vocab, decode_vocab)); + ET_ASSIGN_OR_RETURN( + backend_id, backend_of(decode_meta, kDecodeMethod)); + ET_ASSIGN_OR_RETURN( + prefill_backend, + backend_of(prefill_meta, kPrefillMethod)); + (void)prefill_backend; + + // Every resident session may fill its budget at once; the pool must hold + // them all, or a step could find no free cell mid-generation. + ET_CHECK_OR_RETURN_ERROR( + static_cast(max_sessions) * max_session_tokens <= max_cells, + InvalidArgument, + "CudaExecutor: %d sessions of %d tokens exceed the program's %" PRId64 + " cells", + max_sessions, + max_session_tokens, + max_cells); + + // The pool is the program's: its size and widest step fix the shapes the + // program declared for the step buffers. + llm_cache::CacheConfig cfg{static_cast(max_cells), kv_dtype}; + cfg.max_write = max_step_tokens; + if (initial_capacity >= 0) { + cfg.initial_capacity = std::min(initial_capacity, cfg.capacity); + } + ET_CHECK_OR_RETURN_ERROR( + llm_cache::valid(geometry, cfg), + InvalidProgram, + "CudaExecutor: the program's layout is unusable"); + auto built = llm_cache::CacheFactory::global().build( + backend_id, llm_cache::kind::kBatchedCell, geometry, cfg); + ET_CHECK_OK_OR_RETURN_ERROR(built.error()); + std::shared_ptr cache = built.get(); + auto* const ctl = cache->as(); + ET_CHECK_OR_RETURN_ERROR( + ctl != nullptr && cache->as() != nullptr, + InvalidType, + "CudaExecutor: the cache is not a batched CUDA cache"); + const std::optional seq_limit = ctl->max_seqs(); + ET_CHECK_OR_RETURN_ERROR( + !seq_limit || max_sessions <= *seq_limit, + InvalidArgument, + "CudaExecutor: %d resident sessions requested, but the layout holds %d", + max_sessions, + seq_limit.value_or(0)); + + ET_CHECK_OK_OR_RETURN_ERROR(set_backend_options(options)); + return std::unique_ptr(new CudaExecutor( + std::move(module), + std::move(cache), + max_sessions, + max_session_tokens, + std::move(backend_id), + vocab_size, + max_step_tokens)); +} + +bool CudaExecutor::initialize() { + // The delegate resolves the cache from this key while each method loads. + ::executorch::runtime::BackendOptions<1> options; + ::executorch::runtime::LoadBackendOptionsMap options_map; + if (install_guard_.set_option(options) != Error::Ok || + options_map.set_options(backend_id_.c_str(), options.view()) != + Error::Ok) { + ET_LOG(Error, "CudaExecutor: could not name the cache to the backend"); + return false; + } + for (const char* method : {kDecodeMethod, kPrefillMethod}) { + if (module_->load_method( + method, + /*planned_memory=*/nullptr, + /*event_tracer=*/nullptr, + &options_map) != Error::Ok) { + ET_LOG(Error, "CudaExecutor: could not load %s", method); + return false; + } + } + return true; +} + +// Copied from ModuleExecutor::build_step; see there for the reasoning. +Result CudaExecutor::build_step(const BatchInput& batch) { + Step step; + const std::size_t total = batch.size(); + step.tokens.reserve(total); + step.positions.reserve(total); + step.seq_ids.reserve(total); + step.logit_indices.reserve(batch.inputs.size()); + std::vector> rewinds; + std::unordered_map cursor; + + for (const Input& input : batch.inputs) { + const auto seq_it = sessions_.find(input.sid); + if (seq_it == sessions_.end()) { + ET_LOG(Error, "build_step: session %" PRId64 " is not open", input.sid); + return Error::InvalidArgument; + } + const std::int32_t seq_id = seq_it->second.seq_id; + if (input.size == 0 || !input.tokens || + input.offset > input.tokens->size() || + input.size > input.tokens->size() - input.offset) { + ET_LOG( + Error, + "build_step: session %" PRId64 " gave a slice its tokens do not hold", + input.sid); + return Error::InvalidArgument; + } + + const std::int64_t start = static_cast(input.position) + + static_cast(input.offset); + const auto [cursor_it, first_for_seq] = + cursor.try_emplace(seq_id, ctl_->pos(seq_id)); + int& at = cursor_it->second; + if (start > at) { + ET_LOG( + Error, + "build_step: session %" PRId64 " starts at %" PRId64 + " over a sequence holding %d", + input.sid, + start, + at); + return Error::InvalidArgument; + } + if (start < at) { + if (!first_for_seq) { + ET_LOG( + Error, + "build_step: session %" PRId64 " overlaps its earlier input", + input.sid); + return Error::InvalidArgument; + } + if (start == 0) { + ET_LOG( + Error, + "build_step: session %" PRId64 " reopens from the start", + input.sid); + return Error::InvalidArgument; + } + rewinds.emplace_back(seq_id, static_cast(start)); + at = static_cast(start); + } + + const std::int64_t end = start + static_cast(input.size); + if (end > max_session_tokens_) { + ET_LOG( + Error, + "build_step: session %" PRId64 " reaches %" PRId64 " of %d cells", + input.sid, + end, + max_session_tokens_); + return Error::OutOfResources; + } + + const Token* slice = input.tokens->data() + input.offset; + for (std::size_t k = 0; k < input.size; ++k) { + step.tokens.push_back(static_cast(slice[k])); + step.positions.push_back(start + static_cast(k)); + } + step.seq_ids.insert(step.seq_ids.end(), input.size, seq_id); + at = static_cast(end); + step.logit_indices.push_back( + input.produce_output ? static_cast(step.tokens.size()) - 1 : -1); + } + + for (const auto& [seq_id, from] : rewinds) { + if (!ctl_->rewind(seq_id, from)) { + ET_LOG(Error, "build_step: sequence %d would not truncate", seq_id); + return Error::Internal; + } + } + return step; +} + +std::optional CudaExecutor::open_session() { + if (sessions_.size() >= static_cast(max_sessions_) || + next_session_ == 0) { + return std::nullopt; + } +#if ET_HAS_EXCEPTIONS + try { +#endif + const std::optional seq_id = ctl_->seq_new(); + return seq_id ? publish_session(*seq_id, 0) : std::nullopt; +#if ET_HAS_EXCEPTIONS + } catch (const std::bad_alloc&) { + return std::nullopt; + } +#endif +} + +std::optional CudaExecutor::publish_session( + std::int32_t seq_id, + Position position) { + SequenceGuard guard{*ctl_, seq_id}; + if (ctl_->pos(seq_id) != position) { + return std::nullopt; + } + const SessionId session = next_session_; + if (!sessions_.emplace(session, SessionState{seq_id, nullptr}).second) { + return std::nullopt; + } + guard.owned = false; + next_session_ = + session == std::numeric_limits::max() ? 0 : session + 1; + return session; +} + +void CudaExecutor::close_session(SessionId session) { + const auto it = sessions_.find(session); + if (it == sessions_.end()) { + return; + } + ctl_->seq_rm(it->second.seq_id); + sessions_.erase(it); +} + +std::optional CudaExecutor::clone(SessionId source, Position upto) { + const auto it = sessions_.find(source); + if (it == sessions_.end() || upto < 0 || upto > max_session_tokens_ || + sessions_.size() >= static_cast(max_sessions_) || + next_session_ == 0) { + return std::nullopt; + } +#if ET_HAS_EXCEPTIONS + try { +#endif + if (upto > ctl_->pos(it->second.seq_id)) { + return std::nullopt; + } + const auto seq_id = ctl_->seq_clone(it->second.seq_id, upto); + return seq_id ? publish_session(*seq_id, upto) : std::nullopt; +#if ET_HAS_EXCEPTIONS + } catch (const std::bad_alloc&) { + return std::nullopt; + } +#endif +} + +void CudaExecutor::set_sampling( + SessionId session, + const SamplingParams& params, + std::optional seed) { + const auto it = sessions_.find(session); + if (it == sessions_.end()) { + return; + } + it->second.sampler = std::make_unique( + vocab_size_, + params.temperature, + params.top_p, + seed.value_or(nondeterministic_seed())); + it->second.sampler->set_topk(params.top_k); +} + +bool CudaExecutor::execute(const BatchInput& batch, BatchOutput& out) { + out.outputs.clear(); + out.outputs.resize(batch.inputs.size()); + + const Result step = build_step(batch); + if (!step.ok()) { + return false; + } + + // In order, so a slice attends the cells its predecessors wrote. Each + // input's logits row falls in exactly one slice. + const int total = static_cast(step->tokens.size()); + for (const StepSlice& slice : + plan_slices(total, max_step_tokens_, kMinPrefillTokens)) { + const int off = slice.offset; + const int n = slice.length; + const char* method = + slice.method == StepMethod::Decode ? kDecodeMethod : kPrefillMethod; + // Placement checks the forward's token count against the declaration, so + // each slice declares its own. + if (!ctl_->declare_step(std::vector( + step->seq_ids.begin() + off, step->seq_ids.begin() + off + n))) { + ET_LOG(Error, "CudaExecutor: the cache refused a slice of %d", n); + return false; + } + auto tokens = make_tensor_ptr( + {1, n}, + std::vector( + step->tokens.begin() + off, step->tokens.begin() + off + n)); + auto positions = make_tensor_ptr( + {n}, + std::vector( + step->positions.begin() + off, step->positions.begin() + off + n)); + std::vector selector_values; + std::vector selected_inputs; + for (std::size_t i = 0; i < step->logit_indices.size(); ++i) { + const int row = step->logit_indices[i]; + if (row >= off && row < off + n) { + selector_values.push_back(row - off); + selected_inputs.push_back(i); + } + } + if (selector_values.empty()) { + // The forward still has to produce a row; nothing reads it. + selector_values.push_back(n - 1); + } + const int rows = static_cast(selector_values.size()); + auto selector = make_tensor_ptr({rows}, std::move(selector_values)); + + auto result = module_->execute(method, {tokens, positions, selector}); + if (!result.ok()) { + ET_LOG( + Error, + "CudaExecutor: %s failed with 0x%x", + method, + static_cast(result.error())); + return false; + } + if (result->empty() || !result->at(0).isTensor()) { + ET_LOG(Error, "CudaExecutor: %s returned no logits", method); + return false; + } + auto logits = result->at(0).toTensor(); + if (logits.dim() < 2 || logits.size(logits.dim() - 1) != vocab_size_ || + logits.numel() != static_cast(rows) * vocab_size_) { + ET_LOG( + Error, + "CudaExecutor: %s returned logits that are not [%d, %d]", + method, + rows, + vocab_size_); + return false; + } + for (std::size_t row = 0; row < selected_inputs.size(); ++row) { + const std::size_t input_index = selected_inputs[row]; + const SessionId session = batch.inputs[input_index].sid; + const std::optional token = + sample_row(logits, static_cast(row), session); + if (!token) { + return false; + } + out.outputs[input_index] = Output{session, {*token}}; + } + } + return true; +} + +std::optional CudaExecutor::sample_row( + ::executorch::aten::Tensor& logits, + int row, + SessionId session) { + const auto it = sessions_.find(session); + if (it == sessions_.end() || it->second.sampler == nullptr) { + ET_LOG( + Error, "CudaExecutor: session %" PRId64 " has no sampling policy", session); + return std::nullopt; + } + auto one_row = make_tensor_ptr( + {vocab_size_}, + static_cast(logits.mutable_data_ptr()) + + static_cast(row) * vocab_size_ * + ::executorch::runtime::elementSize(logits.scalar_type()), + logits.scalar_type()); + return static_cast( + ::executorch::extension::llm::sample_from_logits( + *one_row, *it->second.sampler)); +} + +OffGraphKVMetrics CudaExecutor::kv_metrics() const { + return kv_->metrics(); +} + +} // namespace executorch::backends::cuda::batching diff --git a/backends/cuda/batching/cuda_executor.h b/backends/cuda/batching/cuda_executor.h new file mode 100644 index 00000000000..d9a7fd71074 --- /dev/null +++ b/backends/cuda/batching/cuda_executor.h @@ -0,0 +1,158 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +// A batching Executor for a program the CUDA backend compiled into static +// graphs. The CUDA counterpart of ModuleExecutor: a session is one sequence of +// a cell-layout off-graph KV cache, and a batch is one or more forwards +// carrying every input's tokens on a single axis. Unlike ModuleExecutor it +// runs two methods, since a static graph serves one shape: a one-token +// `decode` a CUDA graph is captured for, and a `prefill` dynamic from two +// tokens up. + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include // ET_EXPERIMENTAL + +namespace executorch { +namespace extension { +namespace llm { +class Sampler; +} // namespace llm +} // namespace extension +} // namespace executorch + +namespace executorch::backends::cuda::batching { + +namespace llm_batching = ::executorch::extension::llm::batching; +namespace llm_cache = ::executorch::extension::llm::cache; + +// The methods and metadata the program must export. +inline constexpr char kDecodeMethod[] = "decode"; +inline constexpr char kPrefillMethod[] = "prefill"; +// The cell layout's pool size, which fixes the shape of its step buffers. +inline constexpr char kMaxCellsMethod[] = "get_offgraph_kv_max_cells"; + +// Process-wide CUDA backend options create() sets before the methods load. +struct CudaExecutorOptions { + // decode and prefill are compiled separately over the same weights; share + // them rather than loading a copy per method. + bool weight_sharing_across_methods = true; + // Capture decode into a CUDA graph. prefill stays eager: its width varies. + bool cuda_graph_for_decode = true; +}; + +class ET_EXPERIMENTAL CudaExecutor : public llm_batching::Executor { + public: + ~CudaExecutor() override; + + // Builds the cell cache from the layout the program publishes and sets the + // CUDA backend options. The methods load in initialize(); the program must + // be loaded and its methods must not be, since the delegate resolves the + // cache while they load. `options` is applied process-wide, as backend + // options are. + // + // `max_sessions` counts every resident session, clones included, and each + // reserves `max_session_tokens` cells; together they must fit the program's + // max_cells. `kv_dtype` is the ET ScalarType K/V is stored in; a negative + // `initial_capacity` starts the pools at the cache's default. + static ::executorch::runtime::Result> create( + std::unique_ptr<::executorch::extension::Module> module, + int max_sessions, + int max_session_tokens, + int kv_dtype, + int initial_capacity = -1, + CudaExecutorOptions options = {}); + + // prefill's widest step. A wider batch runs as several forwards. + std::size_t preferred_batch_tokens() const override { + return static_cast(max_step_tokens_); + } + + // Loads both methods here, so the delegate that resolves the cache binds on + // the thread that runs it. + bool initialize() override; + + std::optional open_session() override; + void close_session(llm_batching::SessionId session) override; + std::optional clone( + llm_batching::SessionId source, + llm_batching::Position upto) override; + void set_sampling( + llm_batching::SessionId session, + const llm_batching::SamplingParams& params, + std::optional seed) override; + bool execute( + const llm_batching::BatchInput& batch, + llm_batching::BatchOutput& out) override; + + // The cache's pool usage, for memory reporting. + OffGraphKVMetrics kv_metrics() const; + + private: + struct SessionState { + std::int32_t seq_id; + std::unique_ptr<::executorch::extension::llm::Sampler> sampler; + }; + + struct Step { + std::vector tokens; + std::vector positions; + std::vector seq_ids; + std::vector logit_indices; + }; + + CudaExecutor( + std::unique_ptr<::executorch::extension::Module> module, + std::shared_ptr cache, + int max_sessions, + int max_session_tokens, + std::string backend_id, + std::int32_t vocab_size, + int max_step_tokens); + + ::executorch::runtime::Result build_step( + const llm_batching::BatchInput& batch); + std::optional publish_session( + std::int32_t seq_id, + llm_batching::Position position); + std::optional sample_row( + ::executorch::aten::Tensor& logits, + int row, + llm_batching::SessionId session); + + // Ordered so the module dies first, releasing the delegates that resolved + // the cache before the registry entry naming it goes. + llm_cache::InstallGuard install_guard_; + std::unique_ptr<::executorch::extension::Module> module_; + llm_cache::BatchControl* const ctl_; + const CudaKVCache* const kv_; + int max_sessions_; + int max_session_tokens_; + std::string backend_id_; + std::int32_t vocab_size_; + int max_step_tokens_; + + llm_batching::SessionId next_session_ = 1; + std::unordered_map sessions_; +}; + +} // namespace executorch::backends::cuda::batching diff --git a/backends/cuda/batching/targets.bzl b/backends/cuda/batching/targets.bzl index 338cb39b77b..44ef58dad36 100644 --- a/backends/cuda/batching/targets.bzl +++ b/backends/cuda/batching/targets.bzl @@ -13,6 +13,38 @@ def define_common_targets(is_fbcode = False): visibility = ["PUBLIC"], ) + runtime.cxx_library( + name = "cuda_executor", + srcs = [ + "cuda_executor.cpp", + ], + exported_headers = [ + "cuda_executor.h", + ], + preprocessor_flags = [ + "-DCUDA_AVAILABLE=1", + ], + visibility = ["PUBLIC"], + exported_deps = [ + ":step_plan", + "//executorch/backends/cuda/runtime:cuda_backend", + "//executorch/extension/llm/batching:batching", + "//executorch/extension/llm/cache:kv_cache", + "//executorch/extension/module:module", + ], + deps = [ + "//executorch/extension/llm/runner:stats", + "//executorch/extension/llm/sampler:sampler", + "//executorch/extension/tensor:tensor", + "//executorch/runtime/backend:interface", + "//executorch/runtime/core/exec_aten/util:scalar_type_util", + "//executorch/runtime/platform:platform", + ], + external_deps = [ + ("cuda", None, "cuda-lazy"), + ], + ) + cpp_unittest( name = "test_step_plan", srcs = [ @@ -22,3 +54,27 @@ def define_common_targets(is_fbcode = False): ":step_plan", ], ) + + # Writes the toy decoder test_cuda_executor runs. Needs a GPU: AOTI + # compiles and autotunes on it. test_cuda_executor itself builds with CMake + # only (see CMakeLists.txt): it runs the compiled AOTI library, which + # resolves its shims against the host binary. + runtime.python_binary( + name = "export_toy_decoder", + srcs = [ + "test/export_toy_decoder.py", + ], + main_function = "executorch.backends.cuda.batching.test.export_toy_decoder.main", + keep_gpu_sections = True, + deps = [ + "//caffe2:torch", + "//executorch/backends/cuda:cuda_backend", + "//executorch/backends/cuda:cuda_partitioner", + "//executorch/backends/cuda:cuda_passes", + "//executorch/exir:lib", + "//executorch/exir/backend:compile_spec_schema", + "//executorch/exir/passes:lib", + "//executorch/extension/llm/cache:cache", + "//executorch/extension/llm/export:model_metadata", + ], + ) diff --git a/backends/cuda/batching/test/__init__.py b/backends/cuda/batching/test/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/backends/cuda/batching/test/export_toy_decoder.py b/backends/cuda/batching/test/export_toy_decoder.py new file mode 100644 index 00000000000..f925d877b56 --- /dev/null +++ b/backends/cuda/batching/test/export_toy_decoder.py @@ -0,0 +1,253 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Export the toy decoder the CudaExecutor GPU test runs. + +Writes ``model.pte`` and ``aoti_cuda_blob.ptd`` -- a two-layer decoder (one +full-history layer, one sliding-window layer) exported as the ``decode`` and +``prefill`` methods CudaExecutor drives, lowered in the cell layout -- plus +``expected.txt``: one ``prompt;continuation`` line per prompt, the greedy +continuation computed eagerly with the neutral reference cache. + +The residual stream carries each token's embedding at a large scale and the +LM head maps it to a fixed successor, while attention adds a smaller term. The +greedy argmax therefore keeps a wide margin over bf16 noise, so the exported +program must reproduce the eager tokens exactly, yet attention still feeds +every logit. + + python -m executorch.backends.cuda.batching.test.export_toy_decoder --output-dir DIR +""" + +import argparse +import json +import os + +import torch +import torch.nn as nn + +VOCAB = 32 +DIM = 64 +N_HEADS = 4 +N_KV_HEADS = 2 +HEAD_DIM = 16 +WINDOWS = (0, 4) # layer 0 keeps its whole history, layer 1 a window of 4 +MAX_CONTEXT = 64 +MAX_STEP = 8 # prefill's widest step: longer prompts slice +MAX_CELLS = 256 +EMBED_SCALE = 4.0 + +PROMPTS = ([3], [5, 9, 1], [7, 2, 2, 8, 4, 6, 1], list(range(1, 21))) +NEW_TOKENS = 6 + + +class ToyDecoder(nn.Module): + def __init__(self) -> None: + super().__init__() + generator = torch.Generator().manual_seed(0) + self.embed = nn.Embedding(VOCAB, DIM) + self.q = nn.ModuleList() + self.k = nn.ModuleList() + self.v = nn.ModuleList() + self.o = nn.ModuleList() + for _ in WINDOWS: + self.q.append(nn.Linear(DIM, N_HEADS * HEAD_DIM, bias=False)) + self.k.append(nn.Linear(DIM, N_KV_HEADS * HEAD_DIM, bias=False)) + self.v.append(nn.Linear(DIM, N_KV_HEADS * HEAD_DIM, bias=False)) + self.o.append(nn.Linear(N_HEADS * HEAD_DIM, DIM, bias=False)) + self.head = nn.Linear(DIM, VOCAB, bias=False) + with torch.no_grad(): + self.embed.weight.zero_() + self.embed.weight[:, :VOCAB] = EMBED_SCALE * torch.eye(VOCAB) + for linear in (*self.q, *self.k, *self.v): + linear.weight.normal_(0.0, 0.3, generator=generator) + for linear in self.o: + linear.weight.normal_(0.0, 0.05, generator=generator) + successor = torch.randperm(VOCAB, generator=generator) + self.head.weight.zero_() + self.head.weight[successor, torch.arange(VOCAB)] = 1.0 + + def forward(self, tokens, input_pos, logits_to_keep): + x = self.embed(tokens) + length = tokens.shape[1] + position = input_pos.reshape(-1, 1) + for layer in range(len(WINDOWS)): + q = self.q[layer](x).view(1, length, N_HEADS, HEAD_DIM).transpose(1, 2) + k = self.k[layer](x).view(1, length, N_KV_HEADS, HEAD_DIM).transpose(1, 2) + v = self.v[layer](x).view(1, length, N_KV_HEADS, HEAD_DIM).transpose(1, 2) + y = torch.ops.kvcache.update_and_attend( + q, k, v, position, layer, HEAD_DIM**-0.5, q.dtype + ) + x = x + self.o[layer](y.transpose(1, 2).reshape(1, length, -1)) + return self.head(x[:, logits_to_keep, :])[0].float() + + +def _manifest() -> str: + return json.dumps( + { + "version": 1, + "layout": "cell", + "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) + ], + } + ) + + +def _greedy(model: ToyDecoder, prompt) -> list: + from executorch.extension.llm.cache.reference_cache import ( + CacheConfig, + LayerPolicy, + SequenceReferenceCache, + ) + from executorch.extension.llm.cache.update_and_attend import REGISTRY + + cache = SequenceReferenceCache( + CacheConfig( + n_layers=len(WINDOWS), + n_kv_heads=N_KV_HEADS, + head_dim=HEAD_DIM, + capacity=MAX_CONTEXT, + layers=tuple(LayerPolicy.ring(w) if w else LayerPolicy.flat() for w in WINDOWS), + ) + ) + key = "toy-decoder-greedy" + REGISTRY.install(key, cache) + try: + with REGISTRY.active(key), torch.no_grad(): + generated = [] + tokens = list(prompt) + start = 0 + for _ in range(NEW_TOKENS): + step = torch.tensor([tokens[start:]]) + logits = model( + step, + torch.arange(start, len(tokens)), + torch.tensor([step.shape[1] - 1]), + ) + start = len(tokens) + token = int(logits[0].argmax()) + generated.append(token) + tokens.append(token) + return generated + finally: + REGISTRY.uninstall(key) + + +def export(output_dir: str) -> None: + import torch._inductor.config as inductor_config + from executorch.backends.cuda.cuda_backend import CudaBackend + from executorch.backends.cuda.cuda_partitioner import CudaPartitioner + from executorch.backends.cuda.passes.lower_offgraph_kv import ( + OFFGRAPH_KV_COMPILE_SPEC, + OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC, + ) + from executorch.exir import ( + EdgeCompileConfig, + ExecutorchBackendConfig, + to_edge_transform_and_lower, + ) + from executorch.exir.backend.compile_spec_schema import CompileSpec + from executorch.exir.passes import MemoryPlanningPass + from executorch.extension.llm.export.model_metadata import ( + write_cache_geometry, + write_logits_to_keep_mode, + write_max_context_len, + write_max_seq_len, + write_vocab_size, + ) + from torch.export import Dim + + inductor_config.aot_inductor.precompile_headers = False + torch.manual_seed(0) + model = ToyDecoder().eval() + expected = [_greedy(model, prompt) for prompt in PROMPTS] + + model = model.to(dtype=torch.bfloat16) + width = Dim("width", min=2, max=MAX_STEP) + rows = Dim("rows", min=1, max=MAX_STEP) + long = {"dtype": torch.long} + with torch.no_grad(): + programs = { + "decode": torch.export.export( + model, + ( + torch.zeros(1, 1, **long), + torch.zeros(1, **long), + torch.zeros(1, **long), + ), + strict=True, + ), + "prefill": torch.export.export( + model, + ( + torch.zeros(1, MAX_STEP, **long), + torch.arange(MAX_STEP, **long), + torch.zeros(MAX_STEP, **long), + ), + dynamic_shapes=({1: width}, {0: width}, {0: rows}), + strict=True, + ), + } + + def partitioner(name: str) -> CudaPartitioner: + return CudaPartitioner( + [ + CudaBackend.generate_method_name_compile_spec(name), + CompileSpec("low_memory_mode", b"ON"), + CompileSpec(OFFGRAPH_KV_COMPILE_SPEC, _manifest().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_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, + } + program = to_edge_transform_and_lower( + programs, + partitioner={name: [partitioner(name)] for name in programs}, + compile_config=EdgeCompileConfig(_check_ir_validity=False, _skip_dim_order=True), + constant_methods=constant_methods, + ).to_executorch( + config=ExecutorchBackendConfig( + extract_delegate_segments=True, + memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False), + ) + ) + + os.makedirs(output_dir, exist_ok=True) + with open(os.path.join(output_dir, "model.pte"), "wb") as f: + 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): + f.write( + " ".join(map(str, prompt)) + + ";" + + " ".join(map(str, continuation)) + + "\n" + ) + print(f"Wrote {output_dir}/model.pte, aoti_cuda_blob.ptd, expected.txt") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--output-dir", required=True) + export(parser.parse_args().output_dir) + + +if __name__ == "__main__": + main() diff --git a/backends/cuda/batching/test/test_cuda_executor.cpp b/backends/cuda/batching/test/test_cuda_executor.cpp new file mode 100644 index 00000000000..a3650c59841 --- /dev/null +++ b/backends/cuda/batching/test/test_cuda_executor.cpp @@ -0,0 +1,250 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// Drives CudaExecutor through the batching Runner on the toy decoder that +// export_toy_decoder.py writes to $ET_CUDA_BATCHING_TOY_DIR, and checks every +// generation against the eager greedy continuation it recorded. + +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace cb = ::executorch::backends::cuda::batching; +namespace batching = ::executorch::extension::llm::batching; +using ::executorch::extension::Module; +using ::executorch::runtime::Error; + +namespace { + +// Must match export_toy_decoder.py. +constexpr int kMaxStep = 8; +constexpr int kMaxCells = 256; +constexpr int kMaxContext = 64; +constexpr int kNewTokens = 6; +constexpr int kBFloat16 = 15; // ScalarType::BFloat16 + +struct Case { + std::vector prompt; + std::vector expected; +}; + +std::vector parse_tokens(const std::string& text) { + std::vector tokens; + std::istringstream in(text); + batching::Token token; + while (in >> token) { + tokens.push_back(token); + } + return tokens; +} + +class CudaExecutorTest : public ::testing::Test { + protected: + void SetUp() override { + int devices = 0; + if (cudaGetDeviceCount(&devices) != cudaSuccess || devices == 0) { + GTEST_SKIP() << "CUDA device required"; + } + const char* dir = std::getenv("ET_CUDA_BATCHING_TOY_DIR"); + if (dir == nullptr || *dir == '\0') { + GTEST_SKIP() << "ET_CUDA_BATCHING_TOY_DIR is not set; run " + "export_toy_decoder.py first"; + } + dir_ = dir; + std::ifstream in(dir_ + "/expected.txt"); + ASSERT_TRUE(in.is_open()) << dir_ << "/expected.txt"; + std::string line; + while (std::getline(in, line)) { + const auto split = line.find(';'); + ASSERT_NE(split, std::string::npos) << line; + cases_.push_back( + {parse_tokens(line.substr(0, split)), + parse_tokens(line.substr(split + 1))}); + } + ASSERT_FALSE(cases_.empty()); + } + + std::unique_ptr module() const { + return std::make_unique( + dir_ + "/model.pte", + std::vector{dir_ + "/aoti_cuda_blob.ptd"}, + Module::LoadMode::File, + /*event_tracer=*/nullptr, + /*memory_allocator=*/nullptr, + /*temp_allocator=*/nullptr, + /*share_memory_arenas=*/false); + } + + std::unique_ptr executor( + cb::CudaExecutorOptions options = {}, + int max_sessions = 4) const { + auto created = cb::CudaExecutor::create( + module(), + max_sessions, + kMaxCells / max_sessions > kMaxContext ? kMaxContext + : kMaxCells / max_sessions, + kBFloat16, + /*initial_capacity=*/16, + options); + EXPECT_EQ(created.error(), Error::Ok); + return created.ok() ? std::move(created.get()) : nullptr; + } + + // Runs every prompt concurrently through one runner; returns each + // generation's tokens and its finish reason. + struct Generation { + std::vector tokens; + std::optional reason; + }; + static std::vector generate( + batching::Runner& runner, + const std::vector>& prompts) { + std::vector generations(prompts.size()); + std::vector sessions; + std::vector handles; + std::mutex mutex; + for (size_t i = 0; i < prompts.size(); ++i) { + auto session = runner.open_session_async().get(); + EXPECT_TRUE(session.has_value()) << i; + if (!session) { + return generations; + } + batching::GenConfig config; + config.max_new_tokens = kNewTokens; + config.sampling.temperature = 0.0f; + config.seed = 0; + handles.push_back(session->generate_async( + prompts[i], config, [&, i](const batching::GenerationUpdate& update) { + std::lock_guard guard(mutex); + auto& tokens = generations[i].tokens; + tokens.insert( + tokens.end(), update.tokens.begin(), update.tokens.end()); + })); + sessions.push_back(std::move(*session)); + } + for (size_t i = 0; i < handles.size(); ++i) { + handles[i].wait(); + generations[i].reason = handles[i].finish_reason(); + } + return generations; + } + + std::vector> prompts() const { + std::vector> out; + for (const Case& c : cases_) { + out.push_back(c.prompt); + } + return out; + } + + static std::unique_ptr scheduler() { + // Half the forward for decodes, half for prefill chunks, so decodes and + // prefills share forwards and a long prompt runs in several chunks. + return batching::DecodeFirstScheduler::create(kMaxStep, 4, kMaxStep / 2); + } + + std::string dir_; + std::vector cases_; +}; + +} // namespace + +TEST_F(CudaExecutorTest, ConcurrentGenerationsMatchEagerGreedy) { + auto exec = executor(); + ASSERT_NE(exec, nullptr); + EXPECT_EQ(exec->preferred_batch_tokens(), static_cast(kMaxStep)); + batching::Runner runner(*exec, scheduler()); + const auto generations = generate(runner, prompts()); + const auto kv = exec->kv_metrics(); + runner.shutdown(); + + for (size_t i = 0; i < cases_.size(); ++i) { + EXPECT_EQ(generations[i].reason, batching::FinishReason::NewTokenLimit) + << i; + EXPECT_EQ(generations[i].tokens, cases_[i].expected) << "prompt " << i; + } + // The pools grew from 16 rows to hold every sequence's tokens at once. + size_t tokens = 0; + for (const Case& c : cases_) { + tokens += c.prompt.size() + kNewTokens; + } + EXPECT_GE(kv.flat_capacity, kv.logical_length); + EXPECT_GE(kv.growth_count, 1); + EXPECT_GT(kv.allocated_bytes, 0); + EXPECT_LE(kv.logical_length, static_cast(tokens)); + + // Batched: fewer forwards than running the generations one after another. + const auto engine = runner.metrics(); + EXPECT_EQ(engine.steps_failed, 0u); + EXPECT_GT(engine.decode_sessions_total, engine.steps / 2); +} + +TEST_F(CudaExecutorTest, EagerDecodeMatchesTheCapturedGraph) { + cb::CudaExecutorOptions options; + options.cuda_graph_for_decode = false; + auto exec = executor(options); + ASSERT_NE(exec, nullptr); + batching::Runner runner(*exec, scheduler()); + const auto generations = generate(runner, prompts()); + runner.shutdown(); + for (size_t i = 0; i < cases_.size(); ++i) { + EXPECT_EQ(generations[i].tokens, cases_[i].expected) << "prompt " << i; + } +} + +TEST_F(CudaExecutorTest, SamePromptTwiceInOneBatchGeneratesTheSame) { + auto exec = executor(); + ASSERT_NE(exec, nullptr); + batching::Runner runner(*exec, scheduler()); + const auto& c = cases_.back(); + const auto generations = generate(runner, {c.prompt, c.prompt}); + runner.shutdown(); + EXPECT_EQ(generations[0].tokens, c.expected); + EXPECT_EQ(generations[1].tokens, c.expected); +} + +TEST_F(CudaExecutorTest, SessionsReuseCellsAcrossRounds) { + auto exec = executor(); + ASSERT_NE(exec, nullptr); + batching::Runner runner(*exec, scheduler()); + // Sessions close between rounds, so the second round refills freed cells + // and must not read what the first left there. + for (int round = 0; round < 2; ++round) { + const auto generations = generate(runner, prompts()); + for (size_t i = 0; i < cases_.size(); ++i) { + EXPECT_EQ(generations[i].tokens, cases_[i].expected) + << "round " << round << " prompt " << i; + } + } + runner.shutdown(); +} + +TEST_F(CudaExecutorTest, RefusesLimitsThePoolCannotHold) { + // 8 sessions of the full 64-token context need 512 cells; the program has + // 256. + EXPECT_EQ( + cb::CudaExecutor::create(module(), 8, kMaxContext, kBFloat16, 16, {}) + .error(), + Error::InvalidArgument); + // Longer than the model's context. + EXPECT_EQ( + cb::CudaExecutor::create(module(), 1, kMaxContext + 1, kBFloat16, 16, {}) + .error(), + Error::InvalidArgument); +}