From f259b1e7a98d9a58c78023631509b599b0a1ad82 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Mon, 12 Oct 2026 22:21:50 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../models/muse-glimmer/BATCHING_SERVING.md | 49 +++++++++++++++++-- .../muse-glimmer/serving/serve_batching.py | 13 +++-- .../muse-glimmer/tests/test_serve_batching.py | 34 ++++++++++++- 3 files changed, 89 insertions(+), 7 deletions(-) diff --git a/examples/models/muse-glimmer/BATCHING_SERVING.md b/examples/models/muse-glimmer/BATCHING_SERVING.md index 68a0676bf42..e3fdc291faa 100644 --- a/examples/models/muse-glimmer/BATCHING_SERVING.md +++ b/examples/models/muse-glimmer/BATCHING_SERVING.md @@ -1,9 +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. -Batching CUDA and DFlash execution are not supported. +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 execution is not +supported. CUDA specifics are in [CUDA](#cuda) below. ## Export @@ -173,3 +174,45 @@ 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 runs `CudaExecutor` (backends/cuda/batching) 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, so their product must fit 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. diff --git a/examples/models/muse-glimmer/serving/serve_batching.py b/examples/models/muse-glimmer/serving/serve_batching.py index 94ab085b017..aba4e4315d1 100644 --- a/examples/models/muse-glimmer/serving/serve_batching.py +++ b/examples/models/muse-glimmer/serving/serve_batching.py @@ -6,8 +6,10 @@ """Serve Muse Glimmer solo text/images with the batching multiplexed worker. -Requires an off-graph solo export. The legacy launcher and runners remain -separate; this entry point does not select or fall back to DFlash. +Requires an off-graph batching export: export_solo.py --use-offgraph-kv-cache +for MLX (text and images), export_solo_batching.py for CUDA (text only). The +legacy launcher and runners remain separate; this entry point does not select +or fall back to DFlash. """ 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("--model-id", default="muse_glimmer") parser.add_argument("--max-context", type=int, required=True) parser.add_argument("--max-sessions", type=int, default=4) @@ -86,6 +88,11 @@ 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") + 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 76d28175a22..a1037f6c9cc 100644 --- a/examples/models/muse-glimmer/tests/test_serve_batching.py +++ b/examples/models/muse-glimmer/tests/test_serve_batching.py @@ -63,7 +63,7 @@ def test_worker_forwards_mg_options(arguments): @pytest.mark.parametrize( "flag,value", [ - ("--backend", "cuda"), + ("--backend", "vulkan"), ("--max-vision-patches", "3"), ("--max-vision-patches", str(1 << 31)), ], @@ -74,6 +74,38 @@ 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(zip(command[1::2], command[2::2])) + assert flags["--backend"] == "cuda" + assert flags["--data_path"] == str(data) + assert "--pos_embed_path" not in flags + + +@pytest.mark.parametrize("pos_embed", [False, True]) +def test_cuda_launcher_rejects_missing_weights_or_vision( + arguments, tmp_path, pos_embed +): + data = tmp_path / "aoti_cuda_blob.ptd" + table = tmp_path / "pos_embed.bin" + data.touch() + table.touch() + argv = arguments + ["--backend", "cuda"] + argv += ( + ["--data-path", str(data), "--pos-embed-path", str(table)] + if pos_embed + else [] + ) + 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.