diff --git a/examples/models/muse-glimmer/BATCHING_SERVING.md b/examples/models/muse-glimmer/BATCHING_SERVING.md index e20422bba73..b045a3f56b2 100644 --- a/examples/models/muse-glimmer/BATCHING_SERVING.md +++ b/examples/models/muse-glimmer/BATCHING_SERVING.md @@ -240,3 +240,42 @@ and `--max-draft-owners`), then add `--dflash` to the serve command above. The draft ABI and `--dflash-n-draft` behave as on MLX. Drafts are greedy; the target accepts or resamples them on the device under each request's sampling. Decode drafts of up to four owners run in CUDA graphs. + +## Native + +The native worker serves a `.ptn` on the native backend's ggml engine (Metal, +CUDA, or CPU, whichever ggml variant ExecuTorch links) under the same serving +runtime, prefix cache, scheduler and launcher. Export the Hugging Face model +from its GGUF with the off-graph cache and selected logits; the package +carries the sliding windows and context length: + +```sh +python -m executorch.backends.native.examples.hf_llm.export_hf_llm \ + --model /path/to/models/mg-hf \ + --gguf /path/to/models/mg/model.gguf \ + --offgraph_cache --logits_to_keep selected \ + --output /path/to/models/mg-native/model.ptn +``` + +`--model` names the Hugging Face `meta-models/Muse-Glimmer-30B` directory for +its config and tokenizer, which needs a transformers release with Muse Glimmer. +Build with `make muse-glimmer-batching-native`, then serve: + +```sh +python -m executorch.examples.models.muse_glimmer.serving.serve_batching \ + --backend native \ + --worker-bin /path/to/executorch/cmake-out/examples/models/muse-glimmer-batching/muse_glimmer_batching_worker \ + --model-path /path/to/models/mg-native/model.ptn \ + --tokenizer-path /path/to/models/mg-hf/tokenizer.json \ + --hf-tokenizer /path/to/models/mg-hf \ + --max-context 8192 +``` + +Sessions are sequences of the engine's cell cache: a batch puts several in one +forward, and prefix snapshots clone a session's cells rather than copy them. +The pool holds `--max-context` cells for every resident session (requests, +`--prefix-cache-entries` snapshots and capture slots), and its storage grows +with the cells in use. The native worker serves text only: `--dflash` and +`--pos-embed-path` are rejected. `--native-device` picks `gpu` (default) or +`cpu`, and `--native-kv-dtype` stores the cache as `fp32` (default) or +`fp16`. diff --git a/examples/models/muse-glimmer/serving/serve_batching.py b/examples/models/muse-glimmer/serving/serve_batching.py index e5d13f10461..44a856a11f5 100644 --- a/examples/models/muse-glimmer/serving/serve_batching.py +++ b/examples/models/muse-glimmer/serving/serve_batching.py @@ -7,7 +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 +for MLX (text and images); add --backend cuda for CUDA (text only). --backend +native serves a .ptn from backends/native's export_hf_llm.py --offgraph_cache +--logits_to_keep selected on the ggml engine (text only). DFlash is opt-in, MLX only, and requires a matching export; the legacy launcher and runners remain separate. """ @@ -47,7 +49,19 @@ 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", "cuda"), default="mlx") + parser.add_argument("--backend", choices=("mlx", "cuda", "native"), default="mlx") + parser.add_argument( + "--native-device", + choices=("gpu", "cpu"), + default="gpu", + help="--backend native: the ggml device.", + ) + parser.add_argument( + "--native-kv-dtype", + choices=("fp32", "fp16"), + default="fp32", + help="--backend native: the KV cache storage.", + ) parser.add_argument("--dflash", action="store_true") parser.add_argument( "--dflash-n-draft", @@ -118,6 +132,16 @@ def _parse_args(argv=None): def _validate_backend(parser, args): + if args.backend == "native": + if args.dflash: + parser.error("--backend native does not support --dflash") + if args.pos_embed_path is not None: + parser.error("--backend native serves text only; omit --pos-embed-path") + if args.data_path is not None: + parser.error( + "--backend native reads its weights from the .ptn; omit --data-path" + ) + return if args.backend != "cuda": return if args.data_path is None: @@ -176,6 +200,11 @@ def _worker_command(args): ("eos_id", args.eos_id), ("data_path", args.data_path), ("pos_embed_path", args.pos_embed_path), + ("native_device", args.native_device if args.backend == "native" else None), + ( + "native_kv_dtype", + args.native_kv_dtype if args.backend == "native" else None, + ), ): if value is not None: # gflags requires '=' to pass false for boolean flags. diff --git a/examples/models/muse-glimmer/tests/test_serve_batching.py b/examples/models/muse-glimmer/tests/test_serve_batching.py index 9b1c767c6de..5b450bbd056 100644 --- a/examples/models/muse-glimmer/tests/test_serve_batching.py +++ b/examples/models/muse-glimmer/tests/test_serve_batching.py @@ -166,6 +166,35 @@ def test_cuda_launcher_rejects_missing_weights_or_vision(arguments, tmp_path, ca assert error.value.code == 2 +def test_native_worker_takes_its_engine_options(arguments): + args = serve_batching._parse_args( + arguments + + ["--backend", "native", "--prefix-cache-entries", "2"] + + ["--native-device", "cpu"] + ) + flags = dict(arg.split("=", 1) for arg in serve_batching._worker_command(args)[1:]) + assert (flags["--backend"], flags["--prefix_cache_entries"]) == ("native", "2") + assert (flags["--native_device"], flags["--native_kv_dtype"]) == ("cpu", "fp32") + + +def test_other_backends_omit_native_options(arguments): + flags = serve_batching._worker_command(serve_batching._parse_args(arguments)) + assert not any(flag.startswith("--native_") for flag in flags) + + +@pytest.mark.parametrize( + "extra", + [["--dflash"], ["--pos-embed-path", "FILE"], ["--data-path", "FILE"]], +) +def test_native_launcher_rejects_dflash_vision_and_data(arguments, tmp_path, extra): + path = tmp_path / "file.bin" + path.touch() + extra = [str(path) if value == "FILE" else value for value in extra] + with pytest.raises(SystemExit) as error: + serve_batching._parse_args(arguments + ["--backend", "native"] + extra) + 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.