cuda_executor.cpp

/*
 * 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 <executorch/backends/cuda/batching/cuda_executor.h>

#include <cuda_runtime.h>

#include <algorithm>
#include <cctype>
#include <cinttypes>
#include <limits>
#include <set>
#include <string_view>
#include <unordered_set>
#include <utility>
#include <vector>

#include <executorch/backends/cuda/runtime/backend_options.h>
#include <executorch/extension/cuda/caller_stream.h>
#include <executorch/extension/llm/batching/executor_utils.h>
#include <executorch/extension/llm/runner/model_metadata.h>
#include <executorch/extension/llm/sampler/sampler.h>
#include <executorch/extension/tensor/tensor_ptr.h>
#include <executorch/runtime/backend/backend_options_map.h>
#include <executorch/runtime/backend/interface.h>
#include <executorch/runtime/backend/options.h>
#include <executorch/runtime/platform/log.h>

namespace executorch::backends::cuda::batching {

using ::executorch::extension::make_tensor_ptr;
using ::executorch::extension::Module;
using ::executorch::extension::llm::LogitsToKeepMode;
using ::executorch::runtime::Error;
using ::executorch::runtime::MethodMeta;
using ::executorch::runtime::Result;
using llm_batching::BatchInput;
using llm_batching::BatchOutput;
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 {

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<int> 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<int>(token_sizes[1]);
}

// The logits width `method` produces.
Result<std::int64_t> 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<std::string> 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;
}

// The token count a static method's name carries: N for forward_{N}.
std::optional<int> static_method_width(std::string_view name) {
  const std::string_view prefix = kStaticMethodPrefix;
  if (name.size() <= prefix.size() || name.substr(0, prefix.size()) != prefix) {
    return std::nullopt;
  }
  int width = 0;
  for (const char c : name.substr(prefix.size())) {
    if (!std::isdigit(static_cast<unsigned char>(c)) ||
        width > (std::numeric_limits<int>::max() - 9) / 10) {
      return std::nullopt;
    }
    width = width * 10 + (c - '0');
  }
  return width > 0 ? std::optional<int>(width) : std::nullopt;
}

bool is_long_vector(const MethodMeta& meta, std::size_t output) {
  const auto tensor = meta.output_tensor_meta(output);
  return tensor.ok() &&
      tensor->scalar_type() == ::executorch::aten::ScalarType::Long &&
      tensor->sizes().size() == 1;
}

// Validates an optional device sampler; false when the program lacks it.
Result<bool> check_sampler(
    Module& module,
    const std::unordered_set<std::string>& names,
    const char* method,
    std::size_t num_inputs,
    std::int64_t vocab_size) {
  if (names.count(method) == 0) {
    return false;
  }
  ET_ASSIGN_OR_RETURN(meta, module.method_meta(method));
  bool valid = meta.num_inputs() == num_inputs && meta.num_outputs() == 1 &&
      is_long_vector(meta, 0);
  for (std::size_t i = 0; valid && i < num_inputs; ++i) {
    const auto input = meta.input_tensor_meta(i);
    valid = input.ok() && input->sizes().size() == 2 &&
        (i == 0
             ? is_supported_logits_type(input->scalar_type()) &&
                 input->sizes()[1] == vocab_size
             : input->scalar_type() == ::executorch::aten::ScalarType::Float &&
                 input->sizes()[1] == 4);
  }
  ET_CHECK_OR_RETURN_ERROR(
      valid,
      InvalidProgram,
      "CudaExecutor: %s must map [R, vocab] logits%s to Long [R] tokens",
      method,
      num_inputs == 2 ? " and Float [R, 4] params" : "");
  return true;
}

Error set_backend_options(
    const CudaExecutorOptions& options,
    const std::vector<ForwardMethod>& methods) {
  // The backend matches a comma-separated list of method names.
  std::string graph_methods;
  if (options.cuda_graph_for_static_methods) {
    for (const ForwardMethod& method : methods) {
      if (method.is_static) {
        graph_methods += (graph_methods.empty() ? "" : ",") + method.name;
      }
    }
  }
  for (const std::string& name : options.additional_cuda_graph_methods) {
    graph_methods += (graph_methods.empty() ? "" : ",") + name;
  }
  ET_CHECK_OR_RETURN_ERROR(
      graph_methods.size() < ::executorch::runtime::kMaxOptionValueLength,
      InvalidArgument,
      "CudaExecutor: too many static methods to name in one backend option");
  ::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", graph_methods.c_str()));
  return ::executorch::runtime::set_option(
      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(
    std::unique_ptr<Module> module,
    std::shared_ptr<llm_cache::Cache> cache,
    int max_sessions,
    int max_session_tokens,
    std::string backend_id,
    std::int32_t vocab_size,
    std::vector<ForwardMethod> methods,
    int min_selected_rows,
    bool has_sampler,
    bool has_argmax)
    : install_guard_(cache),
      module_(std::move(module)),
      ctl_(cache->as<llm_cache::BatchControl>()),
      kv_(cache->as<CudaKVCache>()),
      backend_id_(std::move(backend_id)),
      vocab_size_(vocab_size),
      max_session_tokens_(max_session_tokens),
      methods_(std::move(methods)),
      calls_(methods_.size(), 0),
      min_selected_rows_(min_selected_rows),
      sessions_(*ctl_, max_sessions, max_session_tokens, vocab_size),
      has_sampler_(has_sampler),
      has_argmax_(has_argmax) {}

CudaExecutor::~CudaExecutor() {
  if (device_params_ != nullptr) {
    (void)cudaFree(device_params_);
  }
}

Result<std::unique_ptr<CudaExecutor>> CudaExecutor::create(
    const CudaExecutorConfig& config) {
  ET_CHECK_OR_RETURN_ERROR(
      !config.model_path.empty(),
      InvalidArgument,
      "CudaExecutor: no model path");
  auto module = std::make_unique<Module>(
      config.model_path,
      config.data_files,
      config.load_mode,
      /*event_tracer=*/nullptr,
      /*memory_allocator=*/nullptr,
      /*temp_allocator=*/nullptr,
      /*share_memory_arenas=*/false);
  return create(std::move(module), config);
}

