Skip to content

Commit dfdd5db

Browse files
authored
[executorch][muse] Serve MG batching on CUDA from the launcher (#24009)
This PR was created by the merge bot to help merge the original PR into the main branch. ghstack PR number: #23733 by @Gasoonjia ^ Please use this as the source of truth for the PR details, comments, and reviews ghstack PR base: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/49/base ghstack PR head: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/49/head Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/48/orig Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/49/orig Differential Revision: [D124871961](https://our.internmc.facebook.com/intern/diff/D124871961/) @diff-train-skip-merge
1 parent 7ed9cd8 commit dfdd5db

3 files changed

Lines changed: 98 additions & 6 deletions

File tree

‎examples/models/muse-glimmer/BATCHING_SERVING.md‎

Lines changed: 47 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
# Muse Glimmer Batching Serving
22

33
The batching runner and multiplexed worker support MLX autoregressive generation
4-
with text only or text plus one image across the submitted conversation. They are
5-
opt-in entry points; the legacy runners and export defaults remain unchanged.
4+
with text only or text plus one image across the submitted conversation, and
5+
CUDA autoregressive generation with text only. They are opt-in entry points; the
6+
legacy runners and export defaults remain unchanged. DFlash batching is MLX only.
7+
CUDA specifics are in [CUDA](#cuda) below.
68

79
## Export
810

@@ -182,3 +184,46 @@ prepared requests retain independent backing. If the uncached suffix skips the
182184
image entirely, it needs no vision encoding; if it intersects the image, that
183185
backing may still encode the full image before selecting the remaining rows.
184186
Prefix identity does not retain RGB or embedding buffers.
187+
188+
## CUDA
189+
190+
The CUDA worker builds `CudaExecutor` through `CudaExecutorConfig` under the same
191+
serving runtime, prefix cache, scheduler and launcher. Export with the off-graph
192+
cache, which on CUDA writes the batching methods: static `forward_{N}` (each
193+
captured into a CUDA graph), a dynamic `forward_others` for wider steps, and
194+
device samplers. Without `--use-offgraph-kv-cache` CUDA keeps its in-graph
195+
prefill and decode export.
196+
197+
```sh
198+
python -m executorch.examples.models.muse_glimmer.export.export_solo \
199+
--backend cuda \
200+
--gguf /path/to/models/mg/model.gguf \
201+
--use-offgraph-kv-cache \
202+
--output-dir /path/to/models/mg-batching-cuda
203+
```
204+
205+
`--forward-sizes` (default `1,2,4,8,16,32,64`), `--max-step-tokens` (default
206+
4096) and `--max-cells` (default 65536) size the methods and the shared KV
207+
pool; `--logits` samples on the host instead. The artifact is text only, so
208+
`--mmproj` is rejected.
209+
210+
Build with `make muse-glimmer-batching-cuda` on Linux, then serve:
211+
212+
```sh
213+
python -m executorch.examples.models.muse_glimmer.serving.serve_batching \
214+
--backend cuda \
215+
--worker-bin /path/to/executorch/cmake-out/examples/models/muse-glimmer-batching/muse_glimmer_batching_worker \
216+
--model-path /path/to/models/mg-batching-cuda/model.pte \
217+
--data-path /path/to/models/mg-batching-cuda/aoti_cuda_blob.ptd \
218+
--tokenizer-path /path/to/models/mg/tokenizer.json \
219+
--hf-tokenizer /path/to/models/mg \
220+
--max-context 8192
221+
```
222+
223+
`--data-path` is required and `--pos-embed-path` is rejected; image requests
224+
fail preparation. Every resident session reserves `--max-context` cells of the
225+
pool: requests, `--prefix-cache-entries` snapshots, and one capture lane when
226+
snapshots are enabled. Their total reservation must leave one cell for padded
227+
writes within the exported `--max-cells`.
228+
The cell pool tracks at most 64 resident sessions. Prefix snapshots clone a
229+
session's cells by reference, sharing the prompt's KV rather than copying it.

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

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,10 @@
66

77
"""Serve Muse Glimmer text/images with the batching multiplexed worker.
88
9-
Requires an off-graph export. DFlash is opt-in and requires a matching export;
10-
the legacy launcher and runners remain separate.
9+
Requires an off-graph batching export: export_solo.py --use-offgraph-kv-cache
10+
for MLX (text and images); add --backend cuda for CUDA (text only). DFlash is
11+
opt-in, MLX only, and requires a matching export; the legacy launcher and
12+
runners remain separate.
1113
"""
1214

1315
import argparse
@@ -45,7 +47,7 @@ def _parse_args(argv=None):
4547
parser.add_argument("--tokenizer-path", required=True)
4648
parser.add_argument("--hf-tokenizer", required=True)
4749
parser.add_argument("--pos-embed-path")
48-
parser.add_argument("--backend", choices=("mlx",), default="mlx")
50+
parser.add_argument("--backend", choices=("mlx", "cuda"), default="mlx")
4951
parser.add_argument("--dflash", action="store_true")
5052
parser.add_argument(
5153
"--dflash-n-draft",
@@ -104,10 +106,22 @@ def _parse_args(argv=None):
104106
parser.error("BOS and EOS IDs must be uint64 values")
105107
if 4 * ((args.max_image_bytes + 2) // 3) + 1024 > args.max_request_bytes:
106108
parser.error("--max-request-bytes must fit the base64 image plus JSON framing")
109+
_validate_backend(parser, args)
107110
_validate_paths(parser, args)
108111
return args
109112

110113

114+
def _validate_backend(parser, args):
115+
if args.backend != "cuda":
116+
return
117+
if args.data_path is None:
118+
parser.error("--backend cuda requires --data-path (aoti_cuda_blob.ptd)")
119+
if args.pos_embed_path is not None:
120+
parser.error("--backend cuda serves text only; omit --pos-embed-path")
121+
if args.dflash:
122+
parser.error("--dflash is not supported with --backend cuda")
123+
124+
111125
def _validate_paths(parser, args):
112126
args.worker_bin = os.path.expanduser(args.worker_bin)
113127
if shutil.which(args.worker_bin) is None:

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

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,7 @@ def test_draft_count_forwarded(arguments, n_draft):
9494
@pytest.mark.parametrize(
9595
"flag,value",
9696
[
97-
("--backend", "cuda"),
97+
("--backend", "vulkan"),
9898
("--dflash-n-draft", "-1"),
9999
("--dflash-n-draft", str((1 << 31) - 1)),
100100
("--dflash-n-draft", "2"),
@@ -110,6 +110,39 @@ def test_launcher_rejects_invalid_options(arguments, flag, value):
110110
assert error.value.code == 2
111111

112112

113+
def test_cuda_worker_takes_the_weights_file(arguments, tmp_path):
114+
data = tmp_path / "aoti_cuda_blob.ptd"
115+
data.touch()
116+
args = serve_batching._parse_args(
117+
arguments + ["--backend", "cuda", "--data-path", str(data)]
118+
)
119+
command = serve_batching._worker_command(args)
120+
flags = dict(arg.split("=", 1) for arg in command[1:])
121+
assert flags["--backend"] == "cuda"
122+
assert flags["--data_path"] == str(data)
123+
assert "--pos_embed_path" not in flags
124+
125+
126+
@pytest.mark.parametrize("case", ["no_weights", "pos_embed", "dflash"])
127+
def test_cuda_launcher_rejects_missing_weights_vision_or_dflash(
128+
arguments, tmp_path, case
129+
):
130+
data = tmp_path / "aoti_cuda_blob.ptd"
131+
table = tmp_path / "pos_embed.bin"
132+
data.touch()
133+
table.touch()
134+
argv = arguments + ["--backend", "cuda"]
135+
if case != "no_weights":
136+
argv += ["--data-path", str(data)]
137+
if case == "pos_embed":
138+
argv += ["--pos-embed-path", str(table)]
139+
if case == "dflash":
140+
argv += ["--dflash"]
141+
with pytest.raises(SystemExit) as error:
142+
serve_batching._parse_args(argv)
143+
assert error.value.code == 2
144+
145+
113146
@pytest.mark.parametrize("shortfall", [0, 1])
114147
def test_image_frame_capacity(arguments, shortfall):
115148
# Four image bytes need eight base64 bytes plus 1024 bytes of JSON framing.

0 commit comments

Comments
 (0)