Skip to content

[executorch][muse] Run the MG batching runtime on CudaExecutor #24007

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Gasoonjia merged 1 commit into from
Oct 20, 2026
Merged

[executorch][muse] Run the MG batching runtime on CudaExecutor #24007

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions Makefile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,7 @@ help:
@echo " muse-glimmer-cuda - Build Muse Glimmer runners and worker with CUDA backend"
@echo " muse-glimmer-mlx - Build Muse Glimmer runners and worker with MLX backend"
@echo " muse-glimmer-batching-mlx - Build batching Muse Glimmer runner and worker with MLX"
@echo " muse-glimmer-batching-cuda - Build batching Muse Glimmer runner and worker with CUDA"
@echo " qwen3_5_moe-cuda - Build Qwen3.5 MoE runner with CUDA backend"
@echo " qwen3_5_moe-metal - Build Qwen3.5 MoE runner with Metal backend"
@echo " qwen3_5_moe-mlx - Build Qwen3.5 MoE runner with MLX backend"
Expand Down Expand Up @@ -532,6 +533,23 @@ muse-glimmer-batching-mlx:
@echo " Runner: $(CURDIR)/cmake-out/examples/models/muse-glimmer-batching/muse_glimmer_batching_runner"
@echo " Worker: $(CURDIR)/cmake-out/examples/models/muse-glimmer-batching/muse_glimmer_batching_worker"

.PHONY: muse-glimmer-batching-cuda
muse-glimmer-batching-cuda:
@echo "==> Building and installing ExecuTorch with CUDA..."
cmake --workflow --preset llm-release-cuda
@echo "==> Configuring batching Muse Glimmer runners with CUDA..."
cd examples/models/muse-glimmer && cmake --preset muse-glimmer-cuda \
-B "$(CURDIR)/cmake-out/examples/models/muse-glimmer-batching" \
-DMUSE_GLIMMER_BUILD_BATCHING=ON \
-DMUSE_GLIMMER_BUILD_BATCHING_TESTS=OFF
cmake --build "$(CURDIR)/cmake-out/examples/models/muse-glimmer-batching" \
--target muse_glimmer_batching_runner muse_glimmer_batching_worker \
--parallel "$(nproc)"
@echo ""
@echo "Build complete!"
@echo " Runner: $(CURDIR)/cmake-out/examples/models/muse-glimmer-batching/muse_glimmer_batching_runner"
@echo " Worker: $(CURDIR)/cmake-out/examples/models/muse-glimmer-batching/muse_glimmer_batching_worker"

