Skip to content

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

Draft
planetf1 wants to merge 2 commits into
mainfrom
feat/startup-tuning-options
Draft

planetf1 wants to merge 2 commits into
mainfrom
feat/startup-tuning-options

Conversation

@planetf1

@planetf1 planetf1 commented Nov 9, 2026

Copy link
Copy Markdown
Collaborator

What this gives users

A deployment that boots faster when it needs to: scale-to-zero endpoints, demo boxes, and any restart that pays vLLM's compile and CUDA-graph capture cost again. Today that cost is fixed; after this PR it is tunable with four environment variables, all defaulting to exactly today's behaviour, so existing deployments see no change at all.

  • STARTUP_MODE=fast — less compilation (--optimization-level 1) and CUDA graphs captured only for batch sizes 1–32. Batches past the largest captured size still run, without a graph: the failure mode of a too-small set is slower decode, never an error.
  • STARTUP_MODE=eager — --enforce-eager, no compilation and no graphs; fastest boot, slowest serving.
  • CUDA_GRAPH_SIZES=1,2,4,8 — trims the capture set on its own.
  • QUANT=kv_fp8 / QUANT=fp8 — KV-cache fp8 (about half the KV memory) and weight quantization on load. Both perturb decision probabilities slightly; the README says to treat published scores as unquantized-only.
  • CACHE_DIR (default /cache) — when writable, vLLM's compile cache and Triton kernels live under it, so a mounted volume persists them across restarts and a restart skips most of the compile cost. Unwritable or unmounted (permanently the case on Inference Endpoints), behaviour is unchanged.

Fail-safe by construction: unset or empty values mean today's defaults, and a set-but-unrecognised value stops the server at startup with the offending value in the log, matching the existing DECISION_SPEECH_INPUT convention. The first stdout line of every boot is a startup: banner naming the active mode, quantization, capture sizes and cache location, so a deployment's own configuration is visible in its logs.

Mechanism

serve.sh builds a STARTUP_ARGS array from validated env values and appends it to the existing vllm serve invocation. No flags reach vLLM unless an option is explicitly set, and every flag name was verified against vllm serve --help in the shipped vllm/vllm-openai:v0.30.0 image (notably --optimization-level, not the -O1 shorthand the vLLM docs use, and --cudagraph-capture-sizes takes space-separated sizes).

Limitations

  • No GPU measurements yet in this PR: the option set is designed so timings can be collected against a test image (published under the branch tag) before any default is reconsidered.
  • QUANT=fp8 and kv_fp8 change decision outputs; they are opt-in and documented as such.
  • Cache persistence requires a writable volume; Hugging Face Inference Endpoints has none, so there the cache options are inert by design.

Testing

  • tests/test_serve_sh.py extended from 8 to 24 cases: defaults add no flags, each mode adds exactly its flags, CUDA_GRAPH_SIZES overrides fast's default set, empty values fall back, every invalid value (bad mode, malformed size list, unknown quant) exits 2 before vllm starts, writable CACHE_DIR exports both cache roots, and an unwritable one falls back to ephemeral with a warning while the server still starts.
  • Full suite: 324 passed, 4 skipped. ruff format --check, ruff check and the pre-commit hooks (SPDX, gitleaks) pass.
  • Test image: the Image workflow was dispatched against this branch with push, publishing ghcr.io/generative-computing/granite-switch-decision:feat-startup-tuning-options (and -asr) for live validation. latest and 0.2.0 are untouched.

…IZES, QUANT, CACHE_DIR)

Every option defaults to today's behaviour; a set-but-unrecognised value stops
the server at startup instead of silently changing the serving configuration.
The first stdout line after startup echoes the active settings so a deployment
shows its own configuration in its logs.

STARTUP_MODE=fast trades steady-state throughput for boot time
(--optimization-level 1 plus a small default CUDA-graph capture set, batches
past the set still run without a graph). STARTUP_MODE=eager skips compilation
and graphs entirely. CUDA_GRAPH_SIZES trims the capture set on its own.
QUANT=kv_fp8 halves KV-cache memory; QUANT=fp8 also quantizes weights on load.
CACHE_DIR (default /cache) persists the vLLM compile cache and Triton kernels
across restarts when it is a mounted volume, and degrades to today's ephemeral
behaviour when it is not writable, which is permanent on Inference Endpoints.

The vLLM flag names are verified against vllm/vllm-openai:v0.30.0's own
'vllm serve --help' output. serve.sh tests extended from 8 to 24 cases.

