From 773bb35ec9ad769670a7df1232b1713b346612e7 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Mon, 12 Oct 2026 22:21:40 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- Makefile | 18 ++++ examples/models/muse-glimmer/CMakeLists.txt | 50 +++++++-- .../muse-glimmer/runtime/backend_factory.cpp | 18 +++- .../runtime/cuda_backend_factory.cpp | 101 ++++++++++++++++++ .../runtime/cuda_backend_factory.h | 21 ++++ .../runtime/runners/batching_bootstrap.cpp | 9 +- 6 files changed, 202 insertions(+), 15 deletions(-) create mode 100644 examples/models/muse-glimmer/runtime/cuda_backend_factory.cpp create mode 100644 examples/models/muse-glimmer/runtime/cuda_backend_factory.h 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/examples/models/muse-glimmer/CMakeLists.txt b/examples/models/muse-glimmer/CMakeLists.txt index 75230372ef6..2f7493ff374 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,21 @@ 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 runtime/cuda_backend_factory.cpp) + 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 +251,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 +259,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 +307,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..740060d8ebd 100644 --- a/examples/models/muse-glimmer/runtime/backend_factory.cpp +++ b/examples/models/muse-glimmer/runtime/backend_factory.cpp @@ -6,14 +6,26 @@ * 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 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); } - return MuseGlimmerMLXExecutor::create(config); +#endif +#if defined(EXECUTORCH_BUILD_CUDA) + if (config.backend == "cuda") { + return create_muse_glimmer_cuda_backend(config); + } +#endif + return runtime::Error::NotSupported; } } // namespace executorch::extension::llm diff --git a/examples/models/muse-glimmer/runtime/cuda_backend_factory.cpp b/examples/models/muse-glimmer/runtime/cuda_backend_factory.cpp new file mode 100644 index 00000000000..7f6fc3dd3e6 --- /dev/null +++ b/examples/models/muse-glimmer/runtime/cuda_backend_factory.cpp @@ -0,0 +1,101 @@ +/* + * 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 + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace executorch::extension::llm { +namespace { +namespace cuda_batching = ::executorch::backends::cuda::batching; +using runtime::Error; + +constexpr int kKVDtype = static_cast(aten::ScalarType::BFloat16); + +// One cell stays reserved as the scratch row that padded static steps write. +int64_t shared_session_tokens(int64_t context, int64_t max_cells, int sessions) { + return std::min(context, (max_cells - 1) / sessions); +} +} // namespace + +runtime::Result create_muse_glimmer_cuda_backend( + const MuseGlimmerBackendConfig& config) { + ET_CHECK_OR_RETURN_ERROR( + config.backend == "cuda" && !config.model_path.empty() && + config.max_sessions > 0 && config.max_session_tokens >= 0 && + config.max_session_tokens != 1, + InvalidArgument, + "Muse Glimmer CUDA: invalid backend configuration"); + std::vector data_files; + if (!config.data_path.empty()) { + data_files.push_back(config.data_path); + } + auto module = std::make_unique( + config.model_path, + data_files, + Module::LoadMode::MmapUseMlockIgnoreErrors, + /*event_tracer=*/nullptr, + /*memory_allocator=*/nullptr, + /*temp_allocator=*/nullptr, + /*share_memory_arenas=*/false); + ET_CHECK_OK_OR_RETURN_ERROR(module->load()); + ET_ASSIGN_OR_RETURN(context, read_max_context_length(*module)); + ET_ASSIGN_OR_RETURN(vocab, read_vocab_size(*module)); + ET_ASSIGN_OR_RETURN( + max_cells, + detail::read_required_positive_int( + *module, cuda_batching::kMaxCellsMethod)); + ET_CHECK_OR_RETURN_ERROR( + vocab > 200092 && vocab <= std::numeric_limits::max() && + config.bos_id < static_cast(vocab), + InvalidProgram, + "Muse Glimmer CUDA: vocabulary %" PRId64 " cannot hold its special tokens", + vocab); + const int64_t session_tokens = config.max_session_tokens + ? config.max_session_tokens + : shared_session_tokens(context, max_cells, config.max_sessions); + ET_CHECK_OR_RETURN_ERROR( + session_tokens > 1 && session_tokens <= context, + InvalidArgument, + "Muse Glimmer CUDA: %" PRId64 " tokens per session for %d sessions does " + "not fit context %" PRId64 " and %" PRId64 " cells", + session_tokens, + config.max_sessions, + context, + max_cells); + + auto executor = cuda_batching::CudaExecutor::create( + std::move(module), + config.max_sessions, + static_cast(session_tokens), + kKVDtype); + ET_CHECK_OK_OR_RETURN_ERROR(executor.error()); + const auto width = (*executor)->preferred_batch_tokens(); + auto spec = std::make_shared( + MuseGlimmerPreparationSpec{ + aten::ScalarType::BFloat16, + /*hidden_dim=*/0, + static_cast(session_tokens), + static_cast(width), + static_cast(vocab), + /*has_vision=*/false, + /*max_soft_tokens=*/0, + config.image_limits, + config.bos_id}); + return MuseGlimmerBackend{std::move(*executor), std::move(spec)}; +} + +} // namespace executorch::extension::llm diff --git a/examples/models/muse-glimmer/runtime/cuda_backend_factory.h b/examples/models/muse-glimmer/runtime/cuda_backend_factory.h new file mode 100644 index 00000000000..07488e86421 --- /dev/null +++ b/examples/models/muse-glimmer/runtime/cuda_backend_factory.h @@ -0,0 +1,21 @@ +/* + * 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. + */ +#pragma once + +#include + +namespace executorch::extension::llm { + +// The CUDA batching backend: a CudaExecutor over the cell-layout artifact +// export_solo_batching.py writes. Text only; the artifact has no vision path. +// A zero max_session_tokens shares the cell pool evenly between max_sessions, +// capped by the exported context. +runtime::Result create_muse_glimmer_cuda_backend( + const MuseGlimmerBackendConfig& config); + +} // 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 035ce94b9f9..81ec2356ae9 100644 --- a/examples/models/muse-glimmer/runtime/runners/batching_bootstrap.cpp +++ b/examples/models/muse-glimmer/runtime/runners/batching_bootstrap.cpp @@ -18,9 +18,16 @@ DEFINE_string(pte, "", "Muse Glimmer solo off-graph program"); 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(