[executorch][muse] Serve images on CUDA batching

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`.

examples/models/muse-glimmer/runtime/backend_factory.cpp

@@ -58,17 +58,8 @@ runtime::Result<MuseGlimmerBackend> create_muse_glimmer_backend(
         session_tokens > 1,
         InvalidArgument,
         "Muse Glimmer CUDA: sessions need room for a prompt and prediction");
-    auto spec = std::make_shared<const MuseGlimmerPreparationSpec>(
-        MuseGlimmerPreparationSpec{
-            aten::ScalarType::BFloat16,
-            /*hidden_dim=*/0,
-            session_tokens,
-            static_cast<int32_t>(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

examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp

@@ -7,6 +7,8 @@
  */
 #include <executorch/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h>
 
+#include <executorch/examples/models/muse-glimmer/runtime/model_embeddings.h>
+
 #include <cuda_runtime.h>
 #include <executorch/backends/cuda/runtime/backend_options.h>
 #include <executorch/extension/cuda/caller_stream.h>
@@ -462,6 +464,77 @@ MuseGlimmerCudaDFlashExecutor::~MuseGlimmerCudaDFlashExecutor() {
   }
 }
 
+Result<std::shared_ptr<const MuseGlimmerPreparationSpec>> make_cuda_preparation(
+    cb::CudaExecutor& target,
+    const MuseGlimmerBackendConfig& config,
+    std::int32_t max_context,
+    std::shared_ptr<MuseGlimmerModelEmbeddings>* embeddings) {
+  ET_ASSIGN_OR_RETURN(names, target.module().method_names());
+  const bool has_vision = names.count("vision_encoder") != 0;
+  const auto hidden =
+      static_cast<std::int32_t>(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<const MuseGlimmerPreparationSpec>(
+      MuseGlimmerPreparationSpec{
+          aten::ScalarType::BFloat16,
+          has_vision ? hidden : 0,
+          max_context,
+          static_cast<std::int32_t>(target.preferred_batch_tokens()),
+          target.vocab_size(),
+          has_vision,
+          max_soft_tokens,
+          config.image_limits,
+          config.bos_id});
+  if (!has_vision) {
+    ET_CHECK_OR_RETURN_ERROR(
+        config.pos_embed_path.empty(),
+        InvalidArgument,
+        "Muse Glimmer CUDA: the program has no vision_encoder, but --pos-embed-path was given (%s)",
+        config.pos_embed_path.c_str());
+    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<MuseGlimmerModelEmbeddings>(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<batching::Input>& inputs) {
+         return model_embeddings->materialize(*module, inputs);
+       }}));
+  if (embeddings) {
+    *embeddings = std::move(model_embeddings);
+  }
+  return spec;
+}
+
 Result<MuseGlimmerBackend> MuseGlimmerCudaDFlashExecutor::create(
     const MuseGlimmerBackendConfig& config) {
   ET_CHECK_OR_RETURN_ERROR(

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<std::shared_ptr<const MuseGlimmerPreparationSpec>>
+make_cuda_preparation(
+    ::executorch::backends::cuda::batching::CudaExecutor& target,
+    const MuseGlimmerBackendConfig& config,
+    std::int32_t max_context,
+    std::shared_ptr<MuseGlimmerModelEmbeddings>* 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 {

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
@@ -116,8 +116,6 @@ def _validate_backend(parser, args):
         return
     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")
 
 
 def _validate_paths(parser, args):

examples/models/muse-glimmer/tests/test_serve_batching.py

@@ -139,20 +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])