Result<std::unique_ptr<CudaExecutor>> CudaExecutor::create(
    std::unique_ptr<Module> module,
    int max_sessions,
    int max_session_tokens,
    int kv_dtype,
    int initial_capacity,
    CudaExecutorOptions options) {
  ET_CHECK_OR_RETURN_ERROR(
      max_session_tokens > 0,
      InvalidArgument,
      "CudaExecutor: invalid token limit");
  CudaExecutorConfig config;
  config.max_sessions = max_sessions;
  config.max_session_tokens = max_session_tokens;
  config.kv_dtype = kv_dtype;
  config.initial_capacity = initial_capacity;
  config.options = std::move(options);
  return create(std::move(module), config);
}

Result<std::unique_ptr<CudaExecutor>> CudaExecutor::create(
    std::unique_ptr<Module> module,
    const CudaExecutorConfig& config) {
  const int max_sessions = config.max_sessions;
  int max_session_tokens = config.max_session_tokens;
  const int kv_dtype = config.kv_dtype;
  const int initial_capacity = config.initial_capacity;
  const auto& options = config.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_session_tokens == 0 ||
           max_sessions <=
               std::numeric_limits<int>::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.
  // Static methods select one row per token, 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_CHECK_OR_RETURN_ERROR(
      max_cells > 1 && max_cells <= std::numeric_limits<int>::max(),
      InvalidProgram,
      "CudaExecutor: cell count must fit an int and leave a padding row");
  if (max_session_tokens == 0) {
    max_session_tokens = static_cast<int>(
        std::min(max_context_length, (max_cells - 1) / max_sessions));
    ET_CHECK_OR_RETURN_ERROR(
        max_session_tokens > 0,
        InvalidArgument,
        "CudaExecutor: no cell budget remains for each resident session");
  }

  // Every forward_{N} and forward_others the program exports, narrowest first.
  ET_ASSIGN_OR_RETURN(names, module->method_names());
  std::vector<ForwardMethod> methods;
  std::string backend_id;
  std::int64_t method_vocab = 0;
  std::set<int> static_widths;
  for (const std::string& name : names) {
    const std::optional<int> declared = static_method_width(name);
    const bool dynamic = name == kDynamicMethod;
    if (!declared && !dynamic) {
      continue;
    }
    ET_ASSIGN_OR_RETURN(meta, module->method_meta(name.c_str()));
    ET_ASSIGN_OR_RETURN(width, step_width(meta, name.c_str()));
    ET_ASSIGN_OR_RETURN(vocab, logits_width(meta, name.c_str()));
    ET_ASSIGN_OR_RETURN(backend, backend_of(meta, name.c_str()));
    ET_CHECK_OR_RETURN_ERROR(
        !declared || width == *declared,
        InvalidProgram,
        "CudaExecutor: %s takes %d tokens",
        name.c_str(),
        width);
    ET_CHECK_OR_RETURN_ERROR(
        method_vocab == 0 || vocab == method_vocab,
        InvalidProgram,
        "CudaExecutor: the forward methods disagree on the vocabulary");
    method_vocab = vocab;
    backend_id = backend;
    if (declared) {
      static_widths.insert(width);
    }
    methods.push_back({name, width, !dynamic});
  }
  ET_CHECK_OR_RETURN_ERROR(
      !static_widths.empty(),
      InvalidProgram,
      "CudaExecutor: the program exports no %sN method",
      kStaticMethodPrefix);
  std::sort(
      methods.begin(),
      methods.end(),
      [](const ForwardMethod& a, const ForwardMethod& b) {
        return a.max_tokens < b.max_tokens;
      });
  // The dynamic method serves only what no static method holds, so it must be
  // the widest.
  const auto dynamic =
      std::find_if(methods.begin(), methods.end(), [](const ForwardMethod& m) {
        return !m.is_static;
      });
  ET_CHECK_OR_RETURN_ERROR(
      dynamic == methods.end() || dynamic->max_tokens > *static_widths.rbegin(),
      InvalidProgram,
      "CudaExecutor: %s must be wider than every static method",
      kDynamicMethod);
  ET_CHECK_OR_RETURN_ERROR(
      std::count_if(
          methods.begin(),
          methods.end(),
          [](const ForwardMethod& m) { return !m.is_static; }) <= 1,
      InvalidProgram,
      "CudaExecutor: more than one dynamic method");
  const int max_step_tokens = methods.back().max_tokens;
  ET_CHECK_OR_RETURN_ERROR(
      max_step_tokens < max_cells,
      InvalidProgram,
      "CudaExecutor: a %d-token step does not fit %" PRId64 " cells",
      max_step_tokens,
      max_cells);
  // A dynamic method whose kernels switch at a small width is exported from
  // more selected rows than one; its selector is padded up to that.
  int min_selected_rows = 1;
  if (dynamic != methods.end()) {
    ET_ASSIGN_OR_RETURN(
        declared_min_rows,
        metadata::detail::read_int_method(*module, kMinSelectedRowsMethod));
    const std::int64_t requested_min_rows =
        std::max<std::int64_t>(1, declared_min_rows.value_or(1));
    ET_CHECK_OR_RETURN_ERROR(
        requested_min_rows <= max_step_tokens,
        InvalidProgram,
        "CudaExecutor: %" PRId64 " selected rows exceed the widest step",
        requested_min_rows);
    min_selected_rows = static_cast<int>(requested_min_rows);
  }
  ET_ASSIGN_OR_RETURN(published_vocab, metadata::read_vocab_size(*module));
  ET_ASSIGN_OR_RETURN(
      vocab_size, metadata::check_vocab_size(published_vocab, method_vocab));
  ET_ASSIGN_OR_RETURN(
      has_sampler, check_sampler(*module, names, kSampleMethod, 2, vocab_size));
  ET_ASSIGN_OR_RETURN(
      has_argmax, check_sampler(*module, names, kArgmaxMethod, 1, vocab_size));
  ET_CHECK_OR_RETURN_ERROR(
      has_sampler || !has_argmax,
      InvalidProgram,
      "CudaExecutor: %s needs %s for non-greedy rows",
      kArgmaxMethod,
      kSampleMethod);

  // 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. One of the
  // program's cells is the padding scratch row, never handed out.
  ET_CHECK_OR_RETURN_ERROR(
      static_cast<std::int64_t>(max_sessions) * max_session_tokens < max_cells,
      InvalidArgument,
      "CudaExecutor: %d sessions of %d tokens exceed the program's %" PRId64
      " cells, one of which is reserved",
      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{};
  cfg.capacity = static_cast<int>(max_cells);
  cfg.kv_dtype = 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<llm_cache::Cache> cache = built.get();
  auto* const ctl = cache->as<llm_cache::BatchControl>();
  ET_CHECK_OR_RETURN_ERROR(
      ctl != nullptr && cache->as<CudaKVCache>() != nullptr,
      InvalidType,
      "CudaExecutor: the cache is not a batched CUDA cache");
  const std::optional<int> 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, methods));
  return std::unique_ptr<CudaExecutor>(new CudaExecutor(
      std::move(module),
      std::move(cache),
      max_sessions,
      max_session_tokens,
      std::move(backend_id),
      vocab_size,
      std::move(methods),
      min_selected_rows,
      has_sampler,
      has_argmax));
}

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 ForwardMethod& method : methods_) {
    if (module_->load_method(
            method.name,
            /*planned_memory=*/nullptr,
            /*event_tracer=*/nullptr,
            &options_map) != Error::Ok) {
      ET_LOG(Error, "CudaExecutor: could not load %s", method.name.c_str());
      return false;
    }
  }
  for (const auto& [present, name] :
       {std::pair{has_sampler_, kSampleMethod},
        std::pair{has_argmax_, kArgmaxMethod}}) {
    if (present && module_->load_method(name) != Error::Ok) {
      ET_LOG(Error, "CudaExecutor: could not load %s", name);
      return false;
    }
  }
  initialized_ = true;
  return true;
}

