feat: optional startup tuning (STARTUP_MODE, CUDA_GRAPH_SIZES, QUANT, CACHE_DIR)

README.md

@@ -90,7 +90,8 @@ says the model weights are stored under `/repository`, which is where this image
 serves on port `8000` and answers `/health`, so those are the values for the endpoint's container port
 and health route. The guide lists public images and private ones on Docker Hub, AWS ECR, Azure ACR
 and Google GCR. GitHub Container Registry is not in that list, but the image is public, so it needs no
-registry credentials.
+registry credentials. The platform has no persistent storage, so `CACHE_DIR` has nothing to survive
+a restart there; the [startup tuning](#fast-startup-optional) options still work.
 
 ### Check that it works
 
@@ -207,6 +208,9 @@ name is passed on to `vllm serve`.
 | `MAX_MODEL_LEN` | `32768` | Longest prompt plus output. |
 | `GPU_MEMORY_UTILIZATION` | `0.90` | Share of GPU memory vLLM may use. |
 | `TENSOR_PARALLEL_SIZE` | every visible GPU | GPUs to split the model across. |
+| `STARTUP_MODE` | `default` | `fast` compiles less and captures CUDA graphs only for small batches; `eager` skips compilation and graphs entirely. See [Fast startup](#fast-startup-optional). |
+| `CUDA_GRAPH_SIZES` | vLLM's own set | Comma-separated batch sizes to capture CUDA graphs for, e.g. `1,2,4,8,16,32`. Overrides `fast` mode's set. |
+| `CACHE_DIR` | `/cache` | If writable, vLLM's compile cache and Triton's kernels are kept under it, so mounting a volume there persists them across restarts. Unwritable or unmounted, both stay ephemeral. |
 | `DECISION_ADAPTER_NAME` | `decision` (the plugin alone uses the checkpoint's only adapter, else `decision`) | Which adapter answers. Checked against the checkpoint at startup; an unknown name stops the server. |
 | `DECISION_INVOCATION_MARKER` | none (`### DECISION` in the image) | The text the adapter was trained to activate on. It is required. Startup checks that a probe prompt activates the adapter exactly once and stops if not, but it cannot tell a working marker from the one the adapter was trained on. |
 | `DECISION_SPEECH_INPUT` | `0` | `1` keeps speech input; see [Speech input](#speech-input-optional). |
@@ -244,6 +248,67 @@ On every decision request the plugin itself sets `temperature=0.0`, `max_complet
 caps `logprob_token_ids` at 128 ids, so a `choice` with more options runs as several requests over the
 same prompt, served from the prefix cache.
 
+### Fast startup (optional)
+
+A vLLM boot spends noticeable time in `torch.compile` and CUDA-graph capture before the first token.
+Three environment variables trade that against steady-state performance, for demos or for
+scale-to-zero deployments where every wake pays the boot again:
+
+- `STARTUP_MODE=fast` passes `--optimization-level 1` (less compilation) and captures CUDA graphs
+  only for batch sizes 1-32. Batches larger than the largest captured size still run, without a
+  graph, so the effect of a too-small set is slower decode, never a failure.
+- `STARTUP_MODE=eager` passes `--enforce-eager`: no compilation and no graphs at all, the fastest
+  boot and the slowest serving. For a one-user demo the decode cost is usually irrelevant.
+- `CUDA_GRAPH_SIZES=1,2,4,8` trims the capture set on its own, with `STARTUP_MODE` left at
+  `default`.
+
+All three default off; an unset or empty value means today's behaviour, and a value that is not one
+of the listed ones stops the server at startup with the offending value in the log. The first log
+line after startup begins `startup:` and names the active mode, capture sizes and cache location, so
+a deployment's own configuration is visible in its logs.
+
+Measured on a DGX Spark (GB10, 121.6 GB unified memory) serving the 8b, time from container start to
+`/health` answering 200, each arm a fresh container with `MAX_MODEL_LEN=4096`,
+`--max-num-seqs 4` and `--max-num-batched-tokens 1024`:
+
+| Arm | To ready | Weight load | Compile | Graph capture |
+|---|---|---|---|---|
+| `STARTUP_MODE=eager` | **232 s** | 99 s | none | none |
+| `STARTUP_MODE=fast` | 255 s | 104 s | 16 s | 97 s |
+| `default` | 279 s | 122 s | 16 s | (vLLM's set) |
+
+Two conclusions, both specific to that hardware. Weight loading dominates the boot, so a warm page
+cache matters more than any flag here. And CUDA-graph capture is expensive on GB10 (97 s for six
+sizes), so `eager` is the fastest startup mode there, not `fast`. On H100-class hardware capture is
+far cheaper (published figures put a trimmed set near 7 s against 54 s for the default set), which is
+where `fast` earns its keep. Treat the table as a starting point to measure against, not a rule.
+
+`CACHE_DIR` (default `/cache`) adds the other half: vLLM writes its `torch.compile` artifacts and
+Triton its compiled kernels under it when it is writable, so mounting a volume there
+
+```bash
+docker run -d --name gss --gpus all --ipc=host -p 8000:8000 \
+    -v /srv/models/granite-switch-4.1-8b:/repository:ro \
+    -v gss-cache:/cache \
+    ghcr.io/generative-computing/granite-switch-decision:0.2.0
+```
+
+keeps both across restarts, and a restart skips most of the compile cost. The caches are keyed on
+the vLLM version, the model, the configuration and the GPU, so they recompile silently after an
+upgrade and a volume can be shared only between identical deployments. Hugging Face Inference
+Endpoints has no persistent storage, so there the caches stay ephemeral whatever this is set to;
+the option is for self-hosted containers.
+
+### Quantization doesn't work here
+
+vLLM's fp8 quantization cannot serve these checkpoints. `--quantization fp8` fails at load with
+`LoRA weight load: 168 parameter(s) absent from the checkpoint would be served UNINITIALIZED`, because
+the aLoRA adapters are not quantized and have no corresponding scale parameters, and
+`--kv-cache-dtype fp8` fails in `input_quant_fp8.forward_cuda` on `assert input.ndim == 2` in the
+switch forward path. Both were measured on a GB10 against the 8b. There is no `QUANT` option in this
+image for that reason; passing one stops the server at startup with the reason rather than failing
+later. For a smaller memory footprint, use a smaller checkpoint.
+
 ## Behaviour and limits
 
 - **Prompt.** Two user turns: `State:\n<state>`, then `### DECISION` with the question and lettered

serve.sh

@@ -5,6 +5,8 @@
 #   MODEL_DIR               default /repository (HF Inference Endpoints' mount point)
 #   TENSOR_PARALLEL_SIZE    default: every visible GPU
 #   MAX_MODEL_LEN (32768), GPU_MEMORY_UTILIZATION (0.90), PORT (8000), SERVED_MODEL_NAME
+#   STARTUP_MODE (default), CUDA_GRAPH_SIZES, CACHE_DIR (/cache): optional startup tuning; see
+#                           README, "Fast startup". All default to today's behaviour.
 #   DECISION_ADAPTER_NAME  default "decision"; DECISION_INVOCATION_MARKER, see README
 #   DECISION_SPEECH_INPUT  default 0: speech input is switched off, so no speech model is needed.
 #                           1 keeps the checkpoint's speech input; see README, "Speech input".
@@ -20,6 +22,51 @@ case "${DECISION_SPEECH_INPUT:-0}" in
 esac
 TP=${TENSOR_PARALLEL_SIZE:-$(python3 -c "import torch; print(max(torch.cuda.device_count(), 1))")}
 export DECISION_ADAPTER_NAME=${DECISION_ADAPTER_NAME:-decision}
+
+# Startup tuning. Every option defaults to today's behaviour; a value that is set but not one
+# of the listed ones stops the server rather than silently serving with unexpected settings.
+STARTUP_ARGS=()
+case "${STARTUP_MODE:-}" in
+    ""|default) ;;
+    fast) STARTUP_ARGS+=(--optimization-level 1) ;;
+    eager) STARTUP_ARGS+=(--enforce-eager) ;;
+    *) echo "STARTUP_MODE must be default, fast or eager, not '$STARTUP_MODE'" >&2; exit 2 ;;
+esac
+if [ -n "${CUDA_GRAPH_SIZES:-}" ]; then
+    if ! [[ "$CUDA_GRAPH_SIZES" =~ ^[0-9]+(,[0-9]+)*$ ]]; then
+        echo "CUDA_GRAPH_SIZES must be comma-separated batch sizes, not '$CUDA_GRAPH_SIZES'" >&2
+        exit 2
+    fi
+    STARTUP_ARGS+=(--cudagraph-capture-sizes ${CUDA_GRAPH_SIZES//,/ })
+elif [ "${STARTUP_MODE:-}" = "fast" ]; then
+    # Batches past the largest captured size still run, without a graph. Capture beyond the
+    # set costs boot time that a demo's traffic never uses.
+    STARTUP_ARGS+=(--cudagraph-capture-sizes 1 2 4 8 16 32)
+fi
+case "${QUANT:-}" in
+    ""|off) ;;
+    *)
+        # Both values were tried on a GB10 against the 8b checkpoint and both fail inside vLLM:
+        # fp8 weights trip "LoRA weight load: 168 parameter(s) absent from the checkpoint would be
+        # served UNINITIALIZED" (the aLoRA adapters are not quantized), and fp8 KV cache trips
+        # "assert input.ndim == 2" in input_quant_fp8.forward_cuda. Refuse rather than fail later.
+        echo "QUANT is not supported for granite-switch checkpoints: vLLM's fp8 weight and KV" \
+             "quantization both fail against the checkpoint's aLoRA adapters. Unset QUANT." >&2
+        exit 2 ;;
+esac
+# A writable CACHE_DIR persists vLLM's compile cache and Triton's kernels across restarts when
+# it is a mounted volume. Unwritable or unmounted, both stay ephemeral, which is today's
+# behaviour (the platform Hugging Face Inference Endpoints runs on has no persistent storage).
+CACHE_DIR=${CACHE_DIR:-/cache}
+if mkdir -p "$CACHE_DIR" 2>/dev/null && [ -w "$CACHE_DIR" ]; then
+    mkdir -p "$CACHE_DIR/vllm" "$CACHE_DIR/triton"
+    export VLLM_CACHE_ROOT="$CACHE_DIR/vllm" TRITON_CACHE_DIR="$CACHE_DIR/triton"
+else
+    echo "CACHE_DIR '$CACHE_DIR' is not writable; compile and Triton caches stay ephemeral" >&2
+fi
+echo "startup: mode=${STARTUP_MODE:-default} cuda_graph_sizes=${CUDA_GRAPH_SIZES:-vllm-default}" \
+     "cache=${VLLM_CACHE_ROOT:-ephemeral}"
+
 # --dtype bfloat16: composed checkpoints carry torch_dtype null (granite-switch#126).
 # flash_attn_version=2: avoids the FA3 schedule bug (granite-switch#139).
 # --max-logprobs 256: lets chat clients that read labels by top-K ask for more than 20.