gemma4_31b-cuda:
@echo "==> Building and installing ExecuTorch with CUDA..."
cmake --workflow --preset llm-release-cuda
Expand Down
58 changes: 56 additions & 2 deletions backends/cuda/batching/cuda_executor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -228,6 +228,7 @@ CudaExecutor::CudaExecutor(
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),
Expand All @@ -241,18 +242,58 @@ CudaExecutor::~CudaExecutor() {
}
}

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 = 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_sessions <= std::numeric_limits<int>::max() / max_session_tokens,
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());
Expand All @@ -279,6 +320,19 @@ Result<std::unique_ptr<CudaExecutor>> CudaExecutor::create(
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;
Expand Down
29 changes: 29 additions & 0 deletions backends/cuda/batching/cuda_executor.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,19 @@ struct CudaExecutorOptions {
bool cuda_graph_for_static_methods = true;
};

struct CudaExecutorConfig {
std::string model_path;
std::vector<std::string> data_files;
::executorch::extension::Module::LoadMode load_mode =
::executorch::extension::Module::LoadMode::MmapUseMlockIgnoreErrors;
int max_sessions = 1;
// Zero divides the available cell pool across sessions, capped by context.
int max_session_tokens = 0;
int kv_dtype = static_cast<int>(::executorch::aten::ScalarType::BFloat16);
int initial_capacity = -1;
CudaExecutorOptions options;
};

// How often each method ran, for reporting.
struct ForwardMethodCalls {
std::string name;
Expand All @@ -76,6 +89,17 @@ class ET_EXPERIMENTAL CudaExecutor : public llm_batching::Executor {
public:
~CudaExecutor() override;

static ::executorch::runtime::Result<std::unique_ptr<CudaExecutor>> create(
const CudaExecutorConfig& config);

int max_session_tokens() const {
return max_session_tokens_;
}

std::int32_t vocab_size() const {
return vocab_size_;
}

// 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
Expand Down Expand Up @@ -129,6 +153,10 @@ class ET_EXPERIMENTAL CudaExecutor : public llm_batching::Executor {
}

private:
static ::executorch::runtime::Result<std::unique_ptr<CudaExecutor>> create(
std::unique_ptr<::executorch::extension::Module> module,
const CudaExecutorConfig& config);

CudaExecutor(
std::unique_ptr<::executorch::extension::Module> module,
std::shared_ptr<llm_cache::Cache> cache,
Expand All @@ -155,6 +183,7 @@ class ET_EXPERIMENTAL CudaExecutor : public llm_batching::Executor {
const CudaKVCache* const kv_;
std::string backend_id_;
std::int32_t vocab_size_;
int max_session_tokens_;
// Sorted by max_tokens, ascending; the dynamic method, if any, is last.
std::vector<ForwardMethod> methods_;
std::vector<std::uint64_t> calls_;
Expand Down
89 changes: 89 additions & 0 deletions backends/cuda/batching/test/test_cuda_executor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -421,6 +421,95 @@ TEST_P(CudaExecutorTest, RefusesLimitsThePoolCannotHold) {
Error::InvalidArgument);
}

TEST(CudaExecutorConfigTest, InvalidConfigIsRejectedBeforeLoading) {
cb::CudaExecutorConfig config;
EXPECT_EQ(cb::CudaExecutor::create(config).error(), Error::InvalidArgument);

// A missing program must not mask invalid session limits, even without CUDA.
config.model_path = "/nonexistent/cuda_executor_test/model.pte";
config.load_mode = Module::LoadMode::File;
config.max_sessions = 0;
EXPECT_EQ(cb::CudaExecutor::create(config).error(), Error::InvalidArgument);

config.max_sessions = 1;
config.max_session_tokens = -1;
EXPECT_EQ(cb::CudaExecutor::create(config).error(), Error::InvalidArgument);
}

TEST_P(CudaExecutorTest, ConfigAutoBudgetRespectsContextAndPaddingCell) {
const std::vector<std::pair<int, int>> limits = {
{1, kMaxContext}, {4, 63}, {8, 31}};
for (const auto& [sessions, expected_budget] : limits) {
SCOPED_TRACE(sessions);
cb::CudaExecutorConfig config;
config.model_path = dir_ + "/model.pte";
config.data_files = {dir_ + "/aoti_cuda_blob.ptd"};
config.load_mode = Module::LoadMode::File;
config.max_sessions = sessions;
config.initial_capacity = 16;
config.options.cuda_graph_for_static_methods = false;
auto created = cb::CudaExecutor::create(config);
ASSERT_EQ(created.error(), Error::Ok);
auto exec = std::move(created.get());
ASSERT_NE(exec, nullptr);
EXPECT_EQ(exec->max_session_tokens(), expected_budget);
EXPECT_LT(exec->max_session_tokens() * sessions, kMaxCells);
// The toy export publishes a 32-entry vocabulary.
EXPECT_EQ(exec->vocab_size(), 32);
}
}

TEST_P(CudaExecutorTest, ConfiguredBudgetsMatchEagerGreedy) {
ASSERT_EQ(cases_.size(), 4u);
for (const int token_budget : {0, 32}) {
SCOPED_TRACE(token_budget);
cb::CudaExecutorConfig config;
config.model_path = dir_ + "/model.pte";
config.data_files = {dir_ + "/aoti_cuda_blob.ptd"};
config.load_mode = Module::LoadMode::File;
config.max_sessions = 4;
config.max_session_tokens = token_budget;
config.kv_dtype = kBFloat16;
config.initial_capacity = 16;
config.options.cuda_graph_for_static_methods = false;
auto created = cb::CudaExecutor::create(config);
ASSERT_EQ(created.error(), Error::Ok);
auto exec = std::move(created.get());
ASSERT_NE(exec, nullptr);
EXPECT_EQ(exec->max_session_tokens(), token_budget == 0 ? 63 : 32);
batching::Runner runner(*exec, scheduler());
const auto generations = generate(runner, prompts());
runner.shutdown();
ASSERT_EQ(generations.size(), cases_.size());
for (size_t i = 0; i < cases_.size(); ++i) {
EXPECT_EQ(generations[i].reason, batching::FinishReason::NewTokenLimit)
<< "prompt " << i;
EXPECT_EQ(generations[i].tokens, cases_[i].expected) << "prompt " << i;
}
EXPECT_EQ(runner.metrics().steps_failed, 0u);
}
}

TEST_P(CudaExecutorTest, ConfigRefusesOvercommitAndLegacyZeroBudget) {
cb::CudaExecutorConfig config;
config.model_path = dir_ + "/model.pte";
config.data_files = {dir_ + "/aoti_cuda_blob.ptd"};
config.load_mode = Module::LoadMode::File;
config.max_sessions = 8;
config.max_session_tokens = kMaxContext;
config.initial_capacity = 16;
config.options.cuda_graph_for_static_methods = false;
EXPECT_EQ(cb::CudaExecutor::create(config).error(), Error::InvalidArgument);

// Reserving the padding row leaves no token per session at 256 sessions.
config.max_sessions = kMaxCells;
config.max_session_tokens = 0;
EXPECT_EQ(cb::CudaExecutor::create(config).error(), Error::InvalidArgument);
EXPECT_EQ(
cb::CudaExecutor::create(module(), 4, 0, kBFloat16, 16, {}).error(),
Error::InvalidArgument);
}

INSTANTIATE_TEST_SUITE_P(
Toy,
CudaExecutorTest,
Expand Down
48 changes: 37 additions & 11 deletions examples/models/muse-glimmer/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -190,19 +190,22 @@ if(TARGET mlxdelegate)
executorch_target_copy_mlx_metallib(muse_glimmer_worker)
endif()

option(MUSE_GLIMMER_BUILD_BATCHING "Build MLX solo/image batching runners" OFF)
option(MUSE_GLIMMER_BUILD_BATCHING
"Build the batching runner and worker (MLX or CUDA)" OFF
)
option(MUSE_GLIMMER_BUILD_BATCHING_TESTS
"Build preparation/materialization tests without the batching executor"
OFF
)
if(MUSE_GLIMMER_BUILD_BATCHING OR MUSE_GLIMMER_BUILD_BATCHING_TESTS)
if(NOT EXECUTORCH_BUILD_MLX
OR NOT TARGET mlxdelegate
OR EXECUTORCH_BUILD_CUDA
)
if(EXECUTORCH_BUILD_CUDA)
set(_batching_backend cuda)
elseif(EXECUTORCH_BUILD_MLX AND TARGET mlxdelegate)
set(_batching_backend mlx)
else()
message(
FATAL_ERROR
"Muse Glimmer batching targets require EXECUTORCH_BUILD_MLX=ON and an MLX-enabled ExecuTorch installation (without CUDA)"
"Muse Glimmer batching targets require EXECUTORCH_BUILD_CUDA=ON or EXECUTORCH_BUILD_MLX=ON with an MLX-enabled ExecuTorch installation"
)
endif()
foreach(_dependency extension_module extension_tensor tokenizers::tokenizers)
Expand All @@ -226,30 +229,41 @@ if(MUSE_GLIMMER_BUILD_BATCHING OR MUSE_GLIMMER_BUILD_BATCHING_TESTS)
muse_glimmer_batching_preparation PUBLIC extension_module extension_tensor
tokenizers::tokenizers
)
if(_batching_backend STREQUAL "cuda")
# The vision runtime copies device-resident outputs back to the host.
target_link_libraries(muse_glimmer_batching_preparation PUBLIC CUDA::cudart)
endif()

if(MUSE_GLIMMER_BUILD_BATCHING)
if(_batching_backend STREQUAL "cuda")
set(_batching_executor)
set(_batching_backend_library cuda_batching)
else()
set(_batching_executor runtime/mlx_executor.cpp)
set(_batching_backend_library mlx)
endif()
set(_batching_dependencies
extension_llm_batching
extension_llm_cache
extension_llm_sampler
extension_module
extension_tensor
extension_llm_serving
mlx
${_batching_backend_library}
)
foreach(_dependency IN ITEMS ${_batching_dependencies}
extension_llm_batching_executor_utils
)
if(NOT TARGET ${_dependency})
message(
FATAL_ERROR
"Muse Glimmer batching runtime requires ${_dependency}; install ExecuTorch with EXECUTORCH_BUILD_MLX=ON, EXECUTORCH_BUILD_EXTENSION_LLM=ON and EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER=ON"
"Muse Glimmer batching runtime requires ${_dependency}; install ExecuTorch with EXECUTORCH_BUILD_CUDA=ON or EXECUTORCH_BUILD_MLX=ON, EXECUTORCH_BUILD_EXTENSION_LLM=ON and EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER=ON"
)
endif()
endforeach()
add_library(
muse_glimmer_batching_runtime STATIC
runtime/backend_factory.cpp runtime/mlx_executor.cpp
runtime/backend_factory.cpp ${_batching_executor}
runtime/runners/batching_bootstrap.cpp
)
target_link_libraries(
Expand Down Expand Up @@ -291,8 +305,20 @@ if(MUSE_GLIMMER_BUILD_BATCHING OR MUSE_GLIMMER_BUILD_BATCHING_TESTS)
muse_glimmer_batching_worker PRIVATE muse_glimmer_batching_runtime
${_batching_transport}
)
executorch_target_copy_mlx_metallib(muse_glimmer_batching_runner)
executorch_target_copy_mlx_metallib(muse_glimmer_batching_worker)
if(_batching_backend STREQUAL "cuda")
if(NOT APPLE AND NOT MSVC)
# CUDA AOTI blobs resolve symbols from the host executable.
target_link_options(
muse_glimmer_batching_runner PRIVATE "LINKER:--export-dynamic"
)
target_link_options(
muse_glimmer_batching_worker PRIVATE "LINKER:--export-dynamic"
)
endif()
else()
executorch_target_copy_mlx_metallib(muse_glimmer_batching_runner)
executorch_target_copy_mlx_metallib(muse_glimmer_batching_worker)
endif()
endif()

if(MUSE_GLIMMER_BUILD_BATCHING_TESTS)
Expand Down
Loading
Loading