From b4fbe74c011c241d14e2256e1f7aff0089a80db7 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Mon, 19 Oct 2026 13:07:56 -0700 Subject: [PATCH] [executorch][muse] Run the MG batching runtime on CudaExecutor Pull Request resolved: https://github.com/pytorch/executorch/pull/23731 The batching runner and multiplexed worker use the same ServingRuntime, Runner and DecodeFirstScheduler stack on CUDA and MLX; only the executor differs. CudaExecutorConfig now provides model/external-data loading, load mode, session limits, KV dtype, initial pool capacity and backend options to CudaExecutor::create(config). The existing module-based overload remains supported. Both entry points share metadata validation, forward-method discovery, cache construction and graph/weight-sharing setup. A zero config token limit divides the cell pool across resident sessions after reserving its padding row, capped by the published context. The existing Muse backend adapter calls this reusable API and retains only Muse-specific vocabulary/BOS validation and text-only preparation. Remove the separate CUDA backend factory files. Image requests remain rejected for CUDA batching artifacts without a vision path. CMake batching runner/preparation-test targets accept CUDA or MLX, and the CUDA make target builds the shared serving integration. Configuration tests cover automatic and explicit limits, greedy generation parity and invalid/overcommitted configurations. ghstack-source-id: 451855080 @exported-using-ghexport Differential Revision: [D124871959](https://our.internmc.facebook.com/intern/diff/D124871959/) --- Makefile | 18 ++++ backends/cuda/batching/cuda_executor.cpp | 58 +++++++++++- backends/cuda/batching/cuda_executor.h | 29 ++++++ .../cuda/batching/test/test_cuda_executor.cpp | 89 +++++++++++++++++++ examples/models/muse-glimmer/CMakeLists.txt | 48 +++++++--- .../muse-glimmer/runtime/backend_factory.cpp | 60 ++++++++++++- .../runtime/runners/batching_bootstrap.cpp | 9 +- 7 files changed, 294 insertions(+), 17 deletions(-) diff --git a/Makefile b/Makefile index 3f5ba32b0e4..26b6cb73ed5 100644 --- a/Makefile +++ b/Makefile @@ -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" @@ -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 diff --git a/backends/cuda/batching/cuda_executor.cpp b/backends/cuda/batching/cuda_executor.cpp index bb2cff21dc1..a64520b447e 100644 --- a/backends/cuda/batching/cuda_executor.cpp +++ b/backends/cuda/batching/cuda_executor.cpp @@ -228,6 +228,7 @@ CudaExecutor::CudaExecutor( kv_(cache->as()), 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), @@ -241,6 +242,23 @@ CudaExecutor::~CudaExecutor() { } } +Result> CudaExecutor::create( + const CudaExecutorConfig& config) { + ET_CHECK_OR_RETURN_ERROR( + !config.model_path.empty(), + InvalidArgument, + "CudaExecutor: no model path"); + auto module = std::make_unique( + 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> CudaExecutor::create( std::unique_ptr module, int max_sessions, @@ -248,11 +266,34 @@ Result> CudaExecutor::create( 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> CudaExecutor::create( + std::unique_ptr 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::max() / max_session_tokens, + max_sessions > 0 && max_session_tokens >= 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()); @@ -279,6 +320,19 @@ Result> 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::max(), + InvalidProgram, + "CudaExecutor: cell count must fit an int and leave a padding row"); + if (max_session_tokens == 0) { + max_session_tokens = static_cast( + 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 methods; diff --git a/backends/cuda/batching/cuda_executor.h b/backends/cuda/batching/cuda_executor.h index 004b16ec580..b50585fdf80 100644 --- a/backends/cuda/batching/cuda_executor.h +++ b/backends/cuda/batching/cuda_executor.h @@ -66,6 +66,19 @@ struct CudaExecutorOptions { bool cuda_graph_for_static_methods = true; }; +struct CudaExecutorConfig { + std::string model_path; + std::vector 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(::executorch::aten::ScalarType::BFloat16); + int initial_capacity = -1; + CudaExecutorOptions options; +}; + // How often each method ran, for reporting. struct ForwardMethodCalls { std::string name; @@ -76,6 +89,17 @@ class ET_EXPERIMENTAL CudaExecutor : public llm_batching::Executor { public: ~CudaExecutor() override; + static ::executorch::runtime::Result> 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 @@ -129,6 +153,10 @@ class ET_EXPERIMENTAL CudaExecutor : public llm_batching::Executor { } private: + static ::executorch::runtime::Result> create( + std::unique_ptr<::executorch::extension::Module> module, + const CudaExecutorConfig& config); + CudaExecutor( std::unique_ptr<::executorch::extension::Module> module, std::shared_ptr cache, @@ -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 methods_; std::vector calls_; diff --git a/backends/cuda/batching/test/test_cuda_executor.cpp b/backends/cuda/batching/test/test_cuda_executor.cpp index d06aea4580c..6c66cd1de15 100644 --- a/backends/cuda/batching/test/test_cuda_executor.cpp +++ b/backends/cuda/batching/test/test_cuda_executor.cpp @@ -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> 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, diff --git a/examples/models/muse-glimmer/CMakeLists.txt b/examples/models/muse-glimmer/CMakeLists.txt index 75230372ef6..03b8f574b4d 100644 --- a/examples/models/muse-glimmer/CMakeLists.txt +++ b/examples/models/muse-glimmer/CMakeLists.txt @@ -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) @@ -226,8 +229,19 @@ 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 @@ -235,7 +249,7 @@ if(MUSE_GLIMMER_BUILD_BATCHING OR MUSE_GLIMMER_BUILD_BATCHING_TESTS) extension_module extension_tensor extension_llm_serving - mlx + ${_batching_backend_library} ) foreach(_dependency IN ITEMS ${_batching_dependencies} extension_llm_batching_executor_utils @@ -243,13 +257,13 @@ if(MUSE_GLIMMER_BUILD_BATCHING OR MUSE_GLIMMER_BUILD_BATCHING_TESTS) 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( @@ -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) diff --git a/examples/models/muse-glimmer/runtime/backend_factory.cpp b/examples/models/muse-glimmer/runtime/backend_factory.cpp index a269ff4f78a..cb2aaa17713 100644 --- a/examples/models/muse-glimmer/runtime/backend_factory.cpp +++ b/examples/models/muse-glimmer/runtime/backend_factory.cpp @@ -6,14 +6,68 @@ * LICENSE file in the root directory of this source tree. */ #include +#if defined(EXECUTORCH_BUILD_MLX) #include +#endif +#if defined(EXECUTORCH_BUILD_CUDA) +#include +#endif + +#include +#include namespace executorch::extension::llm { runtime::Result create_muse_glimmer_backend( const MuseGlimmerBackendConfig& config) { - if (config.backend != "mlx") { - return runtime::Error::NotSupported; +#if defined(EXECUTORCH_BUILD_MLX) + if (config.backend == "mlx") { + return MuseGlimmerMLXExecutor::create(config); + } +#endif +#if defined(EXECUTORCH_BUILD_CUDA) + if (config.backend == "cuda") { + namespace cb = ::executorch::backends::cuda::batching; + ET_CHECK_OR_RETURN_ERROR( + config.max_session_tokens != 1, + InvalidArgument, + "Muse Glimmer CUDA: a session needs room for a prompt and prediction"); + ET_CHECK_OR_RETURN_ERROR( + !config.use_dflash && config.dflash_n_draft == 0, + NotSupported, + "Muse Glimmer CUDA: DFlash batching is not supported"); + cb::CudaExecutorConfig executor_config; + executor_config.model_path = config.model_path; + if (!config.data_path.empty()) { + executor_config.data_files.push_back(config.data_path); + } + executor_config.max_sessions = config.max_sessions; + executor_config.max_session_tokens = config.max_session_tokens; + ET_ASSIGN_OR_RETURN(executor, cb::CudaExecutor::create(executor_config)); + const auto vocab = executor->vocab_size(); + const auto session_tokens = executor->max_session_tokens(); + ET_CHECK_OR_RETURN_ERROR( + vocab > 200092 && config.bos_id < static_cast(vocab), + InvalidProgram, + "Muse Glimmer CUDA: vocabulary %d cannot hold its special tokens", + vocab); + ET_CHECK_OR_RETURN_ERROR( + session_tokens > 1, + InvalidArgument, + "Muse Glimmer CUDA: sessions need room for a prompt and prediction"); + auto spec = std::make_shared( + MuseGlimmerPreparationSpec{ + aten::ScalarType::BFloat16, + /*hidden_dim=*/0, + session_tokens, + static_cast(executor->preferred_batch_tokens()), + vocab, + /*has_vision=*/false, + /*max_soft_tokens=*/0, + config.image_limits, + config.bos_id}); + return MuseGlimmerBackend{std::move(executor), std::move(spec)}; } - return MuseGlimmerMLXExecutor::create(config); +#endif + return runtime::Error::NotSupported; } } // namespace executorch::extension::llm diff --git a/examples/models/muse-glimmer/runtime/runners/batching_bootstrap.cpp b/examples/models/muse-glimmer/runtime/runners/batching_bootstrap.cpp index 128b3a67929..b3a9bc32c0f 100644 --- a/examples/models/muse-glimmer/runtime/runners/batching_bootstrap.cpp +++ b/examples/models/muse-glimmer/runtime/runners/batching_bootstrap.cpp @@ -20,9 +20,16 @@ DEFINE_int32(dflash_n_draft, 0, "DFlash proposals; zero uses exported maximum"); DEFINE_string(tokenizer, "", "Matching Hugging Face tokenizer JSON"); DEFINE_string(data_path, "", "Optional external weights file"); DEFINE_string(pos_embed_path, "", "Optional vision positional table path"); +#if defined(EXECUTORCH_BUILD_CUDA) +DEFINE_string(backend, "cuda", "Executor backend"); +#else DEFINE_string(backend, "mlx", "Executor backend"); +#endif DEFINE_int32(max_sessions, 4, "Maximum logical sessions"); -DEFINE_int32(max_session_tokens, 0, "Context cap; zero uses model metadata"); +DEFINE_int32( + max_session_tokens, + 0, + "Context cap; zero uses model metadata (CUDA: the cell pool shared evenly)"); DEFINE_int32(max_decode_sequences, 4, "Maximum decodes in a batch"); DEFINE_int32(max_inflight_requests, 4, "Maximum admitted requests"); DEFINE_int32(