@@ -38,4 +85,5 @@ exec vllm serve "$MODEL_DIR" \
     --enable-prompt-tokens-details \
     --enable-auto-tool-choice --tool-call-parser granite4 \
     "${SPEECH_ARGS[@]}" \
+    "${STARTUP_ARGS[@]}" \
     "$@"

tests/test_serve_sh.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
-"""No-GPU checks for serve.sh's speech-input switch, with a stub `vllm` that records its arguments.
+"""No-GPU checks for serve.sh's environment switches, with a stub `vllm` that records its
+arguments and the cache environment it would run with.
 
 python3 -m pytest tests/test_serve_sh.py -v --tb=short
 """
@@ -20,74 +21,187 @@ def _run(tmp_path: Path, extra_env: dict[str, str], *args: str, asr_enabled: boo
     bin_dir = tmp_path / "bin"
     bin_dir.mkdir()
     stub = bin_dir / "vllm"
-    stub.write_text('#!/bin/sh\nprintf "%s\\n" "$@" > "$STUB_OUT"\n')
+    stub.write_text(
+        "#!/bin/sh\n"
+        'printf "%s\\n" "$@" > "$STUB_OUT"\n'
+        'env | grep -E "^(VLLM_CACHE_ROOT|TRITON_CACHE_DIR)=" > "$STUB_ENV" || true\n'
+    )
     stub.chmod(0o755)
     model_dir = tmp_path / "model"
     model_dir.mkdir()
     (model_dir / "config.json").write_text(json.dumps({"asr_enabled": asr_enabled}))
     out = tmp_path / "args.txt"
+    env_out = tmp_path / "env.txt"
     env = {
         "PATH": f"{bin_dir}{os.pathsep}{Path(sys.executable).parent}{os.pathsep}{os.environ['PATH']}",
         "MODEL_DIR": str(model_dir),
         "TENSOR_PARALLEL_SIZE": "1",
         "STUB_OUT": str(out),
+        "STUB_ENV": str(env_out),
+        "CACHE_DIR": str(tmp_path / "cache"),
         "HF_HUB_OFFLINE": "1",
         "HF_HUB_CACHE": str(tmp_path / "empty-cache"),
         # these tests are about the speech model, not the audio libraries
         "DECISION_PREFLIGHT_AUDIO_MODULES": "",
         **extra_env,
     }
     done = subprocess.run(["bash", str(SERVE_SH), *args], env=env, capture_output=True, text=True)
-    return done, (out.read_text().splitlines() if out.exists() else [])
+    argv = out.read_text().splitlines() if out.exists() else []
+    cache_env = dict(line.split("=", 1) for line in env_out.read_text().splitlines()) if env_out.exists() else {}
+    return done, argv, cache_env
 
 
 def test_speech_input_is_off_by_default_and_needs_no_speech_model(tmp_path):
-    done, argv = _run(tmp_path, {})
+    done, argv, _ = _run(tmp_path, {})
     assert done.returncode == 0, done.stderr
     assert argv[argv.index("--hf-overrides") + 1] == OVERRIDE
 
 
 def test_speech_input_on_without_the_model_stops_before_vllm(tmp_path):
-    done, argv = _run(tmp_path, {"DECISION_SPEECH_INPUT": "1"})
+    done, argv, _ = _run(tmp_path, {"DECISION_SPEECH_INPUT": "1"})
     assert done.returncode == 1
     assert argv == []  # vllm never started
     assert "WITH_ASR=1" in done.stderr
 
 
 def test_speech_input_on_with_the_model_cached_starts_without_the_override(tmp_path):
-    cache = tmp_path / "empty-cache" / "models--ibm-granite--granite-speech-5.0-470m-turboctc" / "snapshots" / "abc"
+    speech = "models--ibm-granite--granite-speech-5.0-470m-turboctc"
+    cache = tmp_path / "empty-cache" / speech / "snapshots" / "abc"
     cache.mkdir(parents=True)
     (cache / "config.json").write_text("{}")
-    done, argv = _run(tmp_path, {"DECISION_SPEECH_INPUT": "1"})
+    done, argv, _ = _run(tmp_path, {"DECISION_SPEECH_INPUT": "1"})
     assert done.returncode == 0, done.stderr
     assert "--hf-overrides" not in argv
 
 
 def test_speech_input_on_without_the_audio_libraries_stops_before_vllm(tmp_path):
-    done, argv = _run(
+    done, argv, _ = _run(
         tmp_path,
         {
             "DECISION_SPEECH_INPUT": "1",
             "DECISION_PREFLIGHT_AUDIO_MODULES": "no_such_audio_module_xyz",
         },
     )
     assert done.returncode == 1
-    assert argv == []  # vllm never started
+    assert argv == []
     assert "no_such_audio_module_xyz" in done.stderr
     assert "HTTP 500" in done.stderr
 
 
 def test_extra_arguments_still_reach_vllm(tmp_path):
-    done, argv = _run(tmp_path, {}, "--api-key", "k")
+    done, argv, _ = _run(tmp_path, {}, "--api-key", "k")
     assert done.returncode == 0, done.stderr
     assert argv[-2:] == ["--api-key", "k"]
 
 
-@pytest.mark.parametrize("value", ["2", "yes", ""])
+@pytest.mark.parametrize("value", ["2", "yes"])
 def test_bad_switch_value_is_refused(tmp_path, value):
-    done, argv = _run(tmp_path, {"DECISION_SPEECH_INPUT": value})
-    if value == "":
-        assert done.returncode == 0  # empty falls back to the default, off
-    else:
-        assert done.returncode == 2
-        assert argv == []
+    done, argv, _ = _run(tmp_path, {"DECISION_SPEECH_INPUT": value})
+    assert done.returncode == 2
+    assert argv == []
+
+
+def test_empty_switch_value_falls_back_to_off(tmp_path):
+    done, argv, _ = _run(tmp_path, {"DECISION_SPEECH_INPUT": ""})
+    assert done.returncode == 0
+
+
+# Startup tuning: every option defaults to today's flags, and a set-but-unrecognised value
+# stops the server before vllm starts.
+
+
+def test_startup_defaults_add_no_flags_and_log_the_banner(tmp_path):
+    done, argv, _ = _run(tmp_path, {})
+    assert done.returncode == 0, done.stderr
+    assert "--optimization-level" not in argv
+    assert "--enforce-eager" not in argv
+    assert "--cudagraph-capture-sizes" not in argv
+    assert "--kv-cache-dtype" not in argv
+    assert "--quantization" not in argv
+    assert "startup: mode=default" in done.stdout
+
+
+def test_startup_mode_fast_compiles_less_and_trims_capture_sizes(tmp_path):
+    done, argv, _ = _run(tmp_path, {"STARTUP_MODE": "fast"})
+    assert done.returncode == 0, done.stderr
+    assert argv[argv.index("--optimization-level") + 1] == "1"
+    i = argv.index("--cudagraph-capture-sizes")
+    assert argv[i + 1 : i + 7] == ["1", "2", "4", "8", "16", "32"]
+
+
+def test_startup_mode_eager_skips_compilation_and_graphs(tmp_path):
+    done, argv, _ = _run(tmp_path, {"STARTUP_MODE": "eager"})
+    assert done.returncode == 0, done.stderr
+    assert "--enforce-eager" in argv
+    assert "--optimization-level" not in argv
+
+
+def test_startup_mode_fast_respects_an_explicit_capture_size_set(tmp_path):
+    done, argv, _ = _run(tmp_path, {"STARTUP_MODE": "fast", "CUDA_GRAPH_SIZES": "4,8"})
+    assert done.returncode == 0, done.stderr
+    i = argv.index("--cudagraph-capture-sizes")
+    assert argv[i + 1 : i + 3] == ["4", "8"]
+
+
+def test_cuda_graph_sizes_alone_trims_capture_without_touching_optimization(tmp_path):
+    done, argv, _ = _run(tmp_path, {"CUDA_GRAPH_SIZES": "1,2,4"})
+    assert done.returncode == 0, done.stderr
+    i = argv.index("--cudagraph-capture-sizes")
+    assert argv[i + 1 : i + 4] == ["1", "2", "4"]
+    assert "--optimization-level" not in argv
+
+
+@pytest.mark.parametrize("value", ["turbo", "FAST"])
+def test_bad_startup_mode_is_refused(tmp_path, value):
+    done, argv, _ = _run(tmp_path, {"STARTUP_MODE": value})
+    assert done.returncode == 2
+    assert argv == []
+    assert value in done.stderr
+
+
+@pytest.mark.parametrize("value", ["1,x", "1;2", ",", "1,,2"])
+def test_bad_cuda_graph_sizes_are_refused(tmp_path, value):
+    done, argv, _ = _run(tmp_path, {"CUDA_GRAPH_SIZES": value})
+    assert done.returncode == 2
+    assert argv == []
+    assert value in done.stderr
+
+
+# Both quantization values were measured failing inside vLLM against the 8b checkpoint's aLoRA
+# adapters, so any value other than unset/off is refused at startup rather than failing later.
+
+
+@pytest.mark.parametrize("value", ["kv_fp8", "fp8", "int4"])
+def test_any_quant_value_is_refused_with_the_reason(tmp_path, value):
+    done, argv, _ = _run(tmp_path, {"QUANT": value})
+    assert done.returncode == 2
+    assert argv == []
+    assert "not supported for granite-switch" in done.stderr
+    assert "aLoRA" in done.stderr
+
+
+@pytest.mark.parametrize("value", ["", "off"])
+def test_quant_empty_and_off_are_accepted(tmp_path, value):
+    done, argv, _ = _run(tmp_path, {"QUANT": value})
+    assert done.returncode == 0, done.stderr
+    assert "--quantization" not in argv
+    assert "--kv-cache-dtype" not in argv
+
+
+def test_writable_cache_dir_persists_compile_and_triton_caches(tmp_path):
+    done, _, cache_env = _run(tmp_path, {})
+    assert done.returncode == 0, done.stderr
+    assert cache_env["VLLM_CACHE_ROOT"] == str(tmp_path / "cache" / "vllm")
+    assert cache_env["TRITON_CACHE_DIR"] == str(tmp_path / "cache" / "triton")
+
+
+def test_unwritable_cache_dir_falls_back_to_ephemeral_caches(tmp_path):
+    # A file where a directory is needed: mkdir fails for any user, root included.
+    blocked = tmp_path / "blocked"
+    blocked.write_text("not a directory")
+    done, argv, cache_env = _run(tmp_path, {"CACHE_DIR": str(blocked)})
+    assert done.returncode == 0, done.stderr
+    assert argv  # vllm still started
+    assert cache_env == {}
+    assert "not writable" in done.stderr
+    assert "startup: mode=default" in done.stdout