From 38245188eec32b5db07474116595676836be13c6 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Wed, 14 Oct 2026 09:11:38 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .ci/scripts/export_model_artifact.sh | 23 ++- .ci/scripts/test_model_e2e.sh | 150 +++++++++++++++++- .../tests/test_muse_glimmer_batching_ci.py | 62 ++++++++ .github/workflows/cuda.yml | 63 ++++++++ 4 files changed, 290 insertions(+), 8 deletions(-) create mode 100644 .ci/scripts/tests/test_muse_glimmer_batching_ci.py diff --git a/.ci/scripts/export_model_artifact.sh b/.ci/scripts/export_model_artifact.sh index 223dd82951c..79776c5a242 100755 --- a/.ci/scripts/export_model_artifact.sh +++ b/.ci/scripts/export_model_artifact.sh @@ -44,6 +44,7 @@ Arguments: - vr-streaming: Voxtral Realtime streaming mode - vr-offline: Voxtral Realtime offline mode - solo-text: Muse Glimmer solo text mode + - solo-text-batching: Muse Glimmer text on the CUDA batching runtime - dflash-image: Muse Glimmer DFlash vision mode Examples: @@ -95,7 +96,7 @@ if [ -n "$MODE" ]; then exit 1 fi ;; - solo-text|dflash-image) + solo-text|solo-text-batching|dflash-image) if [ "$HF_MODEL" != "meta-models/Muse-Glimmer-30B-GGUF" ]; then echo "Error: Mode '$MODE' can only be used with Muse Glimmer model" echo "Provided model: $HF_MODEL" @@ -104,7 +105,7 @@ if [ -n "$MODE" ]; then ;; *) echo "Error: Unsupported mode '$MODE'" - echo "Supported modes: vr-streaming, vr-offline, solo-text, dflash-image" + echo "Supported modes: vr-streaming, vr-offline, solo-text, solo-text-batching, dflash-image" exit 1 ;; esac @@ -288,8 +289,8 @@ if [ "$MODEL_NAME" = "muse_glimmer" ]; then echo "Error: Muse Glimmer requires quantization 'kquant-17gb' or 'kquant-dynamic'" exit 1 fi - if [ "$MODE" != "solo-text" ] && [ "$MODE" != "dflash-image" ]; then - echo "Error: Muse Glimmer requires mode 'solo-text' or 'dflash-image'" + if [ "$MODE" != "solo-text" ] && [ "$MODE" != "solo-text-batching" ] && [ "$MODE" != "dflash-image" ]; then + echo "Error: Muse Glimmer requires mode 'solo-text', 'solo-text-batching' or 'dflash-image'" exit 1 fi fi @@ -565,6 +566,18 @@ if [ "$MODEL_NAME" = "muse_glimmer" ]; then --backend cuda \ --output-dir "${OUTPUT_DIR}" ;; + solo-text-batching) + # The off-graph cache exports the batching methods CudaExecutor drives: + # forward_{1..64}, forward_others and the device samplers. + EXPORT_START_SECONDS=$SECONDS + TMPDIR="$INDUCTOR_TMPDIR" \ + TORCHINDUCTOR_CACHE_DIR="$INDUCTOR_CACHE" \ + python -m executorch.examples.models.muse_glimmer.export.export_solo \ + --gguf "$TARGET_GGUF_PATH" \ + --backend cuda \ + --use-offgraph-kv-cache \ + --output-dir "${OUTPUT_DIR}" + ;; dflash-image) DRAFT_GGUF_FILE="dflash-Muse-Glimmer-30B-Q4_K_M.gguf" MMPROJ_GGUF_FILE="mmproj-Muse-Glimmer-30B-Q4_K_M.gguf" @@ -581,7 +594,7 @@ if [ "$MODEL_NAME" = "muse_glimmer" ]; then --output-dir "${OUTPUT_DIR}" ;; *) - echo "Error: Muse Glimmer requires mode 'solo-text' or 'dflash-image'" + echo "Error: Muse Glimmer requires mode 'solo-text', 'solo-text-batching' or 'dflash-image'" exit 1 ;; esac diff --git a/.ci/scripts/test_model_e2e.sh b/.ci/scripts/test_model_e2e.sh index 427329ae961..193ebe14c64 100755 --- a/.ci/scripts/test_model_e2e.sh +++ b/.ci/scripts/test_model_e2e.sh @@ -46,6 +46,7 @@ Arguments: - vr-streaming: Voxtral Realtime streaming mode - vr-offline: Voxtral Realtime offline mode - solo-text: Muse Glimmer solo text mode + - solo-text-batching: Muse Glimmer text on the CUDA batching runtime - dflash-image: Muse Glimmer DFlash vision mode Environment: @@ -111,7 +112,7 @@ if [ -n "$MODE" ]; then exit 1 fi ;; - solo-text|dflash-image) + solo-text|solo-text-batching|dflash-image) if [ "$HF_MODEL" != "meta-models/Muse-Glimmer-30B-GGUF" ]; then echo "Error: Mode '$MODE' can only be used with Muse Glimmer model" echo "Provided model: $HF_MODEL" @@ -120,7 +121,7 @@ if [ -n "$MODE" ]; then ;; *) echo "Error: Unsupported mode '$MODE'" - echo "Supported modes: vr-streaming, vr-offline, solo-text, dflash-image" + echo "Supported modes: vr-streaming, vr-offline, solo-text, solo-text-batching, dflash-image" exit 1 ;; esac @@ -290,13 +291,20 @@ case "$HF_MODEL" in EXPECTED_OUTPUT="Paris" IMAGE_URL="" ;; + solo-text-batching) + # Built by `make muse-glimmer-batching-cuda` into its own directory. + RUNNER_PATH="muse-glimmer-batching" + RUNNER_TARGET="muse_glimmer_batching_runner" + EXPECTED_OUTPUT="Paris" + IMAGE_URL="" + ;; dflash-image) RUNNER_TARGET="dflash_runner" EXPECTED_OUTPUT="dog" IMAGE_URL="https://github.com/pytorch/hub/raw/master/images/dog.jpg" ;; *) - echo "Error: Muse Glimmer requires mode 'solo-text' or 'dflash-image'" + echo "Error: Muse Glimmer requires mode 'solo-text', 'solo-text-batching' or 'dflash-image'" exit 1 ;; esac @@ -453,6 +461,11 @@ EOF fi RUNNER_ARGS="$RUNNER_ARGS --tokenizer_path ${MODEL_DIR}/$TOKENIZER_FILE --prompt_file \"$PROMPT_FILE\" --max_new_tokens 512" RUNNER_ARGS="$RUNNER_ARGS --cuda_graph" + if [ "$MODE" = "solo-text-batching" ]; then + # The batching runner goes through ServingRuntime and takes its own flags. + printf '%s' '<|start|>user<|message|>What is the capital of France?<|eot|><|start|>assistant' > "$PROMPT_FILE" + RUNNER_ARGS="--pte ${MODEL_DIR}/model.pte --data_path ${MODEL_DIR}/aoti_cuda_blob.ptd --tokenizer ${MODEL_DIR}/$TOKENIZER_FILE --prompt \"\$(cat $PROMPT_FILE)\" --max_new_tokens 128 --temperature 0 --max_session_tokens 4096" + fi if [ "$MODE" = "dflash-image" ]; then RUNNER_ARGS="$RUNNER_ARGS --image_path ${MODEL_DIR}/test_image.jpg" fi @@ -540,6 +553,137 @@ else fi echo "::endgroup::" +if [ "$DEVICE" = "cuda" ] && [ "$MODEL_NAME" = "muse_glimmer" ] && [ "$MODE" = "solo-text-batching" ]; then + echo "::group::Run $MODEL_NAME batching serving smoke" + # The CI-only workload tool is not part of the normal serving build. + bash .ci/scripts/build_muse_glimmer_batching_bench.sh + # Concurrent staggered requests through ServingRuntime on CudaExecutor. + BENCH_REPORT="${MODEL_DIR}/batching_requests.json" + cmake-out/examples/models/muse-glimmer-batching/muse_glimmer_batching_bench \ + --pte "${MODEL_DIR}/model.pte" --data_path "${MODEL_DIR}/aoti_cuda_blob.ptd" \ + --tokenizer "${MODEL_DIR}/$TOKENIZER_FILE" --max_sessions 8 --max_session_tokens 2048 \ + --mode requests --requests 5 --prompt_tokens 600 --max_new_tokens 64 --stagger_ms 200 \ + --report_json "$BENCH_REPORT" + python - "$BENCH_REPORT" <<'PY' +import json +import sys + +report = json.load(open(sys.argv[1])) +rows = report["requests"] +assert report["max_batch_sessions"] > 1, report +assert report["multi_session_steps"] > 0, report +assert len(rows) == 5, rows +assert report["failed_requests"] == 0, report +for row in rows: + assert not row["failed"] and row["completion_tokens"] == 64, row +print(f"5 staggered requests: {report['tokens_per_s']:.1f} tok/s") +PY + + pip install -r examples/llm_server/python/requirements.txt + python -c "from huggingface_hub import hf_hub_download; [hf_hub_download('${HF_MODEL}', f, local_dir='${MODEL_DIR}') for f in ('tokenizer_config.json', 'chat_template.jinja')]" + PORT=$(python - <<'PY' +import socket + +with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + print(s.getsockname()[1]) +PY +) + SERVER_LOG=$(mktemp) + python -u -m executorch.examples.models.muse_glimmer.serving.serve_batching \ + --backend cuda \ + --worker-bin cmake-out/examples/models/muse-glimmer-batching/muse_glimmer_batching_worker \ + --model-path "${MODEL_DIR}/model.pte" \ + --data-path "${MODEL_DIR}/aoti_cuda_blob.ptd" \ + --tokenizer-path "${MODEL_DIR}/$TOKENIZER_FILE" \ + --hf-tokenizer "${MODEL_DIR}" \ + --max-context 4096 \ + --prefix-cache-entries 2 \ + --host 127.0.0.1 \ + --port "$PORT" >"$SERVER_LOG" 2>&1 & + SERVER_PID=$! + cleanup_mg_server() { + if kill -0 "$SERVER_PID" 2>/dev/null; then + kill "$SERVER_PID" 2>/dev/null || true + wait "$SERVER_PID" 2>/dev/null || true + fi + rm -f "$SERVER_LOG" + } + trap cleanup_mg_server EXIT + + if ! python - "$PORT" "$SERVER_LOG" <<'PY' +import json +import sys +import time +import urllib.request + +port, log_path = sys.argv[1], sys.argv[2] +base = f"http://127.0.0.1:{port}" + + +def request(path, payload=None): + data = None + headers = {} + if payload is not None: + data = json.dumps(payload).encode("utf-8") + headers["Content-Type"] = "application/json" + req = urllib.request.Request(base + path, data=data, headers=headers) + with urllib.request.urlopen(req, timeout=300) as resp: + return resp.read().decode("utf-8") + + +last = None +for _ in range(600): + try: + request("/health") + break + except Exception as e: + last = e + time.sleep(1) +else: + print(open(log_path, encoding="utf-8", errors="replace").read()) + raise RuntimeError(f"server did not become healthy: {last}") + +body = { + "model": "muse_glimmer", + "messages": [{"role": "user", "content": "What is the capital of France?"}], + "max_tokens": 128, + "temperature": 0, +} +# The second identical request reuses the first one's prompt snapshot. +for attempt in range(2): + resp = json.loads(request("/v1/chat/completions", body)) + message = resp["choices"][0]["message"] + text = (message.get("content") or "") + (message.get("reasoning_content") or "") + if "Paris" not in text: + raise AssertionError(f"expected Paris in serving response, got: {message!r}") + details = resp["usage"].get("prompt_tokens_details") or {} + cached = details.get("cached_tokens", 0) + print(f"request {attempt}: prompt {resp['usage']['prompt_tokens']} cached {cached}") +if cached == 0: + raise AssertionError("the repeated request reused no prompt tokens") + +stream = request("/v1/chat/completions", dict(body, stream=True)) +pieces = [] +for line in stream.splitlines(): + if line.startswith("data: ") and line != "data: [DONE]": + delta = json.loads(line[6:])["choices"][0]["delta"] + pieces.append((delta.get("content") or "") + (delta.get("reasoning_content") or "")) +if len([p for p in pieces if p]) < 2 or "Paris" not in "".join(pieces): + raise AssertionError(f"expected Paris streamed in several chunks, got: {pieces!r}") +print("Muse Glimmer CUDA batching serving smoke passed") +PY + then + echo "Muse Glimmer batching serving smoke failed; server log:" + cat "$SERVER_LOG" + exit 1 + fi + + cleanup_mg_server + trap - EXIT + echo "::endgroup::" +fi + if [ "$DEVICE" = "cuda" ] && [ "$MODEL_NAME" = "qwen3_5_moe" ]; then echo "::group::Run $MODEL_NAME OpenAI serving smoke" pip install -r examples/llm_server/python/requirements.txt "transformers==5.0.0rc1" diff --git a/.ci/scripts/tests/test_muse_glimmer_batching_ci.py b/.ci/scripts/tests/test_muse_glimmer_batching_ci.py new file mode 100644 index 00000000000..b701a057136 --- /dev/null +++ b/.ci/scripts/tests/test_muse_glimmer_batching_ci.py @@ -0,0 +1,62 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import json +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[3] +SCRIPTS = ROOT / ".ci/scripts" + + +class MuseGlimmerBatchingCITest(unittest.TestCase): + def test_e2e_builds_the_ci_tool_before_running_it(self): + script = (SCRIPTS / "test_model_e2e.sh").read_text() + build = script.index("bash .ci/scripts/build_muse_glimmer_batching_bench.sh") + run = script.index( + "cmake-out/examples/models/muse-glimmer-batching/muse_glimmer_batching_bench" + ) + self.assertLess(build, run) + workflow = (ROOT / ".github/workflows/cuda.yml").read_text() + unit_job = workflow.split(" unittest-cuda-muse-glimmer-batching:", 1)[1].split( + " unittest-cuda-quantized-gemm:", 1 + )[0] + self.assertIn("bash .ci/scripts/build_muse_glimmer_batching_bench.sh", unit_job) + self.assertIn("test_muse_glimmer_batching_*.py", unit_job) + self.assertIn(".ci/scripts/*muse_glimmer_batching*", workflow) + + def test_completion_without_actual_batching_fails_the_smoke_check(self): + script = (SCRIPTS / "test_model_e2e.sh").read_text() + checker = script.split(' python - "$BENCH_REPORT" <<\'PY\'\n', 1)[1].split( + "\nPY\n", 1 + )[0] + for changed, succeeds in ( + ({}, True), + ({"max_batch_sessions": 1, "multi_session_steps": 0}, False), + ({"failed_requests": 1}, False), + ({"requests": [{"failed": False, "completion_tokens": 63}] * 5}, False), + ): + with self.subTest(changed=changed), tempfile.TemporaryDirectory() as directory: + report = { + "requests": [{"failed": False, "completion_tokens": 64}] * 5, + "failed_requests": 0, + "max_batch_sessions": 3, + "multi_session_steps": 10, + "tokens_per_s": 100.0, + } + report.update(changed) + path = Path(directory) / "report.json" + path.write_text(json.dumps(report)) + result = subprocess.run( + [sys.executable, "-c", checker, str(path)], + capture_output=True, + text=True, + ) + self.assertEqual(result.returncode == 0, succeeds, result.stderr) diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 25d35c8fdfb..2dd43909d2c 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -34,6 +34,8 @@ on: - .ci/scripts/cuda_benchmark.py - .ci/scripts/cuda_benchmark_from_e2e.sh - .ci/scripts/cuda_perf_prompts/** + - .ci/scripts/*muse_glimmer_batching* + - .ci/scripts/tests/*muse_glimmer_batching* - examples/models/muse-glimmer/** - extension/pybindings/** - runtime/__init__.py @@ -554,6 +556,61 @@ jobs: pip install gguf python -m pytest examples/models/gemma4_31b/tests/ --ignore=examples/models/gemma4_31b/tests/test_mlx_pipeline.py -v -o "addopts=" + unittest-cuda-muse-glimmer-batching: + name: unittest-cuda-muse-glimmer-batching + needs: [changed-files, run-decision] + if: | + contains(needs.changed-files.outputs.changed-files, 'backends/cuda') || + contains(needs.changed-files.outputs.changed-files, 'backends/aoti') || + contains(needs.changed-files.outputs.changed-files, 'extension/llm/batching') || + contains(needs.changed-files.outputs.changed-files, 'extension/llm/serving') || + contains(needs.changed-files.outputs.changed-files, 'examples/llm_server/cpp') || + contains(needs.changed-files.outputs.changed-files, '.ci/scripts/muse_glimmer_batching') || + contains(needs.changed-files.outputs.changed-files, '.ci/scripts/build_muse_glimmer_batching') || + contains(needs.changed-files.outputs.changed-files, '.ci/scripts/tests/test_muse_glimmer_batching') || + contains(needs.changed-files.outputs.changed-files, '.ci/scripts/tests/muse_glimmer_batching') || + contains(needs.changed-files.outputs.changed-files, 'examples/models/muse-glimmer') || + contains(needs.changed-files.outputs.changed-files, '.github/workflows/cuda.yml') || + needs.run-decision.outputs.is-full-run == 'true' + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + permissions: + id-token: write + contents: read + with: + timeout: 120 + runner: mt-l-x86aavx2-11-41-a10g + gpu-arch-type: cuda + gpu-arch-version: "13.0" + submodules: recursive + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + script: | + set -eux + bash ./install_executorch.sh + conda install -y -c conda-forge 'libstdcxx-ng>=12' + export LD_LIBRARY_PATH=/opt/conda/lib:$LD_LIBRARY_PATH + + python -m unittest discover -s .ci/scripts/tests -p 'test_muse_glimmer_batching_*.py' -v + + # The normal serving binaries and the opt-in CI workload tool share + # CudaExecutor; MG preparation tests also run on Linux. + make muse-glimmer-batching-cuda + BATCHING_BIN_DIR="${PWD}/cmake-out/examples/models/muse-glimmer-batching" + cmake -S "${PWD}/examples/models/muse-glimmer" -B "${BATCHING_BIN_DIR}" \ + -DMUSE_GLIMMER_BUILD_BATCHING_TESTS=ON + bash .ci/scripts/build_muse_glimmer_batching_bench.sh "${BATCHING_BIN_DIR}" + cmake --build "${BATCHING_BIN_DIR}" --target muse_glimmer_batching_test --parallel + ctest --test-dir "${BATCHING_BIN_DIR}" \ + -R '^muse_glimmer_batching$' --output-on-failure --no-tests=error + test -x "${BATCHING_BIN_DIR}/muse_glimmer_batching_runner" + test -x "${BATCHING_BIN_DIR}/muse_glimmer_batching_worker" + test -x "${BATCHING_BIN_DIR}/muse_glimmer_batching_bench" + + # The batching export on a tiny model, and the CLI's backend dispatch. + pip install gguf + python -m pytest -v -o "addopts=" \ + examples/models/muse-glimmer/tests/test_cuda_batching_pipeline.py \ + examples/models/muse-glimmer/tests/test_checkpoint_context.py + unittest-cuda-quantized-gemm: name: unittest-cuda-quantized-gemm (${{ matrix.shard }}) needs: [changed-files, run-decision] @@ -1042,6 +1099,10 @@ jobs: name: test-muse-glimmer-cuda-e2e-${{ matrix.variant }}-${{ matrix.mode }} needs: [changed-files, run-decision] if: | + contains(needs.changed-files.outputs.changed-files, '.ci/scripts/muse_glimmer_batching') || + contains(needs.changed-files.outputs.changed-files, '.ci/scripts/build_muse_glimmer_batching') || + contains(needs.changed-files.outputs.changed-files, '.ci/scripts/tests/test_muse_glimmer_batching') || + contains(needs.changed-files.outputs.changed-files, '.ci/scripts/tests/muse_glimmer_batching') || contains(needs.changed-files.outputs.changed-files, 'backends/cuda') || contains(needs.changed-files.outputs.changed-files, 'backends/aoti') || contains(needs.changed-files.outputs.changed-files, 'extension/cuda') || @@ -1061,6 +1122,8 @@ jobs: include: - variant: kquant-17gb mode: solo-text + - variant: kquant-17gb + mode: solo-text-batching - variant: kquant-17gb mode: dflash-image - variant: kquant-dynamic