Configuration

Plugin Setup

To load the plugin, set the VLLM_PLUGINS environment variable before running vLLM:

export VLLM_PLUGINS=spyre_inference

spyre_inference names both the platform plugin and the general plugin: the first activates the platform, and the second registers the OOT custom ops and the Spyre model adaptations, including the Spyre Transformers backend (used for model_impl="transformers").

Usage

You can then use vLLM as usual:

from vllm import LLM

llm = LLM(
    model="ibm-ai-platform/micro-g3.3-8b-instruct-1b",
    max_model_len=128,
    max_num_seqs=2,
)

See the Examples page for more usage patterns.

Gemma-4: text-only use of a vision checkpoint

Every Gemma-4 repository carries a vision_config, so google/gemma-4-31B, google/gemma-4-26B-A4B, google/gemma-4-E2B and google/gemma-4-E4B load as Gemma4ForConditionalGeneration and build a vision tower — weights to load and graphs to warm up that a text-only workload never runs.

To use one of those repositories for text only, pin its decoder architecture:

llm = LLM(
    model="google/gemma-4-26B-A4B",
    hf_overrides={"architectures": ["Gemma4ForCausalLM"]},
    tensor_parallel_size=2,
)

That is also the configuration the tensor-parallel and compile e2e tests run these checkpoints under. A repository with no vision tower gets the override by default, so this is only needed for the multimodal ones.

Gemma-4: per-request image sizes

The vision tower is compiled per token budget at warmup, for the image size the engine runs at: max_soft_tokens from --mm-processor-kwargs, default 280. A request can ask for a larger size (560 or 1120) in its own mm_processor_kwargs: offline through LLM.chat or LLM.generate, and online only when the server runs with --trust-request-mm-kwargs.

Each larger size needs its own budget, compiled at warmup, so they are opt-in through encoder_cudagraph_token_budgets. To serve requests up to 1120:

vllm serve google/gemma-4-26B-A4B-it --tensor-parallel-size 4 \
  --trust-request-mm-kwargs \
  -cc '{"encoder_cudagraph_token_budgets": [64, 128, 256, 280, 576, 1152]}'

The list replaces the default budgets, so keep the small ones in it: 64 to 280 cover the engine's own 280, 576 covers 560 and 1152 covers 1120. Each budget is rounded up to a whole patch group. An image larger than the largest entry, as listed, is rejected with a 400 (VLLMValidationError offline) rather than compiled mid-request, and the encoder cache is raised to hold one at that size. An entry above max_num_batched_tokens (512 unless --max-num-batched-tokens sets it) needs chunked MM input, so --disable-chunked-mm-input with it fails at startup. Raising the engine-wide --mm-processor-kwargs '{"max_soft_tokens": 560}' instead makes every image that size.

Decoder compile buckets

The body pads the packed token count to the next compile_sizes bucket, and warmup dummies every bucket. The lm_head sits outside every body graph and compiles its own, so it needs the same treatment: it projects one row per sampled request, a width that would otherwise take every value in 1..--max-num-seqs as requests finish. Those rows pad onto the same buckets clipped to --max-num-seqs, and warmup projects each width, so no shape reaches the lm_head uncompiled. Pad rows are dropped before sampling.

Encoder / pooling compile buckets

Spyre compile is on by default (STOCK_TORCH_COMPILE, dynamic=False). Pass --enforce-eager to disable it.

Everything derives from one number, R — the token budget. It is --max-num-batched-tokens, capped at 2048 (the measured throughput argmax across pooling models), floored at --max-model-len rounded up to a power-of-two multiple of 64, and capped again at what --max-num-seqs sequences of that length could carry, then floored to a whole multiple of it so that every length divides R. --max-num-seqs is then lowered to R / 64 if it was higher, since no batch wider than that fits.

