diff --git a/backends/cuda/batching/CMakeLists.txt b/backends/cuda/batching/CMakeLists.txt index 32c8f9a9f7f..fdd1e5e50ec 100644 --- a/backends/cuda/batching/CMakeLists.txt +++ b/backends/cuda/batching/CMakeLists.txt @@ -20,6 +20,7 @@ target_link_libraries( cuda_batching PUBLIC aoti_cuda_backend extension_llm_batching extension_llm_batching_util extension_llm_cache extension_module extension_tensor + PRIVATE extension_llm_batching_executor_utils extension_llm_sampler ) install( diff --git a/backends/cuda/batching/cuda_executor.cpp b/backends/cuda/batching/cuda_executor.cpp index a64520b447e..a47197f82aa 100644 --- a/backends/cuda/batching/cuda_executor.cpp +++ b/backends/cuda/batching/cuda_executor.cpp @@ -22,8 +22,10 @@ #include #include +#include #include -#include +#include +#include #include #include #include @@ -209,6 +211,17 @@ Error set_backend_options( kCudaBackendId, backend_options.view()); } +bool is_dense(const ::executorch::aten::Tensor& tensor) { + std::int64_t stride = 1; + for (auto d = tensor.dim() - 1; d >= 0; --d) { + if (tensor.size(d) > 1 && tensor.strides()[d] != stride) { + return false; + } + stride *= tensor.size(d); + } + return true; +} + } // namespace CudaExecutor::CudaExecutor( @@ -516,6 +529,7 @@ bool CudaExecutor::initialize() { return false; } } + initialized_ = true; return true; } @@ -538,86 +552,221 @@ void CudaExecutor::set_sampling( sessions_.set_sampling(session, params, seed); } -bool CudaExecutor::execute(const BatchInput& batch, BatchOutput& out) { - out.outputs.clear(); - out.outputs.resize(batch.inputs.size()); - - const Result step = sessions_.pack(batch); - if (!step.ok()) { +bool CudaExecutor::execute_raw( + const RawTargetStep& step, + const std::function& consume) { + if (!initialized_ || raw_executing_ || !consume || + step.tokens.size() > + static_cast(std::numeric_limits::max()) || + step.positions.size() != step.tokens.size() || + step.seq_ids.size() != step.tokens.size()) { 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, methods_)) { + const int total = static_cast(step.tokens.size()); + for (int i = 0; i < total; ++i) { + if (step.tokens[i] < 0 || step.tokens[i] >= vocab_size_ || + step.positions[i] < 0 || step.positions[i] >= max_session_tokens_) { + return false; + } + } + for (const auto index : step.logit_indices) { + if (index < -1 || index >= total) { + return false; + } + } + struct Reset { + explicit Reset(bool& target) : flag(target) { + flag = true; + } + Reset(const Reset&) = delete; + Reset& operator=(const Reset&) = delete; + Reset(Reset&&) = delete; + Reset& operator=(Reset&&) = delete; + ~Reset() { + flag = false; + } + bool& flag; + } reset(raw_executing_); + for (const auto& slice : plan_slices(total, methods_)) { const int off = slice.offset; const int n = slice.length; - const ForwardMethod& method = methods_[slice.method]; - // Only the real tokens are declared; the cache sends a static method's - // padding rows to its scratch row. + const auto& method = methods_[slice.method]; + std::vector selected; + std::vector indices; + for (std::size_t i = 0; i < step.logit_indices.size(); ++i) { + const auto index = step.logit_indices[i]; + if (index >= off && index < off + n) { + selected.push_back(index - off); + indices.push_back(i); + } + } + // Static selector widths are fixed; dynamic selectors have an independent + // axis but are still bounded by the exported maximum step. + const auto row_limit = static_cast( + method.is_static ? slice.width : method.max_tokens); + if (selected.size() > row_limit) { + ET_LOG( + Error, + "CudaExecutor: too many raw selectors for %s", + method.name.c_str()); + return false; + } + const std::size_t rows = method.is_static + ? slice.width + : std::max(selected.size(), min_selected_rows_); + selected.resize(rows, selected.empty() ? 0 : selected.back()); 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); + step.seq_ids.begin() + off, step.seq_ids.begin() + off + n))) { return false; } - // Padding tokens are token 0 at position 0: valid inputs whose rows - // nothing reads. std::vector token_values( - step->tokens.begin() + off, step->tokens.begin() + off + n); + step.tokens.begin() + off, step.tokens.begin() + off + n); std::vector position_values( - step->positions.begin() + off, step->positions.begin() + off + n); + step.positions.begin() + off, step.positions.begin() + off + n); token_values.resize(slice.width, 0); position_values.resize(slice.width, 0); auto tokens = make_tensor_ptr({1, slice.width}, std::move(token_values)); auto positions = make_tensor_ptr({slice.width}, std::move(position_values)); - - // A static method selects exactly its width in rows, and the dynamic one - // at least its exported minimum. Extra rows repeat the last real one (or - // row 0 when the slice finishes no input); nothing reads them. - auto selected = llm_batching::util::select_rows(*step, off, n); - const std::size_t rows_needed = method.is_static - ? static_cast(slice.width) - : std::max( - selected.selector.size(), - static_cast(min_selected_rows_)); - const std::int64_t fill = - selected.selector.empty() ? 0 : selected.selector.back(); - selected.selector.resize(rows_needed, fill); - const int rows = static_cast(selected.selector.size()); - auto selector = make_tensor_ptr({rows}, std::move(selected.selector)); - + auto selector = make_tensor_ptr( + {static_cast<::executorch::aten::SizesType>(rows)}, + std::move(selected)); auto result = module_->execute(method.name, {tokens, positions, selector}); - if (!result.ok()) { - ET_LOG( - Error, - "CudaExecutor: %s failed with 0x%x", - method.name.c_str(), - static_cast(result.error())); + if (!result.ok() || result->empty() || !result->at(0).isTensor()) { + ET_LOG(Error, "CudaExecutor: raw %s failed", method.name.c_str()); return false; } ++calls_[slice.method]; - if (result->empty() || !result->at(0).isTensor()) { - ET_LOG(Error, "CudaExecutor: %s returned no logits", method.name.c_str()); - return false; - } - auto logits = result->at(0).toTensor(); + const auto logits = result->at(0).toTensor(); if (logits.dim() < 2 || logits.size(logits.dim() - 1) != vocab_size_ || - logits.numel() != static_cast(rows) * vocab_size_) { + logits.numel() != static_cast(rows) * vocab_size_ || + !is_supported_logits_type(logits.scalar_type()) || !is_dense(logits)) { ET_LOG( Error, - "CudaExecutor: %s returned logits that are not [%d, %d]", + "CudaExecutor: %s logits are not a dense [%zu, %d] tensor", method.name.c_str(), rows, vocab_size_); return false; } - if (selected.inputs.empty()) { - continue; + std::optional<::executorch::aten::Tensor> taps; + if (result->size() > 1) { + if (!result->at(1).isTensor()) { + return false; + } + taps = result->at(1).toTensor(); + if (taps->dim() != 3 || taps->size(0) != 1 || + taps->size(1) != slice.width || taps->size(2) <= 0 || + !is_supported_logits_type(taps->scalar_type()) || !is_dense(*taps)) { + ET_LOG( + Error, + "CudaExecutor: %s taps are not a dense [1, %d, H] tensor", + method.name.c_str(), + slice.width); + return false; + } + } + if (!consume({off, n, logits, taps ? &*taps : nullptr, indices})) { + return false; + } + } + return true; +} + +std::optional> CudaExecutor::argmax_rows( + const ::executorch::aten::Tensor& logits, + std::size_t rows) { + using ::executorch::aten::ScalarType; + using ::executorch::aten::SizesType; + if (!initialized_ || logits.dim() < 2 || !is_dense(logits) || + !is_supported_logits_type(logits.scalar_type()) || + logits.size(logits.dim() - 1) != vocab_size_ || + rows > static_cast(std::numeric_limits::max()) || + rows > static_cast(logits.numel() / vocab_size_)) { + return std::nullopt; + } + if (rows == 0) { + return std::vector{}; + } + if (logits.device_type() != ::executorch::aten::DeviceType::CUDA) { + std::vector result; + metadata::Sampler sampler(vocab_size_, 0.0f, 1.0f, 1); + for (std::size_t row = 0; row < rows; ++row) { + auto token = + llm_batching::sample_from_row(logits, static_cast(row), sampler); + if (!token) { + return std::nullopt; + } + result.push_back(*token); + } + return result; + } + if (!has_argmax_ || logits.scalar_type() != ScalarType::Float) { + ET_LOG(Error, "CudaExecutor: device raw logits require %s", kArgmaxMethod); + return std::nullopt; + } + auto view = make_tensor_ptr( + {static_cast(rows), static_cast(vocab_size_)}, + logits.mutable_data_ptr(), + logits.scalar_type(), + ::executorch::aten::Device( + ::executorch::aten::DeviceType::CUDA, logits.device_index())); + auto result = + module_->execute(kArgmaxMethod, std::vector{view}); + if (!result.ok() || result->size() != 1 || !result->at(0).isTensor()) { + return std::nullopt; + } + const auto tokens = result->at(0).toTensor(); + if (tokens.scalar_type() != ScalarType::Long || + tokens.numel() != static_cast(rows)) { + return std::nullopt; + } + std::vector host(rows); + const auto stream = ::executorch::extension::cuda::getCallerStream().value_or( + cudaStreamPerThread); + const auto kind = tokens.device_type() == ::executorch::aten::DeviceType::CUDA + ? cudaMemcpyDeviceToHost + : cudaMemcpyHostToHost; + if (cudaMemcpyAsync( + host.data(), + tokens.const_data_ptr(), + rows * sizeof(std::int64_t), + kind, + stream) != cudaSuccess || + cudaStreamSynchronize(stream) != cudaSuccess) { + return std::nullopt; + } + if (std::any_of(host.begin(), host.end(), [this](auto token) { + return token < 0 || token >= vocab_size_; + })) { + return std::nullopt; + } + return std::vector(host.begin(), host.end()); +} + +bool CudaExecutor::execute(const BatchInput& batch, BatchOutput& out) { + out.outputs.clear(); + out.outputs.resize(batch.inputs.size()); + + Result step = sessions_.pack(batch); + if (!step.ok()) { + return false; + } + + RawTargetStep raw; + raw.tokens = std::move(step->tokens); + raw.positions = std::move(step->positions); + raw.seq_ids = std::move(step->seq_ids); + raw.logit_indices.assign( + step->logit_indices.begin(), step->logit_indices.end()); + return execute_raw(raw, [&](const RawTargetOutput& result) { + auto logits = result.logits; + const auto& inputs = result.selected_indices; + if (inputs.empty()) { + return true; } std::vector row_sessions; - row_sessions.reserve(selected.inputs.size()); - for (const std::size_t input_index : selected.inputs) { + row_sessions.reserve(inputs.size()); + for (const std::size_t input_index : inputs) { row_sessions.push_back(batch.inputs[input_index].sid); } std::vector tokens_out; @@ -626,8 +775,7 @@ bool CudaExecutor::execute(const BatchInput& batch, BatchOutput& out) { if (device_logits_.value_or(device_logits) != device_logits) { ET_LOG( Error, - "CudaExecutor: %s returns %s logits, unlike the forwards before it", - method.name.c_str(), + "CudaExecutor: target returns %s logits, unlike the forwards before it", device_logits ? "device" : "host"); return false; } @@ -636,8 +784,7 @@ bool CudaExecutor::execute(const BatchInput& batch, BatchOutput& out) { if (!has_sampler_) { ET_LOG( Error, - "CudaExecutor: %s returns device logits but the program has no %s", - method.name.c_str(), + "CudaExecutor: target returns device logits but the program has no %s", kSampleMethod); return false; } @@ -656,12 +803,11 @@ bool CudaExecutor::execute(const BatchInput& batch, BatchOutput& out) { tokens_out.push_back(*token); } } - for (std::size_t row = 0; row < selected.inputs.size(); ++row) { - out.outputs[selected.inputs[row]] = - Output{row_sessions[row], {tokens_out[row]}}; + for (std::size_t row = 0; row < inputs.size(); ++row) { + out.outputs[inputs[row]] = Output{row_sessions[row], {tokens_out[row]}}; } - } - return true; + return true; + }); } std::optional> CudaExecutor::sample_on_device( diff --git a/backends/cuda/batching/cuda_executor.h b/backends/cuda/batching/cuda_executor.h index b50585fdf80..8acc77c82e2 100644 --- a/backends/cuda/batching/cuda_executor.h +++ b/backends/cuda/batching/cuda_executor.h @@ -20,6 +20,7 @@ #include #include +#include #include #include #include @@ -85,6 +86,26 @@ struct ForwardMethodCalls { std::uint64_t calls; }; +// A raw packed target step. The caller owns the cache sequence IDs and performs +// any rewind before calling. Do not mix caller-owned IDs with session APIs. +struct RawTargetStep { + std::vector tokens; + std::vector positions; + std::vector seq_ids; + // Global token-row indices; -1 discards a result. Repeats and reordering are + // allowed. Callback selected_indices maps rows back into this vector. + std::vector logit_indices; +}; + +struct RawTargetOutput { + int offset; + int length; + const ::executorch::aten::Tensor& logits; + // All exported token rows, including static padding. Only [0, length) count. + const ::executorch::aten::Tensor* taps; + const std::vector& selected_indices; +}; + class ET_EXPERIMENTAL CudaExecutor : public llm_batching::Executor { public: ~CudaExecutor() override; @@ -140,6 +161,26 @@ class ET_EXPERIMENTAL CudaExecutor : public llm_batching::Executor { const llm_batching::BatchInput& batch, llm_batching::BatchOutput& out) override; + // Engine-thread-only raw access for model adapters. Output storage is + // borrowed until the callback returns: copy/consume every needed output + // before running another method. The callback must not recursively call + // execute_raw(). A failed step/callback may have written KV; its owners must + // not continue. + bool execute_raw( + const RawTargetStep& step, + const std::function& consume); + // Samples no policy and advances no per-session RNG. Device logits require + // the exported sample_argmax method; only the selected token IDs return home. + std::optional> argmax_rows( + const ::executorch::aten::Tensor& logits, + std::size_t rows); + ::executorch::extension::Module& module() { + return *module_; + } + llm_cache::BatchControl& cache_control() { + return *ctl_; + } + // The cache's pool usage, for memory reporting. OffGraphKVMetrics kv_metrics() const; @@ -197,6 +238,8 @@ class ET_EXPERIMENTAL CudaExecutor : public llm_batching::Executor { // Where the first forward's logits lived. Host and device sampling advance // separate generators, so every forward must agree. std::optional device_logits_; + bool initialized_ = false; + bool raw_executing_ = false; }; } // namespace executorch::backends::cuda::batching diff --git a/backends/cuda/batching/targets.bzl b/backends/cuda/batching/targets.bzl index 2ad6a1f0541..3e018c236ad 100644 --- a/backends/cuda/batching/targets.bzl +++ b/backends/cuda/batching/targets.bzl @@ -38,7 +38,9 @@ def define_common_targets(is_fbcode = False): ], deps = [ "//executorch/extension/cuda:caller_stream", + "//executorch/extension/llm/batching:executor_utils", "//executorch/extension/llm/runner:stats", + "//executorch/extension/llm/sampler:sampler", "//executorch/extension/tensor:tensor", "//executorch/runtime/backend:interface", "//executorch/runtime/platform:platform",