[executorch][muse] Serve MG batching on CUDA from the launcher

examples/models/muse-glimmer/BATCHING_SERVING.md

@@ -1,8 +1,10 @@
 # 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. They are
-opt-in entry points; the legacy runners and export defaults remain unchanged.
+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
+legacy runners and export defaults remain unchanged. DFlash batching is MLX only.
+CUDA specifics are in [CUDA](#cuda) below.
 
 ## Export
 
@@ -182,3 +184,46 @@ prepared requests retain independent backing. If the uncached suffix skips the
 image entirely, it needs no vision encoding; if it intersects the image, that
 backing may still encode the full image before selecting the remaining rows.
 Prefix identity does not retain RGB or embedding buffers.
+
+## CUDA
+
+The CUDA worker builds `CudaExecutor` through `CudaExecutorConfig` under the same
+serving runtime, prefix cache, scheduler and launcher. Export with the off-graph
+cache, which on CUDA writes the batching methods: static `forward_{N}` (each
+captured into a CUDA graph), a dynamic `forward_others` for wider steps, and
+device samplers. Without `--use-offgraph-kv-cache` CUDA keeps its in-graph
+prefill and decode export.
+
+```sh
+python -m executorch.examples.models.muse_glimmer.export.export_solo \
+  --backend cuda \
+  --gguf /path/to/models/mg/model.gguf \
+  --use-offgraph-kv-cache \
+  --output-dir /path/to/models/mg-batching-cuda
+```
+
+`--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.
+
+Build with `make muse-glimmer-batching-cuda` on Linux, then serve:
+
+```sh
+python -m executorch.examples.models.muse_glimmer.serving.serve_batching \
+  --backend cuda \
+  --worker-bin /path/to/executorch/cmake-out/examples/models/muse-glimmer-batching/muse_glimmer_batching_worker \
+  --model-path /path/to/models/mg-batching-cuda/model.pte \
+  --data-path /path/to/models/mg-batching-cuda/aoti_cuda_blob.ptd \
+  --tokenizer-path /path/to/models/mg/tokenizer.json \
+  --hf-tokenizer /path/to/models/mg \
+  --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
+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`.
+The cell pool tracks at most 64 resident sessions. Prefix snapshots clone a
+session's cells by reference, sharing the prompt's KV rather than copying it.

examples/models/muse-glimmer/serving/serve_batching.py

@@ -6,8 +6,10 @@
 
 """Serve Muse Glimmer text/images with the batching multiplexed worker.
 
-Requires an off-graph export. DFlash is opt-in and requires a matching export;
-the legacy launcher and runners remain separate.
+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.
 """
 
 import argparse
@@ -45,7 +47,7 @@ def _parse_args(argv=None):
     parser.add_argument("--tokenizer-path", required=True)
     parser.add_argument("--hf-tokenizer", required=True)
     parser.add_argument("--pos-embed-path")
-    parser.add_argument("--backend", choices=("mlx",), default="mlx")
+    parser.add_argument("--backend", choices=("mlx", "cuda"), default="mlx")
     parser.add_argument("--dflash", action="store_true")
     parser.add_argument(
         "--dflash-n-draft",
@@ -104,10 +106,22 @@ def _parse_args(argv=None):
         parser.error("BOS and EOS IDs must be uint64 values")
     if 4 * ((args.max_image_bytes + 2) // 3) + 1024 > args.max_request_bytes:
         parser.error("--max-request-bytes must fit the base64 image plus JSON framing")
+    _validate_backend(parser, args)
     _validate_paths(parser, args)
     return args
 
 
+def _validate_backend(parser, args):
+    if args.backend != "cuda":
+        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")
+    if args.dflash:
+        parser.error("--dflash is not supported with --backend cuda")
+
+
 def _validate_paths(parser, args):
     args.worker_bin = os.path.expanduser(args.worker_bin)
     if shutil.which(args.worker_bin) is None:

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

@@ -94,7 +94,7 @@ def test_draft_count_forwarded(arguments, n_draft):
 @pytest.mark.parametrize(
     "flag,value",
     [
-        ("--backend", "cuda"),
+        ("--backend", "vulkan"),
         ("--dflash-n-draft", "-1"),
         ("--dflash-n-draft", str((1 << 31) - 1)),
         ("--dflash-n-draft", "2"),
@@ -110,6 +110,39 @@ def test_launcher_rejects_invalid_options(arguments, flag, value):
     assert error.value.code == 2
 
 
+def test_cuda_worker_takes_the_weights_file(arguments, tmp_path):
+    data = tmp_path / "aoti_cuda_blob.ptd"
+    data.touch()
+    args = serve_batching._parse_args(
+        arguments + ["--backend", "cuda", "--data-path", str(data)]
+    )
+    command = serve_batching._worker_command(args)
+    flags = dict(arg.split("=", 1) for arg in command[1:])
+    assert flags["--backend"] == "cuda"
+    assert flags["--data_path"] == str(data)
+    assert "--pos_embed_path" not in flags
+
+
+@pytest.mark.parametrize("case", ["no_weights", "pos_embed", "dflash"])
+def test_cuda_launcher_rejects_missing_weights_vision_or_dflash(
+    arguments, tmp_path, case
+):
+    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)]
+    if case == "dflash":
+        argv += ["--dflash"]
+    with pytest.raises(SystemExit) as error:
+        serve_batching._parse_args(argv)
+    assert error.value.code == 2
+
+
 @pytest.mark.parametrize("shortfall", [0, 1])
 def test_image_frame_capacity(arguments, shortfall):
     # Four image bytes need eight base64 bytes plus 1024 bytes of JSON framing.