Assisted-by: pi
Signed-off-by: Nigel Jones <jonesn@uk.ibm.com>
…pters

Measured on a DGX Spark (GB10) against the 8b checkpoint:

- --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 carry no scale parameters.
- --kv-cache-dtype fp8 fails in input_quant_fp8.forward_cuda on
  "assert input.ndim == 2" in the switch forward path.

Neither can serve the only model this image supports, so the option is
removed rather than shipped broken. Any QUANT value now stops the server at
startup with that reason instead of failing deep in engine init.

README gains the measured boot table for the same hardware: eager 232 s,
fast 255 s, default 279 s. Weight loading (99-122 s) dominates, and CUDA
graph capture costs 97 s for six sizes on GB10, so eager is the fastest
startup mode there while fast remains the right choice where capture is
cheap. Banner no longer reports a quantization setting.

Assisted-by: pi
Signed-off-by: Nigel Jones <jonesn@uk.ibm.com>
@planetf1

Copy link
Copy Markdown
Collaborator Author

Measured on a DGX Spark (GB10, 121.6 GB unified memory), arm64 image, 8b checkpoint

Ran the test image against the local 8b mount, each arm a fresh container, max_model_len=4096, --max-num-seqs 4, --max-num-batched-tokens 1024, gpu_memory_utilization=0.19 (sharing the GPU with another engine on the box). Time from docker run to /health returning 200:

Arm To ready Weight load Compile Graph capture
STARTUP_MODE=eager 232.2 s 99.5 s none none
STARTUP_MODE=fast 254.8 s 104.3 s 16.3 s 97 s (six sizes)
default 279.4 s 122.2 s 16.1 s vLLM's set

Findings, in order of how much they change the picture:

  1. Weight loading dominates the boot. 99 to 122 s of a 232 to 279 s total, at roughly 170 MB/s off local NVMe. No compile or capture flag touches it. On this class of hardware a warm page cache is worth more than the entire STARTUP_MODE feature.
  2. fast is slower than eager on GB10. CUDA-graph capture cost 97 s for six sizes (1,2,4,8,16,32), which swamps the 16 s saved by -O1. Capture appears to scale with the size and count of graphs. This inverts the expected ordering, and it is the single most useful thing this test produced: eager is the startup mode for this hardware, fast is for hardware where capture is cheap (published H100 figures put a trimmed set near 7 s).
  3. Both quantization paths are dead. --quantization fp8 fails at load with LoRA weight load: 168 parameter(s) absent from the checkpoint would be served UNINITIALIZED (the aLoRA adapters are not quantized and have no scale parameters). --kv-cache-dtype fp8 fails in input_quant_fp8.forward_cuda on assert input.ndim == 2 in the switch forward path. QUANT has been removed from the image; any value now stops the server at startup with that reason rather than failing deep in engine init. Removing a documented-but-always-failing option is better than shipping it.

Verified working on this hardware

  • The arm64 image boots and serves with the plugin, all three STARTUP_MODE values, and CUDA_GRAPH_SIZES.
  • The startup: banner appears in real boot logs: startup: mode=eager cuda_graph_sizes=vllm-default cache=/cache/vllm.
  • CACHE_DIR resolves and is writable (no fallback warning). The compile-cache hit benefit was not measured, since eager produces no compile artifacts to reuse.

What this changes in the PR

Commit 8d3533f removes QUANT, rewrites the README section with the measured table and the hardware-specific caveat, and updates the tests to assert that any quantization value is refused with the reason. Suite: 326 passed, 4 skipped.

Still open

  • No watchdogs tripped, peak MemAvailable never fell below 55.8 GB, and nothing was stopped on the host; the other engine on the box kept running throughout.
  • The absolute numbers are bounded-shape (4096 ctx, 4 seqs), so they are a proxy. The Hugging Face endpoint run remains the place for production-shaped boots, and it is also the only way to measure image-pull and platform-start overhead.

@planetf1

Copy link
Copy Markdown
Collaborator Author

End-to-end request check (same GB10, arm64 image, STARTUP_MODE=eager)

Booted, then sent real requests rather than stopping at /health:

POST /v1/decisions — HTTP 200 in 0.84 s

{"model":"granite-switch","answers":[
 {"type":"choice","name":"queue","choice":"payments",
  "probabilities":[{"value":"payments","probability":0.9669},
                   {"value":"account","probability":0.0331}],
  "confidence":0.9338},
 {"type":"predicate","name":"escalate","probability":0.8933}],
 "usage":{"input_tokens":123,"input_tokens_details":{"cached_tokens":16,"cache_write_tokens":80},
          "output_tokens":2,"total_tokens":125}}

