Plugin Architecture

spyre-inference is a vLLM out-of-tree (OOT) platform plugin that enables inference on IBM's Spyre AI accelerator. It integrates with vLLM's plugin system to replace key compute layers with Spyre-optimized implementations while preserving the rest of the vLLM execution pipeline.

System Overview

The diagram below shows how spyre-inference fits into vLLM's process architecture. Blue boxes are Spyre-specific classes provided by this plugin; dark boxes are vLLM base classes; the gold box is the model loaded from vLLM's model registry with Spyre custom ops injected via OOT registration.

Process-level view of vLLM with the spyre-inference plugin. Dashed arrows (▷) indicate inheritance; solid arrows indicate composition or dependency.

The plugin registers via two entry points:

Entry Point	Target	Purpose
vllm.platform_plugins	spyre_inference:register	Registers TorchSpyrePlatform — sets dtype, worker class, attention backend, and distributed backend
vllm.general_plugins	spyre_inference:register_ops	Calls register_all() — importing the ops package triggers every @register_oot() layer swap, and register_all() additionally registers the spyre_convert custom op (RoPE registers no op — its rotation runs in-graph). Also overrides vLLM's TransformersForCausalLM with SpyreTransformersForCausalLM

vLLM is built from source with VLLM_TARGET_DEVICE=empty (no device-specific C kernels), so the platform overrides a few CPU-backend assumptions: import_kernels() is a no-op (there is no vllm._C), and the model runner reimplements the slot-mapping kernel in pure PyTorch. The host-side sampling kernels vLLM's CPU backend would supply are built by the plugin itself, as spyre_inference._C (see Host-side sampling).

Component view of a Granite model

Static architecture of the spyre-inference plugin showing how it integrates with vLLM and which model layers are replaced for Spyre execution.

Custom Op Replacement

Most layers that require Spyre-specific handling are replaced via vLLM's @ClassName.register_oot() decorator (a few, like SiluAndMul, need no replacement and run upstream in the compiled graph). Most replacements are pure class swaps that run when the ops package is imported. register_all() additionally registers the spyre_convert custom op — the convert helper keeps device transfers invisible to torch.compile. RoPE registers no op — its rotation-cache gather and rotation (2×2, or split-free for a sub-stick half) run directly in the compiled graph (see below).