• Body (Linear / LN): one shape, R rows. Every pooling step pads to it. Fixing it is what keeps the attention kernels keyed on sequence shapes alone.
• Lengths: powers of two from 64 (one Spyre stick) up to --max-model-len, rounded up to a power-of-two multiple of 64. Every length is then 64 * 2^k, so each divides R and each rectangle covers the body exactly. SPYRE_ATTN_QUERY_BUCKETS overrides the ladder, rounded the same way.
• Attention, rectangular path: for each length L, one rectangle B = R / L. The runner pads every sequence to L and the batch to B, so Q/K/V are the grid: one reshape, one F.scaled_dot_product_attention, one store, no data movement inside the layer. Taken whenever num_seqs <= B.
• Attention, ragged path: for a batch too wide for any rectangle, Q/K/V stay packed and requests are split into groups — the requests sharing one padded length. Each group is one fused gather/attend/scatter keyed on (group width, extent), so a ragged step makes one kernel call per group rather than per request. Widths are powers of two up to B; a wider group is chunked into descending powers of two.

With --max-model-len 512 --max-num-seqs 32 --max-num-batched-tokens 2048 that is 23 shapes: one body, four rectangles ((64,32) (128,16) (256,8) (512,4), each exactly 2048 rows), and 18 group pairs. At --max-num-seqs 4 the group family is empty — no batch that narrow can miss the rectangular path — leaving five shapes. Those counts assume the rectangular path is opaque; traced into the block graph (see below), each rectangle is a block graph of its own rather than an attention kernel.

The choice is made once per step from the step's metadata, before the forward, and the runner counts it in spyre_encoder_stats (rect_steps / ragged_steps). The ragged path always runs behind the opaque attention op. When compiled and the head size fills whole 64-element sticks, the rectangular path is traced into the block graph instead, so each rectangle is its own block graph. In eager mode both paths stay opaque, with no block graph; a compiled, unpadded sub-stick head size also keeps both paths opaque, sharing one block graph. Pooling models such as granite-30m pad their native 32-wide heads to 64 before construction.

Compiled pooling warmup runs one dummy at the body shape. With the rectangular path opaque, the first attention call in it traces every declared rectangle and group pair, against that call's own tensors (a Spyre tensor's device layout is part of its cache key). With it traced in, warmup adds one dummy per rectangle, plus one forced onto the ragged path (when any group is reachable) that traces the group pairs. Eager pooling uses one short dummy and always takes the packed path.

Example:

vllm serve ibm-granite/granite-embedding-125m-english \
  --runner pooling --max-num-seqs 4 --max-model-len 512

Tuning buckets for padding

Bucketing trades warmup time for per-request padding. A request is padded up to the next bucket on each axis and the padding is masked out, so buckets far above your real shapes waste compute, while buckets that hug your workload cut that waste but add graphs to compile at warmup. Per-sequence attention pairs each query bucket only with KV buckets that fit its smallest real query length. Batched decode uses KV-length/num-sequences pairs; both recording paths also enforce the allocated page-capacity limit. Extra buckets can still multiply the recorded variants, so keep those lists short.

Decoder body (packed token count). Override the defaults with compile_sizes; the platform clamps --max-num-batched-tokens to the largest entry. A decode-heavy run at --max-num-seqs 8 rarely needs the full power-of-two ladder:

from vllm import LLM

llm = LLM(
    model="ibm-ai-platform/micro-g3.3-8b-instruct-1b",
    max_num_seqs=8,
    max_model_len=2048,
    compilation_config={"compile_sizes": [1, 8, 512]},
)

1 and 8 cover decode steps (one token per running sequence, up to 8); 512 is the prefill bucket.

The default prefill bucket is 512 tokens, but larger ones compile. For prompts of many thousands of tokens, a 1024-token prefill chunk shortens time to first token, on dense models such as granite-3.3-8b as well as on MoE models such as gemma-4, where each prefill step reads every expert's weights however many tokens it carries. Pass a larger --max-num-batched-tokens to get one; the default buckets then keep 512 too, so short prefills don't pad up to the larger one:

vllm serve google/gemma-4-26B-A4B -tp 4 --max-model-len 32768 \
    --max-num-batched-tokens 1024 \
    --limit-mm-per-prompt '{"image": 0, "video": 0, "audio": 0}'

