From e0beee286a2a8756f3f2703dd93e3949dd36dc54 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Mon, 12 Oct 2026 22:22:00 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .ci/scripts/export_model_artifact.sh | 23 ++++- .ci/scripts/test_model_e2e.sh | 146 ++++++++++++++++++++++++++- .github/workflows/cuda.yml | 51 ++++++++++ 3 files changed, 212 insertions(+), 8 deletions(-) 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..efef587ebbd 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,133 @@ else fi echo "::endgroup::" +if [ "$DEVICE" = "cuda" ] && [ "$MODEL_NAME" = "muse_glimmer" ] && [ "$MODE" = "solo-text-batching" ]; then + echo "::group::Run $MODEL_NAME batching serving smoke" + # 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 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/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 25d35c8fdfb..acf684b6337 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -554,6 +554,55 @@ 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, '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 + + # The CUDA counterpart of the MLX batching build job: the batching + # runner, worker and bench build against CudaExecutor, and the shared + # MG preparation tests 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 + cmake --build "${BATCHING_BIN_DIR}" --target muse_glimmer_batching_test muse_glimmer_batching_bench --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] @@ -1061,6 +1110,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