vLLM Layer	Spyre Replacement	Device	Notes
GemmaRMSNorm	SpyreGemmaRMSNorm	Spyre	A maybe_compile(force=True) forward_native, so the fp32 promotion is kept. Plain RMSNorm needs no replacement — upstream's fp32 forward_native lowers eagerly — but Gemma's trailing fp32 weight multiply does not: a STANDARD [hidden] operand that torch-spyre can neither broadcast against a staggered-EA activation nor de-stagger, so the kernel has to be compiled even under enforce_eager.
RotaryEmbedding, Llama3RotaryEmbedding, YaRNScalingRotaryEmbedding, Gemma4RotaryEmbedding	SpyreRotaryEmbedding, SpyreLlama3RotaryEmbedding, SpyreYaRNScalingRotaryEmbedding, SpyreGemma4RotaryEmbedding	Spyre	Fully on-device, no opaque op. A [max_pos, 2, 2, rotary_dim//2] rotation cache is built from cos_sin_cache, flattened to 2-D with the position axis outermost, and primed on-device in _apply before torch.compile; forward_oot then gathers this pass's per-token slice with index_select and applies the 2×2 rotation-matrix formulation (_rotate_neox_2x2) — both traced directly into the enclosing block graph. Priming before compile is the requirement: building the cache lazily inside the traced forward segfaults libsenlib during warmup, whereas a cache already materialized on-device indexes cleanly. Only neox-style full rotary is supported — other configs raise NotImplementedError at construction. The 2×2 form views a head as two halves, so it needs a stick-aligned rotary_dim//2. A head whose half is sub-stick (head_size=64) takes _rotate_neox_split_free instead: x * cos + rotate_half(x) * sin, with rotate_half as a matmul by a fixed signed permutation, so every view is a whole head. Its cache is [max_pos, 2 * head_size], laid out [cos \| cos \| sin \| sin]. Head-dim padding then only has to make a head stick-aligned, so the native path pads to the next 64-multiple
VocabParallelEmbedding	SpyreVocabParallelEmbedding	Spyre (TP tables built on CPU at load)	The weight moves to Spyre with the model and the embedding gather runs on-device (aten.embedding now has a Spyre kernel, torch-spyre#420). TP=1 gathers directly. When TP>1, the per-vocab reindex/keep tables are built once on CPU at load and registered as device buffers; forward derives masked_input/keep from them on-device (index_select/F.embedding), applies the keep mask, and all_reduces — no per-step CPU round-trip
ColumnParallelLinear, MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear, ReplicatedLinear	SpyreColumnParallelLinear, SpyreMergedColumnParallelLinear, SpyreQKVParallelLinear, SpyreRowParallelLinear, SpyreReplicatedLinear	Spyre	All five swap in SpyreUnquantizedLinearMethod (the transposed-weight fast path below). SpyreQKVParallelLinear additionally asserts gather_output=False; SpyreRowParallelLinear (o_proj, down_proj) inherits upstream's all_reduce when reduce_results=True under TP>1
SiluAndMul	— (not replaced)	Spyre	No OOT class: vLLM's own SiluAndMul is traced into the compiled graph, so silu(gate)·up runs on Spyre and slices the fused [..., 2*d] on-device. The Spyre-specific piece is mlp_pad.py, which zero-pads intermediate_size to the 64-element stick at load time so that slice lands at a lowerable offset (inert since silu(0) = 0)
NewGELU	— (not replaced)	Spyre	No OOT class: vLLM's own gelu_new is traced into the compiled graph, cube term included — torch-spyre decomposes its torch.pow(x, 3.0) into a chain of mul ops (torch-spyre#4479)
ParallelLMHead	SpyreParallelLMHead	Spyre	TP≥1 with vocab sharding; per-rank weight padded to a multiple of 64×32 and pre-transposed; apply runs x @ Wᵀ then the un-pad slice on Spyre, compiled as its own graph (maybe_compile, see Decoder compile buckets) unless compilation is off; logits stay on Spyre for the TP all_gather
LogitsProcessor	SpyreLogitsProcessor	Spyre → CPU	Moves logits to CPU so all downstream sampling runs on the host. _apply_head D2Hs on the single-card path; when TP>1 _gather_logits runs the all_gather on Spyre and then converts the result. Either way the sampler's logits.to(torch.float32) never runs on Spyre, where it would crash torch-spyre's copy_from_d2d
GateLinear	SpyreGateLinear	Spyre	Clears out_dtype so MoE router logits stay in the weight dtype because Spyre cannot restickify fp32 (spyre::ReStickifyOpHBM is unsupported for IEEE_FP32). The MoE backend promotes stick-aligned full-softmax reductions to fp32 and returns them to the transport dtype
FP8 ColumnParallelLinear / fused QKV / MergedColumnParallelLinear / RowParallelLinear	SpyreFp8LinearKernel	Spyre	Checkpoint FP8 is DMA'd straight into qfp8wt at load (torch-spyre#4663), so model.to("spyre") skips it. quantscalepertokenfp8 + qfp8ch + spyre.scaled_mm trace into the block graph, or compile as their own graph (maybe_compile) outside one. The LM head stays FP16 unless a compressed-tensors checkpoint ships it in FP8, which SpyreFp8LMHeadMethod DMAs into qfp8wt as is, compiled even under enforce_eager

Transposed linear weights

F.linear(x, W) computes x @ Wᵀ however W is laid out, and on Spyre that transposed matmul is ~3.5× slower than a plain x @ A (torch-spyre#3512). SpyreTransposedWeightMethod (custom_ops/linear.py) is the shared base that closes that gap in two overrides:

• process_weights_after_loading replaces the loaded [out, in] weight with a contiguous [in, out] Wᵀ (optionally padding the output rows first), so the transpose is paid once at load time rather than every forward. A plain w.t().contiguous() runs on one thread and is slow on big weights, so by default (SPYRE_BLOCKED_TRANSPOSE=1) the same Wᵀ is built 128 source rows at a time, on torch's CPU threads (bit-identical).
• apply runs spyre_linear_t — torch.matmul(x, Wᵀ) plus optional bias — instead of F.linear, dropping any trailing pad columns with an on-device un-pad slice.

SpyreUnquantizedLinearMethod uses the base defaults (transpose in place, no padding); the five linear subclasses install it in __init__, but only when quant_method is an UnquantizedLinearMethod; quantized layers keep their own method. This is the pure-PyTorch equivalent of torch-spyre's [1,0] weight layout, which only fires for nn.Linear and so misses every vLLM parallel-linear. FP8 checkpoints use SpyreFp8LinearKernel (custom_ops/fp8_linear_kernel.py) rather than F.linear: load DMAs the checkpoint FP8 weight into qfp8wt, and the GEMM is qfp8ch + spyre.scaled_mm. SpyreUnquantizedLMHeadMethod reuses the same base with WEIGHT_T_ATTR="padded_weight_t" and ROW_ALIGN=64*32, so the fast path and the padding/un-pad logic are defined once.

Fused projections stay fused. SpyreQKVParallelLinear returns the whole [..., q+k+v] tensor and the unmodified upstream idiom q, k, v = qkv.split(...) slices it, exactly as SpyreMergedColumnParallelLinear's [..., 2*d] output feeds the upstream SiluAndMul, which slices gate/up on-device. Earlier revisions instead split the QKV weight on CPU at load time into three per-part GEMMs — a SplitQKV container built by an analyze_and_unfuse pass — so that no fused output ever had to be sliced; one fused GEMM is faster than three, so that pass is gone. Attention gathers query rows rather than passing offset views into compiled kernels, and materializes fused K/V views before cache writes (see Attention Backend).

Model adaptations

Some models need more than a swapped-out layer: a different transport for an input, a buffer that has to follow .to("spyre"), an expert dispatch Spyre can lower. Those live in spyre_inference/models/, one module per architecture, as subclasses of the upstream vLLM class rather than runtime monkey-patches. models/__init__.py holds spyre_models() (architecture string → Spyre class, built from the _ADAPTED_MODULES and _ADAPTED_ARCHS tables) and register_models(), which points vLLM's ModelRegistry at them; _-prefixed modules hold machinery those subclasses share or delegate to and register no architecture of their own. Registration is lazy — nothing is imported until vLLM resolves the architecture — and register_models() first checks every key against vLLM's own registry, so an upstream rename fails loudly instead of silently falling through to the unadapted class.

Where upstream hardcodes a class and offers no hook (the BERT wrappers hardcode embedding_class; gemma-4 builds its self-decoder inline), the already-built instance is retyped to its Spyre subclass — same __init__, same parameters and module tree, but a Spyre forward and armed compile state. Prefer a documented upstream extension point where one exists: CustomOp.register_oot / PluggableLayer.register_oot for a layer, and — for a MoE — the quant-method seam the unquantized oracle leaves open for an out-of-tree platform.

Three adaptations worth knowing:

• BERT / RoBERTa (models/_token_type.py) carry token_type_ids in a side buffer owned by the embedding instead of vLLM's bit-pack into the high bits of input_ids, which Spyre cannot unpack (torch-spyre#3509). Their word/segment/position gathers and layer norm can form a separate compiled embedding prologue when no enclosing graph traces them; RoBERTa's position offset is applied on the host before the gather.

• MoE (moe.py) supplies the routed-expert backend vLLM's unquantized MoE oracle lacks for an out-of-tree platform (it selects UnquantizedMoeBackend.OOT — no kernel — and leaves process_weights_after_loading to the plugin). A CustomOp.register_oot replacement for UnquantizedFusedMoEMethod computes the experts in two Spyre forms — gathered, which reads only a token's routed experts but lowers one token at a time, and all-expert persistent. Dispatch looks only at the packed token count: up to SPYRE_MOE_GATHERED_MAX_TOKENS — a decode batch, or a prompt that buckets that short — the gathered form is driven once per token, and above it the all-expert form takes the whole batch in one call. A multi-token batch whose rows are not stick-addressable takes the all-expert form whatever its size, since each row's storage offset has to span whole sticks for the compiled loop to address it; a single token is handed to the gathered form whole, so it never needs that. In the post-load hook it also rebuilds each layer's w13 [E,2M,H] / w2 [E,H,M] stacks into the [E,H,M] / [E,M,H] stacks those forms contract on (see MoE expert weights), freeing each source stack as it goes, since the device cannot hold both layouts at once. Tensor parallelism needs nothing further: upstream shards each expert's intermediate dim, so the forms just see a narrower M — zero-widened to whole sticks where a shard lands mid-stick — and MoERunner all-reduces the per-rank partial sums. Each model's own adaptation module supplies its recipe and any model-owned scaling (configure_gemma4_moe_layers in models/gemma4.py). Gemma4DecoderLayer.forward is untouched. SpyreMoERunner, an OOT MoERunner, takes vLLM's direct _moe_forward entry — as upstream does on CPU and TPU — instead of the opaque torch.ops.vllm.moe_forward custom op, so a compiled block traces the expert forms into its own graph; an eager run dispatches them to compiled regions.

• Gemma-4 per-layer embeddings (models/gemma4.py, E2B / E4B) retype the self-decoder to SpyreGemma4SelfDecoderLayers. Its upstream PLE lookup, including the in-vocab mask, always runs in a compiled region, including under --enforce-eager; an enclosing graph absorbs it rather than nesting a compile. The lookup and row-splitting methods cache their compiled callables independently. It cuts the backbone's per-layer rows in one compiled graph and hands each block its own offset-0 row through the [:, layer_idx, :] the backbone loop already does. A view of the packed tensor carries a different storage_offset per layer, which torch-spyre guards on, so every layer would compile its own block. The multimodal wrapper allocates its PLE cache on the host as a plain attribute, so models/gemma4_mm.py re-registers it as a buffer that follows the model to the device; that runs even text-only, since the default load keeps the wrapper with every modality at 0.

MoE expert weights

_prepare_layer moves each layer's experts to the device once, as three stacks, and hands the gathered form its own view of each: a reinterpret_tensor_with_layout alias over the same storage. Prefill and decode therefore share every byte of expert memory; there is no second weight pool.

Attribute	Shape	Device layout	Read by
spyre_moe_gate, spyre_moe_up	[E, H, M]	[E, H, M/s, s] (gather)	persistent
spyre_moe_gate_alias, spyre_moe_up_alias	[E·C, H/C, M]	[E·C, H/C, M/s, s]	gathered
spyre_moe_down	[E, M, H]	[E, H/s, M, s] (nn.Linear order)	persistent
spyre_moe_down_alias	[E·C, M, H/C]	[E·C, H/(C·s), M, s]	gathered

s is the stick (64 fp16 elements), M the intermediate dim zero-widened to whole sticks, and C is selected once during load by _derive_moe_chunks: the smallest divisor of H/s that spreads the top_k·C entries over the most cores. A core then holds top_k·H·M / cores weight elements whatever C is, so that split also has the smallest per-core working set. For Gemma 4 with top_k=8 and 32 cores this selects C=4. SPYRE_MOE_CHUNKS overrides the derived value; the override must still split H into whole-stick chunks. Alias entry e·C + c is chunk c of expert e. TP narrows only M, so the slicing is unchanged at every TP degree.

Decode orders the top_k·C entries row-major (row, then chunk): the entry index is the widened expert index times C plus a [1, C, 1] offset, both device tensors made at load, so the input is one broadcast, the gate/up chunk sum reduces a contiguous axis and the down parts reshape straight into H. Built from per-chunk torch.cats instead, the region splits into four kernels plus per-call constants, whose launch cost outweighs the narrower slices at TP4.

The gathered form slices the hidden dim because a gather divides work only along its entries. Gate/up slice their contraction dim (rows of H); down slices its output dim (columns of H). Persistent dispatch remains available for token batches above the gathered bound.

An alias is only valid when the sliced axis sits directly inside the expert axis in device order, so that each slice is one contiguous block and slicing merely relabels the flat order. That is what fixes each stack's layout. Gate/up keep H outer to the M sticks (the gather layout); down keeps the H sticks outer to M (the nn.Linear order its persistent matmul wants anyway — in the gather layout its long H rows would stream in short transfers). The other layout would not alias for either stack. The selected C guarantees H % (C·s) == 0 before any weights move, and each alias checks the pool's actual device layout before reinterpreting it.

The two forms differ numerically in one place: the gathered form sums gate/up over C fp16 partial products before the activation, so decode rounds up to three more times there than the persistent form. Down's output slices are concatenated, not summed.

Compilation Granularity

Under CompilationMode.STOCK_TORCH_COMPILE, compile_for_spyre compiles each entry of the model's block ModuleList in place via block.compile(backend="inductor", fullgraph=True, dynamic=False). In place matters: rebinding the list entry to the OptimizedModule that torch.compile returns would re-parent the block under an _orig_mod child and rename every parameter, breaking weight save/reload.

Every graph is compiled dynamic=False because torch-spyre's Inductor backend rejects SymInt shapes: a compiled graph is specialized to one concrete input shape. This is the root reason the plugin buckets shapes everywhere — variable request shapes are padded up to a small fixed set of compiled shapes (the padding masked out), and warmup pre-compiles every reachable bucket so no request pays an Inductor compile mid-serving. Two shape axes are bucketed independently: the packed num_tokens for the block graph (below) and (num_blocks, query_len) for attention (see Attention Backend).

Blocks are found structurally — a ModuleList whose non-PPMissingLayer entries own an Attention somewhere, and are not themselves Attention layers — so decoder stacks (model.layers) and encoder stacks (bert.encoder.layer) are both covered, as are hybrid Mamba+attention stacks that mix layer classes in one list. A ModuleList of bare Attention layers (Zamba2's shared dpa_list) is skipped: it is not a block stack. So is any list under a vision tower (vision_encoder, vision_tower, vision_model, visual), which stays eager unless marked (below) or compiled per encoder budget (see Vision encoders). Models whose attention is not a vLLM Attention — MLA (DeepSeek, Kimi) — match nothing and fall back to a whole-model graph.

A vision tower opts in. A block cannot assemble a constant on the host: a CPU-resident buffer has no device layout, and layout propagation rejects the graph. So a tower qualifies only once its block body builds nothing there: a load-time constant such as a rope permutation becomes a non-persistent buffer on the layer, lifted as a graph input, and a per-image one such as the attention mask is built once in the encoder forward and passed to each layer as an argument. A multimodal patch that has done that marks the layers with mark_blocks_compilable (multimodal/utils.py), and they are discovered like any other block stack; every other tower stays eager. CLIP's tower is marked (multimodal/clip.py): its attention is vLLM's MMEncoderAttention, which Spyre replaces out of tree (custom_ops/vit_attn.py) with a padded_sdpa call that traces into the block rather than running behind upstream's opaque SDPA wrapper, and its only mask is the key-padding row, which padded_sdpa keeps a graph input. Marked tower blocks compile beside the decoder blocks but never stand in for them, so a model with no decoder blocks still falls back to the whole-model graph. Warmup runs a marked tower at every image-batch width up to mm_max_items_per_batch, for pooling and generative models alike, and the compile guard watches its blocks like the decoder's: a tower compile after warmup is a shape warmup missed. A tower whose per-image shape varies beyond the batch width (a patch budget, an aspect ratio) has to make that shape enumerable before it is marked. Gemma 4's tower is not marked: SpyreGemma4ForConditionalGeneration (models/gemma4_mm.py) runs it on the vision encoder mixin below, compiled per token budget, and rejects video. Upstream's doubling covers the engine's max_soft_tokens. Larger per-request sizes get budgets only when listed in encoder_cudagraph_token_budgets (see per-request image sizes); TorchSpyrePlatform.validate_request rejects an image above the largest budget at the frontend, and the platform raises the encoder cache and compute budget to hold one. Several images pack into one call only up to the budget covering the engine's own largest image, since attention spans the whole packed sequence; a larger image runs alone at its own budget.

Blocks of one class share one forward code object, so Dynamo traces the first and the rest reuse that entry; whatever it re-traces hits the Inductor FX graph cache. The backend compile count is independent of depth, but it is not 1: layer 0 specializes separately because residual is None there, so a Llama-shaped stack yields two artifacts, and stacks that vary per layer yield more — Gemma 3 alternates sliding-window and full attention, giving four. A fresh num_tokens bucket then costs one block recompile rather than a whole-model one. For decoder attention, num_tokens is the block graph's only shape dependence: kv-cache length and its block-count buckets live inside unified_attention_with_output, which is opaque to this graph and compiles its own kernels (see Kineto profiling). Inline rectangular encoder attention also specializes the block graph by rectangle.

Depth independence relies on vLLM hoisting the per-layer attention name out of the graph, which needs torch >= 2.11 and VLLM_USE_LAYERNAME=1. Without it each block bakes in its own layer name and compiles separately, which is worse than the whole-model graph; the runner logs a warning when it detects this. Inductor freezing (enabled by max_autotune) defeats sharing the same way, by folding each block's weights into its own graph.

Embeddings and the final norm sit outside the block list, so no enclosing block graph covers them. Both stay eager when compilation is off, except a Gemma final norm: SpyreGemmaRMSNorm passes force=True, so it compiles as its own one-op graph in every mode, including enforce_eager, because its fp32 weight multiply has no working eager form. lm_head is never in a block graph either — compute_logits is a separate call on the wrapper — so its projection compiles its own graph, over row widths warmup pads onto (see Decoder compile buckets).

SPYRE_COMPILE_GRANULARITY=model restores the whole-model fullgraph, whose compile cost grows with layer count.

Vision encoders

A multimodal architecture gets a compiled, bucketed vision encoder by subclassing its upstream class with SpyreVisionEncoderMixin (models/_vision_encoder.py) ahead of it. The mixin implements upstream's SupportsEncoderCudaGraph protocol, so the bucketing is upstream's: EncoderCudaGraphManager picks power-of-two token budgets (or takes encoder_cudagraph_token_budgets), packs a step's images into them, pads the inputs and splits the output per image. SpyreEncoderGraphManager (v1/worker/spyre_encoder_graph.py) replaces only capture and replay: capturing a budget compiles the encoder at that static shape during warmup, and replay uploads the padded inputs and runs those graphs. Warmup compiles every budget, largest first, after the attention kernels, so the compile guard covers the vision graphs. Upstream's GPUModelRunner._execute_mm_encoder already sends encoder work through the manager, and the embedding merge, encoder cache and decoder are unchanged.

Its source is multimodal-vision-encoder.d2; regenerate it from docs/architecture/ with d2 --layout elk --pad 40 multimodal-vision-encoder.d2 multimodal-vision-encoder.svg.

Adding a model means one file in models/ (plus its _ADAPTED_ARCHS entry): the subclass implements the hooks below and nothing in the runner or manager changes. The units follow upstream: an image token is one position of vision_project's output, counted by upstream's token budgets and EncoderItemSpec.output_tokens, and host inputs hold one entry per patch, patches_per_token per image token.

Hook	Required	Purpose	Native Pixtral (models/pixtral.py)	Gemma 4 (models/gemma4_mm.py)
vision_buffers, patches_per_token, out_hidden_size	yes	host inputs (fp16, or integer indices) and their padding values; patches per image token; output width	patches, one-hot segments, positions; M²; text hidden size	patches, one-hot segments, x and y positions; k² (pooling kernel); text hidden size
vision_item_tokens(mm_kwargs)	yes	image tokens per item, for packing into budgets	patches / M² per image	valid patches / k² per image
vision_inputs(mm_kwargs)	yes	host buffers, unpadded; validated against vision_buffers	offset-major patches, segments, RoPE positions	valid patches only, offset-major by position; segments; x and y
vision_embed(values)	yes	graph 1: patches -> block input embeddings and shared block inputs	patch linear, ln_pre, image mask	patch linear plus position embeddings and RoPE, gathered per axis; image mask
vision_blocks()	yes	graph 2: blocks sharing one artifact	TransformerBlocks with SpyrePixtralAttention	stock encoder layers, their attention and norm instances switched to Spyre subclasses
vision_project(x, taps)	yes	graph 3: merge or pool, then project into the text space	projector norm, patch merge, adapter	the k×k pool only; the tail runs in postprocess_encoder_output
vision_call_block	no	block call convention (default block(x, **kw))	default	default
vision_taps	no	intermediate hidden states passed to vision_project	none	none
vision_pad_groups, vision_patch_multiple	no	padding layout and alignment (default: one group per merge offset, 64-wide sticks)	default	default (576 patches at k = 3)
vision_max_items_per_call	no	layout bound on items per encoder call	63 (one-hot width)	63 (one-hot width)
vision_items_limited_by_tokens	no	limit items per call by budget tokens and vision_max_items_per_call (then required) instead of upstream's smallest budget	False	False
vision_pack_tokens	no	packing stops at the budget covering this; larger items run alone (default None: up to the largest budget)	default	the engine's largest image
vision_modalities	no	modalities run through the graphs; others go to upstream's embed_multimodal	("image",)	("image",); video is rejected

Two upstream protocol methods are also the extension points for model-specific shapes, used as upstream models use them: override postprocess_encoder_output when placeholder embeddings differ from the encoder's image tokens (unpadding, newline tokens; its input is already on the host), and get_encoder_cudagraph_budget_range when the processor's largest item is not counted in the encoder's image tokens. The mixin owns the other protocol methods tied to the layout (replay buffers, item specs, capture inputs, forward), so an upstream class's own versions of those are always shadowed. select_encoder_cudagraph_items and get_input_modality come from the upstream class when it implements the protocol (Gemma 3/4, Qwen-VL), else from the model; upstream Pixtral does not, so SpyrePixtralForConditionalGeneration writes them. Gemma 4 overrides postprocess_encoder_output to run its tail on the host in fp32, as upstream does: the pooler's sqrt(hidden_size) scale can leave the fp16 range, so the scale, standardize and embed_vision never run on the card, and its graphs end at the pooled vision hidden size. The eager (enforce_eager) path runs the same hooks uncompiled. Upstream's processor, weight mapper, placeholders and embed_input_ids are inherited unchanged.

Constraints the layout follows:

• The doubling's largest budget is the largest item the model's upstream processor can produce (MultiModalBudget.mm_max_toks_per_item, the same number the scheduler's encoder budget uses, or the model's get_encoder_cudagraph_budget_range), so no item needs a shape warmup skipped; user budgets may go above it, and the platform rejects a request image above the largest. Each budget's patch count is a whole multiple of vision_patch_multiple.
• Patches are offset-major: one group per offset inside a merge group, each padded at its own tail. The patch merge is then a free view to [M², tokens, D], a batched matmul and a sum. Upstream's unfold regroup has no lowering, and a [L, D] → [L/M², M²·D] reshape feeding a matmul cannot be laid out.
• The image mask is segments @ segmentsᵀ over a one-hot [L, 64] segment matrix (the last column reserved for padding slots), not a column-minus-row broadcast, which torch-spyre cannot lay out.
• RoPE positions are uploaded as integers and cos/sin are gathered on the card from a GatherTable (placed rows-outermost, as the decoder's RoPE cache is), so the host uploads one index per patch. The rotation is a matmul against a constant pair-swap matrix, since a 64-wide head's 32-wide halves cannot be sliced.
• The encoder output comes back to the host in one copy and is sliced there: device views at offsets off the stick read as offset 0 (torch-spyre#3770), and the embedding merge runs on the host anyway.
• Tensor parallelism only; mm_encoder_tp_mode="data" is rejected.

Attention Backend

SpyreAttentionBackend is the decoder attention backend. It implements paged attention using pure PyTorch operations (no custom CUDA kernels). The KV cache holds two dense tensors per layer on Spyre, one each for K and V, both shaped [num_blocks, num_kv_heads, block_size, head_size] — the shape SpyreAttentionBackend.get_kv_cache_shape advertises. It runs a FlashAttention-style online softmax that iterates over pages without any compact-gather step, reading each page by indexing the dense tensor with an int32 device tensor (an indirect access, so the compiled bundle carries a real index rather than a constant slice). A page arrives in the shape the matmuls want, so no per-page permute runs before them. The transpose moves to the write instead: a token's KV heads are block_size rows apart, so the store is one index_copy_ per KV head (kv_write_index publishes one index per head) over a source materialized contiguously first.

Step	Device	Operation
1. Build metadata & masks	CPU → Spyre	The metadata builder reads query_start_loc/seq_lens, pads each query_len and KV block count onto their buckets, builds the per-sequence query-row index tables and the additive mask on CPU, then copies them to the device
2. Write new K/V to cache	Spyre	Compiled index_copy_ scatter, one per KV head, through a row view of the paged cache (fused into one kernel); only the row-index vectors are computed host-side and copied over
3. Page attention	Spyre	Eligible decode rows take a batched kernel, even for one sequence; prefill and ineligible decode rows use the host-driven per-sequence loop. Query rows are gathered on-device: Q @ Kᵀ · scale → optional soft-cap → + tile_mask → online softmax → @ V
4. Write-back	Spyre	Results are written into the output buffer in-kernel when possible, otherwise by a device-to-device copy

The compiled kernels themselves — the per-sequence page attention, the batched decode path, the KV store, and the cache's device layout — live under spyre_inference/v1/attention/ops/; the backend module holds the metadata builder and the host-side orchestration that calls them. The worker allocates each layer's cache through its impl (allocate_pages), so the shape, the layout and the kernels cannot disagree.

Cross-layer KV sharing

A layer with a kv_sharing_target_layer_name (gemma-4 E2B / E4B) allocates no cache: the runner aliases its target's pages in initialize_kv_cache_tensors. vLLM adds such a layer to its target's KV cache group but not to that group's UniformTypeKVCacheSpecs, so the per-layer spec lookup in initialize_attn_backend would KeyError; both runners hand the attention backend a copy (with_kv_sharing_layer_specs) that maps it to the target's spec, leaving the retained config's specs, and so its page_size_bytes, physical. YOCO fast prefill (kv_sharing_fast_prefill) swaps in an attention backend Spyre does not provide and is rejected.

Because attention kernels are dynamic=False too, they are pre-compiled during warmup rather than lazily on first use: by default (SPYRE_ATTN_RECORD=1) warmup records per-sequence variants from SpyreAttnBucketer.variants(). A query bucket pairs only with block counts that fit its smallest real query length; variants exceeding the allocated page count are skipped. Failed recordings are logged and compile on first use instead. When the batched-decode kernel is enabled (SPYRE_BATCHED_DECODE=1, the default) warmup enumerates its variants, the product of the KV-length (num_blocks) and num-sequences buckets; combinations that reach or exceed the allocated page count are skipped and declined at dispatch. A single step can carry a mix of prefill and decode sequences; each sequence is padded to its own query bucket (decodes use the length-1 bucket) before dispatch. SPYRE_ATTN_RECORD=0 restores lazy per-variant compilation.

Under dynamic=False the Python loop over a sequence's KV pages is unrolled at trace time, so a graph holds one copy of the attention body per page and compile time grows with KV length. SPYRE_ATTN_FOR_EACH_TILE=1 walks that axis with torch-spyre's for_each_tile instead, leaving one body plus a tile spec, and the same applies to the batched-decode kernel's walk over block chunks. The kernels carry the online softmax as a (tile_max, tile_sum, tile_output) triple either way. walk_tiles selects the walk; the same flag also controls fullgraph and the batched-decode page-index upload. SPYRE_ATTN_FOR_EACH_TILE=1 is the default; setting it to 0 uses Python loops as a rollback path. The flag is read once at import, so set it before importing spyre_inference.

LX-resident pages

The per-sequence decode kernel (page_attn_decode, one query row) keeps a gathered page in the LX scratchpad from its gather to its last use instead of round-tripping through HBM. A query wider than one row amortises the page transfer over every row, so it takes page_attn_prefill instead: batched GQA over the unfolded cache, one whole page per gather. Two shape choices get the decode kernel to LX residency. The page is gathered on (page, kv_head) with a [num_kv_heads, 1] index, so the gather's split lands per KV head — an output axis of probs @ V the consumer can mirror; behind a 1-D index the entry axis instead splits in whole 32-entry sticks. And the query groups fold into the query's row axis — a reshape, since heads are KV-major — so each matmul carries a single batch dim: the batched GQA form leaves the page with two batch dims and Inductor clones it out to a query-group axis it does not have (torch-spyre#4123). The cache fold is free too — [num_blocks, KV, block_size, D] reshapes to [num_blocks * KV, block_size, D] — but the cache is allocated with that folded axis at device dim 0, which is where an indexed axis has to sit for the gather to cost one page rather than the whole tensor.

Two things follow from those choices. The gather is a 2-D subscript, which lowers to aten.index and fails eager on Spyre by upcasting its int32 index (torch-spyre#4334, tracked by test_spyre_indirect_page_gather_subscript_needs_compile), so attention is always compiled on Spyre, even under --enforce-eager — attention compiles in its own domain, so the rest of the model still runs eager. Off Spyre an eager run calls the same kernels uncompiled, which is how the CPU-reference tests exercise them. Because the bmm's output axes (num_kv_heads * padded_query_len) cannot fill 32 cores at decode, and filling them would mean K-splitting a reduction a gather cannot mirror, the attention compile alone is capped at 8 cores; SPYRE_ATTN_MAX_CORES overrides that. The batched decode kernel's page index is uploaded one entry per stick, which is what the tiled walk needs.

Residency is a property of the layout plan, not of a result, so it is measured off the planner's own verdicts. Read those from the allocation — LifetimeBoundBuffer.address, where None means HBM — and not from the allocator's lx_pinning: debug log, whose → lx is the default for any buffer the spill-reason dict does not mention, and which never reports a graph-input clone at all. K's residency needs torch-spyre#4153: q @ Kᵀ lowers the transpose to a restickify, whose cross-frame barrier bars an LX-resident input without that PR's local-read proof. V is read directly by probs @ V and stays resident either way.

Batched decode needs none of this. Its whole-page gather is LX-resident as it stands, so the per-sequence kernel's (page, kv_head) fold has nothing to buy there: folding anyway divides blocks_per_chunk by num_kv_heads and multiplies the chunk trips, which measured roughly 1.5–2x more kernel time across the decode batches and contexts tried.

FP8 KV cache

--kv-cache-dtype fp8 (or fp8_e4m3) stores the pages as float8_e4m3fn, halving KV memory and the bytes each page read moves. Upstream sizes the cache as uint8 bytes and allocate_pages maps that back to e4m3. The store writes K / k_scale and V / v_scale, clamped to e4m3's range first because the float8 cast maps an overflow to NaN, and converts with spyre.qfp8ch: a plain .to has no device lowering and falls back to CPU, which writes bytes in sequential order while the fp8 -> fp16 read expects qfp8ch's order — so the pages are tagged QFP8CH in head_major_kv_layout. The three read kernels cast a gathered page back to the model dtype and apply v_scale to the output; k_scale is folded into the softmax scale the impl passes them. The attention math itself stays fp16.

The scales are per-tensor and per-layer, read from the layer's _k_scale_float / _v_scale_float (1.0 for a checkpoint without a kv_cache_scheme) at model load and again after reload_weights, which re-warms when they change, and passed to the kernels as Python floats, skipped entirely at 1.0. A KV-sharing layer takes its target's scales. A [1] scale tensor would let layers with different scales share one compiled block, but its broadcast fails torch-spyre's pointwise layout pass. Baking them in costs one attention variant set per distinct scale pair at warmup, and one block compile per distinct pair in the model graph: free for a checkpoint without KV scales, where every layer is 1.0, and up to one per layer for a calibrated one.

Key constraints:

• KV length bucketing: padded block count on power-of-two buckets from block_size to max_model_len (avoids per-step recompilation on Spyre); SPYRE_ATTN_KV_LADDER selects a denser preset. A layer with one KV head floors the ladder at two blocks: one KV head on one page leaves the decode kernel no spare dim for the stick
• Query length bucketing: [1] + multiples of min(512, max_num_batched_tokens) (consistent tensor shapes for compilation)
• Num-sequences bucketing (batched-decode kernel only, SPYRE_BATCHED_DECODE=1, the default): powers of two from 1 to max_num_seqs (SPYRE_ATTN_NUM_SEQS_BUCKETS); eligible (num_blocks, num_seqs) combinations are recorded during warmup
• Head size: Must be a multiple of 64 (128-byte Spyre stick ÷ 2-byte float16) — of 128 with an FP8 KV cache (1-byte elements)
• Block size: Must be a power of two and a multiple of 64. The default is 128; a user-supplied block_size is rounded up to the next multiple of 64, and one that is then not a power of two is rejected
• Head configurations: GQA, MHA (num_queries_per_kv = 1) and MQA (num_kv_heads = 1) are all exercised (tests/attention/test_spyre_attn.py)
• Supported: sliding-window masking (per layer, so a hybrid stack's full-attention layers stay unwindowed) and logits soft-capping
• Not supported: ALiBi

Encoder-only attention

Encoder-only (embedding) models take a separate path. For ENCODER/ENCODER_ONLY layers, TorchSpyrePlatform.get_attn_backend_cls selects SpyreEncoderAttentionBackend → SpyreEncoderAttentionImpl (both subclass the decoder backend/impl in spyre_encoder_attn.py). This path has no KV cache — attention is bidirectional over the full sequence — so it skips the paged-cache machinery. There is no online-softmax loop either: a sequence's whole K/V fits one tensor, so nothing forces the block-wise walk the decoder needs.

Everything hangs off one number, R = encoder_budget_rows(...): the token budget, capped at 2048, floored at max_model_len rounded up to a power-of-two multiple of 64, capped at what max_num_seqs sequences of that length could carry, and finally floored to a whole multiple of that longest length — which is what makes every declared length divide R. The pooling body is always R rows. Fixing it is what reduces the attention kernels' cache keys to the sequence shapes alone — both kernels take the body buffer as an argument, so a varying buffer size would multiply every attention graph.

On top of that one buffer sit two paths over a single power-of-two length ladder (ENCODER_LEN_ALIGNMENT = 64 doubling up to max_model_len, rounded up to a power-of-two multiple of 64 — the same rounding a request's own extent gets, so the ladder declares exactly the extents a request can be assigned):

1. Rectangular path — one rectangle per length, B = R / L, so a rectangle is exactly the body buffer. The runner pads each sequence to L and the batch to B in _preprocess (host-side integer tensors for input_ids; multimodal inputs_embeds are copied to CPU, padded there and copied back), so Q/K/V are the grid: _encoder_rect_kernel reshapes, runs one F.scaled_dot_product_attention, and stores — no data movement inside the layer. For packed-order pooling, unpad_encoder_grid compacts the grid at a fixed row count so the gather does not specialise per token total. Homogeneous CLS pooling instead reads each sequence's first grid row and skips the gather.
2. Ragged path — for a batch too wide for any rectangle. Q/K/V stay packed and requests are grouped by their own padded extent; _encoder_fused_kernel does gather, attend and scatter for one group in a single graph, keyed on (width, extent). Request boundaries ride in int32 row-index tables, so a card never does offset arithmetic on shapes — offsets are data. A group wider than the widest declared width is chunked into descending powers of two, so every dispatch lands on a warmed pair.

The runner picks between them once per step in _build_attention_metadata, through plan_encoder_step, and records the choice as the type of attn_metadata.encoder_plan (EncoderRectPlan versus a list of EncoderGroupPlan). It has to run there rather than in forward: the builder does a D2H read and H2D conversions, which inside a traced region become graph nodes.

The ragged path always stays behind the opaque unified_attention_with_output: it dispatches once per group, so tracing it in would make the block graph's structure depend on the step's group multiset. The rectangular path is traced into the block graph instead whenever every encoder layer qualifies — attention compiled, and a head size that fills whole sticks (install_encoder is the only gate). The runner publishes the rectangle's mask (publish_encoder_grid), and the layer's forward reads (width, extent) off that mask's shape, so inductor fuses attention with the projections and the FFN. The rectangle then is part of the block graph's shape: one block graph per declared rectangle, plus the opaque-attention one the ragged path runs. An unpadded 32-wide encoder layer and any eager run keep both paths opaque; pooling models such as granite-30m and MiniLM pad their original 32-wide heads to 64 before construction. Either way path selection is never a branch inside a compiled region: it is decided before the forward, and the traced graph only sees the shape it was handed.

Three torch-spyre constraints shape the rest: a compile input's storage_offset is a Dynamo guard (torch-spyre#4449, which closed #3770) and for int32 is still dropped outright, so rows are gathered with index_select rather than sliced — a slice would either recompile per offset or read the wrong rows; there is no on-device arange or full, so every index and mask tensor is host-built and converted separately (one mask for a rectangle, a row table and a mask per ragged group); and SDPA's decomposition does amax then exp(scores - max), which NaNs a fully masked row — hence the finfo.min / 2 mask fill and the single attendable key a batch-pad lane gets.

Encoder / embedding models: compile shape axes

The body always runs at R rows. Attention adds one rectangle per declared length on the rectangular path and one (width, extent) pair per group on the ragged one (a group being the requests that share one padded extent, attended together in one call). With max_model_len=512, max_num_seqs=32 and a 2048-token budget that is four rectangles and 18 group pairs; at max_num_seqs=4 the group family is empty, since no batch that narrow can miss the rectangular path.

Where those shapes are compiled depends on whether the rectangular path is traced in (see above). Opaque, the body compiles once and each rectangle and group pair is its own attention kernel — 23 shapes in the example above. Traced in, each rectangle is instead its own block graph, warmed by one dummy run per rectangle (_warm_encoder_inline_paths), plus the opaque-attention block graph whenever a group is reachable; the 18 group pairs remain attention kernels.

Encoder / embedding models under STOCK_TORCH_COMPILE, ragged path only: attention over the packed list, grouped by each request's padded extent. Predates the rectangular path and the single R-row body, so read the body bucketing and the warmup sweep as historical; the grouping and the row-index tables are still current. Regenerating it needs the d2 toolchain.

Device Placement Strategy

TorchSpyreModelRunner inherits from vLLM's GPUModelRunner and treats Spyre as the "GPU" in the CpuGpuBuffer pattern. Buffers are created via a SpyreCpuGpuBuffer override:

• Float dtypes: .cpu on CPU (numpy staging for the scheduler), .gpu on Spyre as float16
• Int / bool dtypes: .gpu aliased to .cpu (Spyre doesn't natively support these)

self.device stays cpu so that scatter, indexing, and block-table ops run on CPU, but float compute tensors land on Spyre via self._spyre_device. Because there is no vllm._C under VLLM_TARGET_DEVICE=empty, the runner also swaps in a pure-PyTorch _compute_slot_mapping_impl for the paged-cache slot-mapping kernel.

At load time, load_model moves every module except Attention scale buffers onto Spyre. The weight transposes happen before that, in each layer's process_weights_after_loading, while the weights are still on CPU.

_SpyreModelWrapper sits between the model runner and the model and converts at the call boundary:

• Input: CPU int32/int64 tensors → Spyre int64 (for embedding lookup)
• Output: Spyre float16 tensors → CPU (for logits indexing and sampling). A pooling model whose pooler runs on Spyre keeps its hidden states there instead, and the pooler copies back only the pooled vectors
• compute_logits: moves the CPU-sliced hidden_states[logits_indices] back onto Spyre for the SpyreParallelLMHead matmul, which returns logits on Spyre

SpyreVocabParallelEmbedding inherits weight loading and shard arithmetic from upstream and overrides forward. The weight moves to Spyre with the rest of the model, and the embedding gather runs on-device now that aten.embedding has a Spyre kernel (torch-spyre#420) — this replaces the earlier silent D2H/H2D CPU fallback that copied the full [vocab, hidden] weight on every decode step. When TP>1 the shard mask is applied on-device: get_masked_input_and_mask runs once at load to build per-vocab reindex/keep lookup tables (its int32 subtraction still cannot lower on Spyre), registered as device buffers; forward then gathers through them with index_select/F.embedding, applies the keep mask, and all_reduces — all on Spyre, no per-step CPU round-trip.

Hidden states flow on Spyre between decoder layers, with CPU round-trips only for work that stays host-side: logits indexing for sampling, and the attention metadata the builder prepares on CPU (slot mapping, the per-sequence index tables and additive mask). Eligible decode prefixes use batched attention, even at one sequence; otherwise the per-sequence loop is host-driven control flow, but its query-row gather and kernels run on Spyre. KV-cache writes and any separate write-back copy stay on-device. RoPE's rotation-cache gather and the embedding gather also run on-device.

Host-side sampling

SpyreLogitsProcessor hands the sampler CPU logits, so every sampling step runs on the host, serially between device forwards. vLLM picks its sampler path by platform: TorchSpyrePlatform is PlatformEnum.OOT, so is_cpu() is False, and TopKTopPSampler falls through to forward_native — eager PyTorch, with none of the fused kernels upstream reserves for forward_cpu. Left as is, that path sorts the full [batch, vocab] logits for top-k/top-p and generates a fresh [batch, vocab] noise tensor for every random draw, each a separate pass over memory. At a 262k vocabulary that is a large share of each decode step.

SpyreModelRunner patches upstream's Sampler in place after construction, pointing its greedy draw at greedy_sample (v1/sample/sampler.py) and its random draw at SpyreTopKTopPSampler, so every holder of it, the rejection sampler included, samples the same way. That brings back the CPU kernels and drops the full-vocabulary sort:

• Random draw — SpyreTopKTopPSampler.forward_native calls torch.ops._spyre_C.fused_gumbel_argmax: one pass of argmax(logits + g), with the Gumbel noise g hashed per element from a per-row seed (splitmix64 of (seed, i)), so there is no softmax and no noise tensor. The noise is bounded, so tokens too far below the running best (including every -inf left by top-k/top-p) are skipped exactly, and most others are rejected with one exp before the two logs of the Gumbel transform. A seeded request draws its row seed from its own torch.Generator, so it stays reproducible. use_fp64_gumbel keeps fresh fp64 noise, drawn in log space (argmax(x - log q), q ~ Exp(1)). vLLM's kernel reads its noise from a fixed 2^20-entry table instead, which limits each row to 2^20 noise windows and leaves much of a large vocabulary unreachable.
• Greedy — greedy_sample calls torch.ops._spyre_C.greedy_argmax, a vectorized max followed by a first-match scan, so ties break like torch.argmax. A single row falls back to torch.argmax, where OpenMP fork/join would cost more than the scan.
• Top-k + top-p — apply_top_k_top_p_sort_free takes a topk window of 2 × max_k instead of sorting the vocabulary, and runs top-p's softmax and cumsum over that window only. vLLM keeps every token tied with a row's k-th value, so the window widens to the exact tie count when ties run past it. Top-p without top-k still sorts.

The kernels live in csrc/cpu/sampling_kernels.cpp and are built with CMake from setup.py, over the cpu_types_{x86,vsx,vxe}.hpp SIMD headers copied unchanged from vLLM, so x86, POWER and s390x all get vector code. x86 ships an AVX512 build (_C) and an AVX2 fallback (_C_AVX2), and v1/sample/sampling_kernels.py imports whichever the host supports, as vLLM's CPU backend does. The ops register under the fixed _spyre_C namespace rather than TORCH_EXTENSION_NAME, so the two builds share one op name and cannot collide with a vLLM _C in the same environment. If the extension is missing or fails to import (a failed build, or a dev checkout that never built it), the import logs a warning and the samplers fall back to torch.argmax and a fresh-noise log-space draw. Because the sampler reimplements the tail of upstream TopKTopPSampler.forward_native, it has to be re-synced on a vLLM bump.

Transformers backend

When model_impl="transformers", register_ops swaps vLLM's TransformersForCausalLM for SpyreTransformersForCausalLM (spyre_inference/transformers_backend.py). vLLM's stock Transformers backend still handles model creation, weight loading, attention routing, the KV cache, and scheduling, and its fusers replace HF's linear/norm/GLU modules with vLLM layers — which the OOT registrations above then pick up automatically.

The subclass covers what upstream leaves to HF's module code. There is no RoPE fuser, so HF's rotary_emb survives and would derive cos/sin inside the forward from int64 position_ids, a cast torch-spyre cannot lower. It is replaced with a precomputed [max_model_len, 2, 2, head_dim/2] rotation cache — built on the host and moved to the device before compile, leaving only an index_select in the graph — plus an apply_rotary_pos_emb that applies the 2×2 rotation by multiply-and-reduce rather than HF's rotate_half slicing. That rotation still views a head as two halves, so here the platform widens head_dim to the next 128-multiple, not the native path's 64-multiple. The weight passes in head_pad.py are shared with the native path and pad Q/K interleaved, so this backend only has to rebuild the rotation cache at the pre-pad frequencies.

RMSNorm and GemmaRMSNorm both run upstream forward_native, compiled as their own graph wherever no block graph encloses them, which is the final norm after the last block; eagerly that norm launches one Spyre program per op. GemmaRMSNorm also compiles under --enforce-eager, because its trailing fp32 weight multiply hits a torch-spyre gap eagerly: a STANDARD [hidden] operand that can neither broadcast against nor de-stagger the staggered-EA activation. Because the fusers key on class names, the OOT registry covers the fused norms (SpyreTPAwareRMSNorm, SpyreTPAwareGemmaRMSNorm) too.

Distributed (TP)

TorchSpyrePlatform.get_device_communicator_cls returns SpyreCommunicator, a DeviceCommunicatorBase override in spyre_inference/distributed/spyre_communicator.py. The installed libspyre_comms.so now implements barrier, broadcast, send/recv, list-form allgather, gather, and allreduce; only reduce remains a throw-stub, and torch-spyre's spyreccl backend still stubs _allgather_base (so dist.all_gather_into_tensor doesn't work).

SpyreCommunicator therefore overrides:

• all_reduce — uses the functional _c10d_functional.all_reduce, which torch-spyre lowers to spyre::all_reduce_async inside a compiled graph and which runs eagerly via libspyre_comms outside one, so one code path serves both modes.
• all_gather — inside a compiled graph the functional collective lowers to spyre::all_gather_async and stays on device; eager keeps native list-form dist.all_gather, because the functional entry point (allgather_into_tensor_coalesced) is rejected outside a graph. CPU tensors route through the gloo half of the multi-backend cpu:gloo,spyre:spyreccl group.
• reduce_scatter — raises; it is not on the TP forward path.

gather is not overridden — it works natively via libspyre_comms. The tests/probes/test_spyre_comms_native_probes.py xfail-strict suite is the canonical signal: when a probe flips green, delete the corresponding override or workaround.

The worker (TorchSpyreWorker) inherits directly from vLLM's Worker (gpu_worker), not CPUWorker — Spyre needs none of the CPU-specific init (NUMA binding, host-RAM profiling). Data and pipeline parallelism (data_parallel_size > 1 or pipeline_parallel_size > 1) are rejected in check_and_update_config.