diff --git a/examples/models/muse-glimmer/BATCHING_SERVING.md b/examples/models/muse-glimmer/BATCHING_SERVING.md index abb7702d525..1c1aaca46b0 100644 --- a/examples/models/muse-glimmer/BATCHING_SERVING.md +++ b/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. diff --git a/examples/models/muse-glimmer/serving/serve_batching.py b/examples/models/muse-glimmer/serving/serve_batching.py index 7c699244f18..9644c3bff73 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 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: diff --git a/examples/models/muse-glimmer/tests/test_serve_batching.py b/examples/models/muse-glimmer/tests/test_serve_batching.py index 351513a3664..5fc51108416 100644 --- a/examples/models/muse-glimmer/tests/test_serve_batching.py +++ b/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.