From f5119a958b6311cbedc592bad577fab763d52228 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Sun, 18 Oct 2026 16:46:19 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../models/muse-glimmer/BATCHING_SERVING.md | 16 ++-- .../muse-glimmer/runtime/backend_factory.cpp | 13 +-- .../runtime/cuda_dflash_executor.cpp | 93 ++++++++++++++++--- .../runtime/cuda_dflash_executor.h | 23 ++++- .../muse-glimmer/serving/serve_batching.py | 8 +- .../muse-glimmer/tests/test_serve_batching.py | 26 +++--- 6 files changed, 128 insertions(+), 51 deletions(-) diff --git a/examples/models/muse-glimmer/BATCHING_SERVING.md b/examples/models/muse-glimmer/BATCHING_SERVING.md index 4bed79eb1ee..fdc2b27974c 100644 --- a/examples/models/muse-glimmer/BATCHING_SERVING.md +++ b/examples/models/muse-glimmer/BATCHING_SERVING.md @@ -1,8 +1,8 @@ # Muse Glimmer Batching Serving -The batching runner and multiplexed worker support MLX autoregressive generation -with text only or text plus one image across the submitted conversation, and -CUDA autoregressive generation with text only. They are opt-in entry points; the +The batching runner and multiplexed worker support MLX and CUDA autoregressive +generation with text only or text plus one image across the submitted +conversation. They are opt-in entry points; the legacy runners and export defaults remain unchanged. DFlash batching runs on MLX and CUDA (text only). CUDA specifics are in [CUDA](#cuda) below. @@ -205,8 +205,10 @@ python -m executorch.examples.models.muse_glimmer.export.export_solo \ `--forward-sizes` (default `1,2,4,8,16,32,64`), `--max-step-tokens` (default 4096) and `--max-cells` (default 65536) size the methods and the shared KV -pool; `--logits` samples on the host instead. The artifact is text only, so -`--mmproj` is rejected. +pool; `--logits` samples on the host instead. Add `--mmproj` for images: the +token forwards are unchanged, and the artifact gains `vision_encoder`, +`embed_text` and a dynamic `forward_from_embeddings` over the same cache, which +runs only the prefill steps that carry an image. Build with `make muse-glimmer-batching-cuda` on Linux, then serve: @@ -221,8 +223,8 @@ python -m executorch.examples.models.muse_glimmer.serving.serve_batching \ --max-context 8192 ``` -`--data-path` is required and `--pos-embed-path` is rejected; image requests -fail preparation. Every resident session reserves `--max-context` cells of the +`--data-path` is required. Image artifacts read `pos_embed.bin` beside +`model.pte`, or `--pos-embed-path`. Every resident session reserves `--max-context` cells of the pool: requests, `--prefix-cache-entries` snapshots, and one capture lane when snapshots are enabled. Their total reservation must leave one cell for padded writes within the exported `--max-cells`. diff --git a/examples/models/muse-glimmer/runtime/backend_factory.cpp b/examples/models/muse-glimmer/runtime/backend_factory.cpp index b833cd85d64..9ff3fba54a0 100644 --- a/examples/models/muse-glimmer/runtime/backend_factory.cpp +++ b/examples/models/muse-glimmer/runtime/backend_factory.cpp @@ -54,17 +54,8 @@ runtime::Result create_muse_glimmer_backend( 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}); + ET_ASSIGN_OR_RETURN( + spec, make_cuda_preparation(*executor, config, session_tokens)); return MuseGlimmerBackend{std::move(executor), std::move(spec)}; } #endif diff --git a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp index 55f8ee16a9d..f0dc4013bb8 100644 --- a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp +++ b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp @@ -7,6 +7,8 @@ */ #include +#include + #include #include #include @@ -65,7 +67,8 @@ float positional_uniform( mix(seed ^ mix(static_cast(purpose)) ^ mix(static_cast(position) + 1)); std::seed_seq words{ - static_cast(value), static_cast(value >> 32)}; + static_cast(value), + static_cast(value >> 32)}; std::mt19937 generator(words); return std::uniform_real_distribution(0.0f, 1.0f)(generator); } @@ -452,6 +455,72 @@ MuseGlimmerCudaDFlashExecutor::~MuseGlimmerCudaDFlashExecutor() { } } +Result> make_cuda_preparation( + cb::CudaExecutor& target, + const MuseGlimmerBackendConfig& config, + std::int32_t max_context, + std::shared_ptr* embeddings) { + ET_ASSIGN_OR_RETURN(names, target.module().method_names()); + const bool has_vision = names.count("vision_encoder") != 0; + const auto hidden = + static_cast(target.embeddings_hidden_size()); + std::int64_t max_soft_tokens = 0; + if (has_vision) { + ET_CHECK_OR_RETURN_ERROR( + hidden > 0, + InvalidProgram, + "Muse Glimmer CUDA: vision needs %s", + cb::kEmbeddingsMethod); + ET_ASSIGN_OR_RETURN( + soft_tokens, + muse_glimmer_abi::check_vision_abi( + target.module(), + hidden, + aten::ScalarType::BFloat16, + config.max_vision_patches, + ::executorch::backends::cuda::kCudaBackendId)); + max_soft_tokens = soft_tokens; + } + auto spec = std::make_shared( + MuseGlimmerPreparationSpec{ + aten::ScalarType::BFloat16, + has_vision ? hidden : 0, + max_context, + static_cast(target.preferred_batch_tokens()), + target.vocab_size(), + has_vision, + max_soft_tokens, + config.image_limits, + config.bos_id}); + if (!has_vision) { + return spec; + } + auto path = config.pos_embed_path; + if (path.empty()) { + const auto slash = config.model_path.find_last_of('/'); + path = (slash == std::string::npos ? "." + : config.model_path.substr(0, slash)) + + "/pos_embed.bin"; + } + auto model_embeddings = + std::make_shared(spec, std::move(path)); + Module* module = &target.module(); + ET_CHECK_OK_OR_RETURN_ERROR(target.set_input_embedder( + {[model_embeddings, module] { + return model_embeddings->initialize(*module); + }, + [model_embeddings](const batching::PreparedInput& input) { + return model_embeddings->accepts(input); + }, + [model_embeddings, module](const std::vector& inputs) { + return model_embeddings->materialize(*module, inputs); + }})); + if (embeddings) { + *embeddings = std::move(model_embeddings); + } + return spec; +} + Result MuseGlimmerCudaDFlashExecutor::create( const MuseGlimmerBackendConfig& config) { ET_CHECK_OR_RETURN_ERROR( @@ -670,8 +739,8 @@ bool MuseGlimmerCudaDFlashExecutor::initialize() { const auto meta = module.method_meta(name); const int width = n * block_; auto shaped = [&](std::size_t index, bool input, std::vector sizes) { - auto tensor = - input ? meta->input_tensor_meta(index) : meta->output_tensor_meta(index); + auto tensor = input ? meta->input_tensor_meta(index) + : meta->output_tensor_meta(index); return tensor.ok() && std::equal( tensor->sizes().begin(), @@ -900,9 +969,7 @@ std::vector MuseGlimmerCudaDFlashExecutor::verify_params( const auto& sampling = *state.sampling; const auto seed = state.seed.value_or(0); std::vector params{ - sampling.temperature, - static_cast(sampling.top_k), - sampling.top_p}; + sampling.temperature, static_cast(sampling.top_k), sampling.top_p}; for (int j = 0; j < proposals_; ++j) { params.push_back(df::positional_uniform(seed, first + j, 2)); } @@ -951,7 +1018,8 @@ bool MuseGlimmerCudaDFlashExecutor::verify( const auto& state = sessions_.at(owners[i]); greedy = greedy && state.sampling->temperature == 0.0f; const auto row = verify_params( - state, static_cast(step.positions[i * width] + 1)); + state, + static_cast(step.positions[i * width] + 1)); params.insert(params.end(), row.begin(), row.end()); } const auto stream = @@ -1056,10 +1124,11 @@ bool MuseGlimmerCudaDFlashExecutor::draft( // draft. Longer contexts drain through the dynamic one. const bool use_static = proposals && static_drafts_[owners.size()] && std::all_of(plan->counts.begin(), plan->counts.end(), [this](int n) { - return n <= block_; - }); + return n <= block_; + }); if (use_static && - !df::pad_to_static(*plan, static_cast(owners.size()), block_, row_bytes_)) { + !df::pad_to_static( + *plan, static_cast(owners.size()), block_, row_bytes_)) { return false; } const int padded = static_cast(plan->hidden.size() / row_bytes_); @@ -1077,8 +1146,8 @@ bool MuseGlimmerCudaDFlashExecutor::draft( const std::string method = use_static ? df::kStaticForwardPrefix + std::to_string(owners.size()) : std::string(df::kForward); - auto result = target_->module().execute( - method, {tokens, hidden, positions, order}); + auto result = + target_->module().execute(method, {tokens, hidden, positions, order}); if (!result.ok() || result->size() != 1 || !result->at(0).isTensor()) { return false; } diff --git a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h index 11dffd5f7fe..2a023aa57cd 100644 --- a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h +++ b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h @@ -15,6 +15,18 @@ namespace executorch::extension::llm { +class MuseGlimmerModelEmbeddings; + +// The preparation spec for prompts to `target`. When the program has +// vision_encoder, also installs an image embedder on it; `embeddings` then +// receives it, for adapters that build their own target steps. +runtime::Result> +make_cuda_preparation( + ::executorch::backends::cuda::batching::CudaExecutor& target, + const MuseGlimmerBackendConfig& config, + std::int32_t max_context, + std::shared_ptr* embeddings = nullptr); + // Exact packed-export metadata contract. Geometry tensors use the ordinary // get_n_caches/get_kv_heads/get_head_dims/get_windows names with get_draft_. namespace cuda_dflash_detail { @@ -108,7 +120,11 @@ std::optional pack_draft( // Widens a plan whose context fits one block per owner to the static // draft_forward_{N} shape: unused context rows go last, after every block, so // the cache writes them to its scratch cell and no query attends them. -bool pad_to_static(DraftPlan& plan, int owners, int block, std::size_t row_bytes); +bool pad_to_static( + DraftPlan& plan, + int owners, + int block, + std::size_t row_bytes); // Reusable device memory for small host-built inputs. class DeviceStaging { @@ -220,9 +236,8 @@ class MuseGlimmerCudaDFlashExecutor final : public batching::Executor { const std::vector& owners, const std::vector& candidates, std::vector& results); - std::vector verify_params( - const State& state, - batching::Position first) const; + std::vector verify_params(const State& state, batching::Position first) + const; bool drain(batching::SessionId owner); bool draft( const std::vector& owners, diff --git a/examples/models/muse-glimmer/serving/serve_batching.py b/examples/models/muse-glimmer/serving/serve_batching.py index 2d3ab3a8014..199ba281980 100644 --- a/examples/models/muse-glimmer/serving/serve_batching.py +++ b/examples/models/muse-glimmer/serving/serve_batching.py @@ -7,9 +7,9 @@ """Serve Muse Glimmer text/images with the batching multiplexed worker. Requires an off-graph batching export: export_solo.py --use-offgraph-kv-cache -for MLX (text and images); add --backend cuda for CUDA (text only). DFlash is -opt-in, MLX only, and requires a matching export; the legacy launcher and -runners remain separate. +for MLX; add --backend cuda for CUDA. Both serve text and images when exported +with --mmproj. DFlash is opt-in and requires a matching export; the legacy +launcher and runners remain separate. """ import argparse @@ -109,8 +109,6 @@ def _parse_args(argv=None): if args.backend == "cuda": if args.data_path is None: parser.error("--backend cuda requires --data-path (aoti_cuda_blob.ptd)") - if args.pos_embed_path is not None: - parser.error("--backend cuda serves text only; omit --pos-embed-path") _validate_paths(parser, args) return args diff --git a/examples/models/muse-glimmer/tests/test_serve_batching.py b/examples/models/muse-glimmer/tests/test_serve_batching.py index 080a6ef6e1b..e5bba6a2ef4 100644 --- a/examples/models/muse-glimmer/tests/test_serve_batching.py +++ b/examples/models/muse-glimmer/tests/test_serve_batching.py @@ -139,22 +139,24 @@ def test_cuda_launcher_forwards_dflash(arguments, tmp_path): ) -@pytest.mark.parametrize("case", ["no_weights", "pos_embed"]) -def test_cuda_launcher_rejects_missing_weights_or_vision( - arguments, tmp_path, case -): +def test_cuda_launcher_requires_the_weights_file(arguments): + with pytest.raises(SystemExit) as error: + serve_batching._parse_args(arguments + ["--backend", "cuda"]) + assert error.value.code == 2 + + +def test_cuda_worker_takes_the_vision_table(arguments, tmp_path): data = tmp_path / "aoti_cuda_blob.ptd" table = tmp_path / "pos_embed.bin" data.touch() table.touch() - argv = arguments + ["--backend", "cuda"] - if case != "no_weights": - argv += ["--data-path", str(data)] - if case == "pos_embed": - argv += ["--pos-embed-path", str(table)] - with pytest.raises(SystemExit) as error: - serve_batching._parse_args(argv) - assert error.value.code == 2 + args = serve_batching._parse_args( + arguments + + ["--backend", "cuda", "--data-path", str(data)] + + ["--pos-embed-path", str(table)] + ) + flags = dict(arg.split("=", 1) for arg in serve_batching._worker_command(args)[1:]) + assert flags["--pos_embed_path"] == str(table) @pytest.mark.parametrize("shortfall", [0, 1])