That is the expected answer for the input used ("My payouts have failed three times. The bank says everything is fine."): the payments queue at 0.967 and escalate at 0.893. Prefix-cache counters are populated too, so --enable-prompt-tokens-details is doing its job.

POST /v1/chat/completions — HTTP 200 in 0.41 s, content "Ready".

GET /v1/models returns granite-switch with max_model_len: 4096, as configured.

So the image does not merely start: the plugin's routes, the adapter's single-token label read, and the shared chat route all work on this hardware.

On the quantization failure: not a regression

QUANT never shipped. It was introduced in 29ee7b5 and removed in 8d3533f, both on this unmerged branch. git log -S QUANT shows only those two commits, and git show <tag>:serve.sh for every tag from v0.1.0 to v0.2.0 matches nothing, as does origin/main. No released image ever had the option, so nothing regressed for users; this is a defect in new code caught before merge.

The underlying limitation appears to be undocumented:

  • The granite-switch 4.1 8b model card contains no mention of quantization, fp8, awq or gptq (fetched whole, 22,763 bytes).
  • No open or closed issue in generative-computing/granite-switch covers quantization; the only adjacent one is #128 (weights loading at ~1.87x on-disk size, unrelated).
  • Notably, granite-switch does ship tests/hf/test_quantization.py for its Hugging Face backend, so quantization is supported through transformers. The failure is specific to vLLM serving these checkpoints, where the aLoRA adapters meet the quantized linear layers.

Our README now documents both failure modes under "Quantization doesn't work here". Worth a follow-up issue upstream, since the assert input.ndim == 2 in input_quant_fp8.forward_cuda looks like a vLLM gap in the switch forward path rather than an inherent model limitation. Not filed as part of this PR.

@planetf1

Copy link
Copy Markdown
Collaborator Author

Correction to two claims in my previous comments

Both were under-verified when I wrote them. Re-derived from the logs and the shipped code.

1. "Graph capture cost 97 s for six sizes" — wrong.

That figure came from vLLM's capture progress bar, which does not isolate capture. From the log timestamps, engine init (profile + KV + warmup + compile + capture) was:

Arm weights load engine init capture configuration
eager 103.6 s 108.2 s none
fast 108.7 s 124.4 s (compile 16.3 s) [1,2,4,8,16,32], PIECEWISE, 6 sizes
default 126.6 s 132.5 s (compile 16.1 s) [1,2,4,8], FULL_AND_PIECEWISE, 4 sizes

So capture plus compilation costs ~16–24 s over eager, not 97 s.

2. "fast is slower than eager on GB10" — true for this test, but the test was rigged.

--max-num-seqs 4 pinned vLLM's own capture set to four sizes (max_cudagraph_capture_size: 8), so the baseline was already trimmed. Our fast preset then substituted six sizes up to 32, i.e. it captured more graphs than default. fast never got the situation it was designed for (a large default set), so the comparison says nothing about fast on a normal deployment.

3. Why the totals are not a usable ranking: weight loading varied 103.6 → 126.6 s across the three runs, unprompted, from page-cache and disk state. That spread (~23 s) is larger than every flag effect measured, so 232 vs 255 vs 279 s to ready is dominated by load noise. Only the engine-init column is a controlled comparison.

4. "Likely a vLLM gap" for the fp8 failure — I had this backwards.

vLLM 0.30's scaled_fp8_quant states the precondition explicitly: # This code assumes batch_dim and num_tokens are flattened / assert input.ndim == 2. And vLLM 0.30 has no granite_switch.py in its model registry (granite.py, granitemoe*.py, granite_speech*.py, but no switch variant), so the switch forward on this path comes from the granite-switch plugin at granite_switch/vllm/, whose decoder code reshapes activations to 3-D around the attention projections.

The most probable reading is therefore that the plugin's vLLM backend hands unflattened activations to a vLLM quantized linear that documents a 2-D precondition — an integration gap on the granite-switch side, or a combination that was never supported, rather than a defect in vLLM core. Same shape of answer for the fp8-weights failure: vLLM's LoRA loader refuses a parameter set it cannot fully initialise, which looks like intended behaviour when online weight quantization meets adapters it does not quantize.

Net effect on this PR: unchanged. Neither quantization path can serve these checkpoints, so QUANT stays removed and the README's "Quantization doesn't work here" section stands. What changes is who to ask about it: the granite-switch plugin's vLLM backend, not vLLM upstream.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant