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(