Warmup takes longer, and for prompts under a few thousand tokens it gains little; decode can get slightly slower. The example serves text only: a larger token budget also raises the number of image-batch widths the vision tower is warmed at, which can run out of device memory (#1302). At a long --max-model-len, a 2048-token bucket can push the largest attention variants past torch-spyre's 60 s backend compile limit at warmup, so they are left to compile on first use (torch-spyre#5250 makes that limit configurable).

Attention (KV length × query length). Set the buckets directly as comma-separated lists. Each is clamped to its limit: entries above --max-model-len (KV) or --max-num-batched-tokens (query) are dropped, and the limit is appended if missing, so every schedulable length keeps a bucket.

export SPYRE_ATTN_KV_BUCKETS=256,1024,2048    # default: the SPYRE_ATTN_KV_LADDER preset (pow2)
export SPYRE_ATTN_QUERY_BUCKETS=1,512         # 1 = decode; 512 = prefill chunk

The default KV buckets are the powers of two from block_size; the denser presets below union them with a further series. The ladder stays geometric rather than uniform because the recorded set grows with both the KV and num-sequences axes, even after filtering, so each extra entry costs warmup time. If your context never exceeds 2048, dropping the higher entries removes variants from warmup at no serving cost.

Choosing a KV ladder

SPYRE_ATTN_KV_LADDER picks how finely the KV ladder subdivides, without having to write the buckets out. The default is pow2. Every preset keeps the powers of two, so a KV length's bucket can only move down relative to pow2, never up. The denser presets do not contain each other, though: at block_size 128, 2600 rounds to 2688 under 8_5 but to 2944 under 4_3. Counts and overpay below are for block_size 128, at the --max-model-len each column names:

preset	spacing	buckets at 4096 / 16384 / 32768	mean overpay at 4096 / 32768
pow2 (default)	powers of two only	6 / 8 / 9	39% / 39%
8_5	8/5 ratio	10 / 15 / 18	21% / 19%
4_3	4/3 ratio	14 / 20 / 24	15% / 13%
9_8	9/8 ratio	20 / 34 / 41	8% / 6%
uniform	odd sixteenths of --max-model-len	13 / 15 / 16	14% / 14%

Overpay is the bucket over the real KV length, averaged over every KV length above block_size, so it describes a workload whose lengths are spread out. Warmup cost tracks the number of distinct recorded kernels, which rises with the bucket count. The ratio presets put their buckets at fixed token counts; uniform places them as fractions of --max-model-len, so its buckets move up with the context while its count stays roughly flat. Short contexts that are not a power of two are the exception: there the step rounds down to a whole block and adds a few buckets (16 at 2047 against 12 at 2048). Match --max-model-len to your real lengths before choosing it.

As a rule of thumb, a denser preset pays off with long KV lengths and large batches, because the saving is a number of KV tokens per sequence per decode step; at short lengths or a few sequences it buys little for its warmup. Once the block count exceeds one chunk of 32 // num_seqs blocks, the batched-decode kernel also pads it up to whole chunks, which takes back part of a denser preset's saving below 32 sequences. 8_5 is the recommended compromise, 9_8 the fastest when warmup time does not matter, and pow2 the default because it has the shortest warmup.

If the KV lengths are known and clustered, SPYRE_ATTN_KV_BUCKETS beats every preset: put the buckets where the lengths actually are. A workload that sits on one power of two gains nothing from any preset.

A denser ladder for large batches only

A denser ladder pays off in proportion to the tokens it saves times the batch size, so its gain concentrates at large batches. SPYRE_ATTN_KV_LADDER_LARGE_BATCH names a second preset that batched decode uses once the batch's num-sequences bucket reaches SPYRE_ATTN_LARGE_BATCH_MIN_SEQS (default 16, which with the default buckets means 9 or more sequences), while the per-sequence path and smaller batches keep SPYRE_ATTN_KV_LADDER. Below the threshold the dense buckets cost warmup without a measured gain. The second ladder is unioned with the first, so no length rounds up further. It is off by default, rejected at startup together with SPYRE_ATTN_KV_BUCKETS, which replaces the ladder it extends, and ignored with a warning when the threshold is above the largest num-sequences bucket. SPYRE_ATTN_KV_LADDER=8_5 with SPYRE_ATTN_KV_LADDER_LARGE_BATCH=9_8 gets close to plain 9_8 at large batches for noticeably less warmup.

With the batched-decode kernel enabled (SPYRE_BATCHED_DECODE=1, the default), warmup records eligible KV-length × num-sequences combinations. SPYRE_ATTN_NUM_SEQS_BUCKETS (default: powers of two from 1 to --max-num-seqs) is the extra lever there, and the same keep-it-short advice applies.

KV cache size

The KV cache is sized by vLLM's own GPU worker code: after the weights load, it reads the card's free memory (torch.accelerator.get_memory_info) and gives the KV cache what is left of --gpu-memory-utilization (default 0.92) × total. The cache keeps every block that memory pays for, usually far more than --max-num-seqs requests at --max-model-len need; with prefix caching on, the spare blocks keep earlier prompts' KV for reuse.

Each Spyre layer allocates its own page per block, so the budget is converted into blocks at the device cost of all layers' pages. That is more than vLLM's shared pool counts when groups would overlay one buffer, as a Mamba hybrid's do. Every layer's K and V pages are one tensor each, and the Spyre allocator never splits a tensor across its 16 GiB memory regions, so a model with few layers gets fewer blocks than its budget: they must still pack into the regions. When memory runs out before a full batch fits, a warning names the knobs, and requests wait for pages instead of all running at full length.

The profile run is skipped, as the first forward on Spyre would trace the model before the KV cache exists: a forward's activations come out of the 1 - --gpu-memory-utilization headroom. Lower it if a forward fails to allocate. --kv-cache-memory-bytes skips the measurement; --num-gpu-blocks-override replaces the block count it yields, so the warning above checks the override instead. VLLM_CPU_KVCACHE_SPACE, which sets the KV cache size on vLLM's CPU backend, is ignored, with a warning when it is set to a size; setting it to a size next to --kv-cache-memory-bytes is an error.

FP8 KV cache

Pass --kv-cache-dtype fp8 to store the KV cache in float8 (e4m3), halving the bytes each cached token takes:

vllm serve ibm-granite/granite-4.1-8b-fp8 --kv-cache-dtype fp8

It needs a head size that is a multiple of 128. Other quantized cache dtypes (fp8_e5m2, fp8_inc, nvfp4, ...) are rejected at startup. Since the KV cache is sized from free device memory (see KV cache size), an fp8 cache fits twice the tokens of an fp16 one. Attention still computes in the model dtype: pages are dequantized on read with the checkpoint's per-layer KV scales, or 1.0 when it has none. The speedup grows with the number of cached tokens each decode step reads, so it is small at short contexts; check accuracy against an fp16 cache before relying on it.

The scales are compiled into the graphs as constants. A checkpoint whose layers carry different KV scales therefore compiles the transformer block and each recorded attention kernel once per distinct pair of scales, which lengthens warmup; one without scales compiles them once, as an fp16 cache does.

Offline engine in the calling process

At tensor_parallel_size=1, set VLLM_ENABLE_V1_MULTIPROCESSING=0 to run the offline LLM engine in the calling process instead of a spawned EngineCore, which skips the spawn and the second round of imports at startup. vllm serve always spawns EngineCore.

The engine then changes the calling process: vLLM freezes the garbage collector's heap, caches its environment variables, seeds the global RNGs and leaves RANK/WORLD_SIZE set, the worker sets torch's thread count, and the card stays open until the process exits. Build an engine with tensor_parallel_size > 1 in a new process. Set OMP_WAIT_POLICY=PASSIVE before Python starts: OpenMP reads it when it loads, so the default the plugin sets only reaches spawned processes.

pyproject.toml Reference

The pyproject.toml includes several key build configurations:

Build Configuration

[tool.uv.extra-build-variables.vllm]
VLLM_TARGET_DEVICE = "empty"
CMAKE_ARGS = "--fresh"

[tool.uv.extra-build-dependencies]
torch-spyre = ["torch==2.14.1"]

These settings ensure:

• torch-spyre is built against the same PyTorch version (2.14.1) the runtime pins, so its C++ extension cannot drift in ABI
• vLLM is built with the empty backend — no device-specific C kernels. This avoids the torch-version coupling of prebuilt CPU wheels and the dependency on vllm._C (whose CPU-optimized ops we don't need; Spyre provides its own)

Source Repositories

The plugin pulls dependencies from specific Git repositories:

[tool.uv.sources]
vllm = { git = "https://github.com/vllm-project/vllm", rev = "..." }
torch-spyre = { git = "https://github.com/torch-spyre/torch-spyre", rev = "..." }

This ensures that torch-spyre and vllm are compiled/installed from source, instead of pulling pre-compiled wheels from PyPI.

PyTorch CPU Index

[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true

This ensures the CPU flavor of PyTorch is installed, as CUDA support is not required.