[executorch][ci] Test Muse Glimmer batching on CUDA in OSS CI

.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

.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"

.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