std::optional<SessionId> CudaExecutor::open_session() {
  return sessions_.open();
}

void CudaExecutor::close_session(SessionId session) {
  sessions_.close(session);
}

std::optional<SessionId> CudaExecutor::clone(SessionId source, Position upto) {
  return sessions_.clone(source, upto);
}

void CudaExecutor::set_sampling(
    SessionId session,
    const SamplingParams& params,
    std::optional<std::uint64_t> seed) {
  sessions_.set_sampling(session, params, seed);
}

bool CudaExecutor::execute_raw(
    const RawTargetStep& step,
    const std::function<bool(const RawTargetOutput&)>& consume) {
  if (!initialized_ || raw_executing_ || !consume ||
      step.tokens.size() >
          static_cast<std::size_t>(std::numeric_limits<int>::max()) ||
      step.positions.size() != step.tokens.size() ||
      step.seq_ids.size() != step.tokens.size()) {
    return false;
  }
  const int total = static_cast<int>(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 auto& method = methods_[slice.method];
    std::vector<std::int64_t> selected;
    std::vector<std::size_t> 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<std::size_t>(
        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<std::size_t>(selected.size(), min_selected_rows_);
    selected.resize(rows, selected.empty() ? 0 : selected.back());
    if (!ctl_->declare_step(std::vector<std::int32_t>(
            step.seq_ids.begin() + off, step.seq_ids.begin() + off + n))) {
      return false;
    }
    std::vector<std::int64_t> token_values(
        step.tokens.begin() + off, step.tokens.begin() + off + n);
    std::vector<std::int64_t> position_values(
        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));
    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() || result->empty() || !result->at(0).isTensor()) {
      ET_LOG(Error, "CudaExecutor: raw %s failed", method.name.c_str());
      return false;
    }
    ++calls_[slice.method];
    const auto logits = result->at(0).toTensor();
    if (logits.dim() < 2 || logits.size(logits.dim() - 1) != vocab_size_ ||
        logits.numel() != static_cast<std::int64_t>(rows) * vocab_size_ ||
        !is_supported_logits_type(logits.scalar_type()) || !is_dense(logits)) {
      ET_LOG(
          Error,
          "CudaExecutor: %s logits are not a dense [%zu, %d] tensor",
          method.name.c_str(),
          rows,
          vocab_size_);
      return false;
    }
    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<std::vector<Token>> 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::size_t>(std::numeric_limits<SizesType>::max()) ||
      rows > static_cast<std::size_t>(logits.numel() / vocab_size_)) {
    return std::nullopt;
  }
  if (rows == 0) {
    return std::vector<Token>{};
  }
  if (logits.device_type() != ::executorch::aten::DeviceType::CUDA) {
    std::vector<Token> 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<int>(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<SizesType>(rows), static_cast<SizesType>(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<runtime::EValue>{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<std::int64_t>(rows)) {
    return std::nullopt;
  }
  std::vector<std::int64_t> 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<Token>(host.begin(), host.end());
}

bool CudaExecutor::execute(const BatchInput& batch, BatchOutput& out) {
  out.outputs.clear();
  out.outputs.resize(batch.inputs.size());

  Result<llm_batching::util::PackedStep> 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<SessionId> row_sessions;
    row_sessions.reserve(inputs.size());
    for (const std::size_t input_index : inputs) {
      row_sessions.push_back(batch.inputs[input_index].sid);
    }
    std::vector<Token> tokens_out;
    const bool device_logits =
        logits.device_type() == ::executorch::aten::DeviceType::CUDA;
    if (device_logits_.value_or(device_logits) != device_logits) {
      ET_LOG(
          Error,
          "CudaExecutor: target returns %s logits, unlike the forwards before it",
          device_logits ? "device" : "host");
      return false;
    }
    device_logits_ = device_logits;
    if (device_logits) {
      if (!has_sampler_) {
        ET_LOG(
            Error,
            "CudaExecutor: target returns device logits but the program has no %s",
            kSampleMethod);
        return false;
      }
      auto sampled = sample_on_device(logits, row_sessions);
      if (!sampled) {
        return false;
      }
      tokens_out = std::move(*sampled);
    } else {
      for (std::size_t row = 0; row < row_sessions.size(); ++row) {
        const std::optional<Token> token =
            sessions_.sample(row_sessions[row], logits, static_cast<int>(row));
        if (!token) {
          return false;
        }
        tokens_out.push_back(*token);
      }
    }
    for (std::size_t row = 0; row < inputs.size(); ++row) {
      out.outputs[inputs[row]] = Output{row_sessions[row], {tokens_out[row]}};
    }
    return true;
  });
}

std::optional<std::vector<Token>> CudaExecutor::sample_rows(
    const ::executorch::aten::Tensor& logits,
    std::size_t rows,
    const std::vector<float>& params) {
  if (params.size() != rows * 4) {
    return std::nullopt;
  }
  bool all_greedy = true;
  for (std::size_t row = 0; row < rows; ++row) {
    all_greedy = all_greedy && params[row * 4] == 0.0f;
  }
  if (all_greedy) {
    return argmax_rows(logits, rows);
  }
  if (logits.device_type() != ::executorch::aten::DeviceType::CUDA) {
    ET_LOG(
        Error,
        "CudaExecutor: sampled rows need device logits and %s",
        kSampleMethod);
    return std::nullopt;
  }
  if (!initialized_ || !has_sampler_ || logits.dim() < 2 || !is_dense(logits) ||
      logits.size(logits.dim() - 1) != vocab_size_ ||
      logits.scalar_type() != ::executorch::aten::ScalarType::Float ||
      rows == 0 ||
      rows > static_cast<std::size_t>(logits.numel() / vocab_size_)) {
    return std::nullopt;
  }
  return run_device_sampler(logits, rows, params, false);
}

std::optional<std::vector<Token>> CudaExecutor::sample_on_device(
    const ::executorch::aten::Tensor& logits,
    const std::vector<SessionId>& sessions) {
  const bool all_greedy =
      std::all_of(sessions.begin(), sessions.end(), [this](SessionId session) {
        return sessions_.greedy(session);
      });
  // Every row's generator advances whichever method runs, as on the host.
  std::vector<float> params;
  params.reserve(sessions.size() * 4);
  for (const SessionId session : sessions) {
    const auto row = sessions_.device_sampling(session);
    if (!row) {
      return std::nullopt;
    }
    params.insert(
        params.end(), {row->temperature, row->top_p, row->top_k, row->coin});
  }
  return run_device_sampler(logits, sessions.size(), params, all_greedy);
}

std::optional<std::vector<Token>> CudaExecutor::run_device_sampler(
    const ::executorch::aten::Tensor& logits,
    std::size_t row_count,
    const std::vector<float>& params,
    bool all_greedy) {
  using ::executorch::aten::ScalarType;
  using ::executorch::aten::SizesType;
  const ::executorch::aten::Device cuda(
      ::executorch::aten::DeviceType::CUDA, logits.device_index());
  const auto rows = static_cast<SizesType>(row_count);
  // The rows that sample lead the logits; the rest are padding.
  auto row_logits = make_tensor_ptr(
      {rows, static_cast<SizesType>(vocab_size_)},
      logits.mutable_data_ptr(),
      logits.scalar_type(),
      cuda);

  const cudaStream_t stream =
      ::executorch::extension::cuda::getCallerStream().value_or(
          cudaStreamPerThread);
  auto run_sampler =
      [&]() -> Result<std::vector<::executorch::runtime::EValue>> {
    if (all_greedy && has_argmax_) {
      return module_->execute(
          kArgmaxMethod,
          std::vector<::executorch::runtime::EValue>{row_logits});
    }
    if (row_count > device_params_rows_) {
      if (device_params_ != nullptr) {
        (void)cudaFree(device_params_);
        device_params_ = nullptr;
        device_params_rows_ = 0;
      }
      ET_CHECK_OR_RETURN_ERROR(
          cudaMalloc(&device_params_, row_count * 4 * sizeof(float)) ==
              cudaSuccess,
          MemoryAllocationFailed,
          "CudaExecutor: could not stage the sampler's params");
      device_params_rows_ = row_count;
    }
    // The backend runs every method on the caller's stream if one is set,
    // else the per-thread stream; ordering the upload there puts it ahead of
    // the sampler.
    ET_CHECK_OR_RETURN_ERROR(
        cudaMemcpyAsync(
            device_params_,
            params.data(),
            params.size() * sizeof(float),
            cudaMemcpyHostToDevice,
            stream) == cudaSuccess,
        Internal,
        "CudaExecutor: could not upload the sampler's params");
    auto device_params =
        make_tensor_ptr({rows, 4}, device_params_, ScalarType::Float, cuda);
    return module_->execute(kSampleMethod, {row_logits, device_params});
  };
  const auto result = run_sampler();
  if (!result.ok() || result->empty() || !result->at(0).isTensor()) {
    ET_LOG(Error, "CudaExecutor: the device sampler failed");
    return std::nullopt;
  }
  const auto& tokens = result->at(0).toTensor();
  if (tokens.numel() != rows || tokens.scalar_type() != ScalarType::Long) {
    ET_LOG(Error, "CudaExecutor: the device sampler returned a bad shape");
    return std::nullopt;
  }
  std::vector<std::int64_t> host(row_count);
  const auto kind = tokens.device_type() == ::executorch::aten::DeviceType::CUDA
      ? cudaMemcpyDeviceToHost
      : cudaMemcpyHostToHost;
  if (cudaMemcpyAsync(
          host.data(),
          tokens.const_data_ptr(),
          host.size() * sizeof(std::int64_t),
          kind,
          stream) != cudaSuccess ||
      cudaStreamSynchronize(stream) != cudaSuccess) {
    ET_LOG(Error, "CudaExecutor: could not read the sampled tokens");
    return std::nullopt;
  }
  return std::vector<Token>(host.begin(), host.end());
}

OffGraphKVMetrics CudaExecutor::kv_metrics() const {
  return kv_->metrics();
}

std::vector<ForwardMethodCalls> CudaExecutor::method_calls() const {
  std::vector<ForwardMethodCalls> calls;
  calls.reserve(methods_.size());
  for (std::size_t i = 0; i < methods_.size(); ++i) {
    calls.push_back({methods_[i].name, calls_[i]});
  }
  return calls;
}

} // namespace executorch::backends::cuda::batching