Inference Backend

Owns the model. The only place in the system that knows what a tensor is. Everything above it talks in token IDs.

The interface· Prefill primitives· Decode (batched)· Batch ops· Hardware detection· Torch backend· MLX backend· → KV Cache

The InferenceBackend interface

Defined in backends/base.py. Anything the scheduler does to "the model" goes through these methods.

Prefill primitives — what each one means

matched tokens (already cached) slice K and V tensors from the model's pool format into block-sized chunks copy each chunk into the BlockManager's pre-allocated layer pools at the allocated block IDs register in radix tree + exact-match index next time anyone asks for [t1..t7], full hit last_cache_hit_tokens backend stamps this on the response so the user (and metrics) can see how much was a cache hit.

decode_step_batched — the hot path

This is what runs every tick. Make it fast and the whole engine gets fast.

Batch operations (in-place row management)

Method	What it does
`splice_into_batched(batched_kv, new_kv, kv_len)`	Adds a row to the batched KV (after prefill finishes). Left-pads to fit.
`remove_row_from_cache(batched_kv, row_idx)`	Drops a row in-place when a request finishes. The remaining rows stay aligned.
`stack_caches_left_padded(rows)`	Builds the batched KV from a list of per-row KVs at the start of decode_step.
`kv_length(kv)`	How many real (non-padded) tokens this KV represents.

Hardware detection

Auto-pick the best available accelerator at startup. Manual override via DEVICE env var.

Custom-forward Torch backend (M1–M2.4 + paged-attention kernel)

backends/custom_torch_backend.py + models/gemma4.py + models/paged_attention_kernel.py. Hand-written Gemma 4 forward — byte-identical to HF on short prompts, correct (incl. sliding window) on long. Scheduler-driven with paged KV, prefix sharing, and a Triton decode kernel.

Loads HF weights into our own GemmaForCausalLM via from_hf(). No transformers model class on the hot path.
Full Gemma 4 architecture: dual-head-dim attention (sliding=256 / global=512), QK-norm + V-norm, GeGLU MLP (tanh), dual RoPE, per-layer embedding gating, KV sharing across the last 20 layers, final-logit softcapping at 30.
Activated via BACKEND=custom-{cuda,mps,cpu} (the BACKEND env wins over device auto-detect). Falls back to TorchBackend by default.
Scheduler primitives (M2.4): prefill, prefill_lookup/chunk/store, decode_step_batched, splice_into_batched, remove_row_from_cache, kv_length — so ContinuousBatchScheduler drives it end-to-end (admit → prefill → batched decode → evict). In-flight KV = list[PagedKVCache], one paged cache per active row (not HF DynamicCache).
Batched decode: one forward over all rows. On CUDA, a Triton paged-attention kernel (paged_decode_attention) reads each row's K/V directly via its block table — no gather, no padding, FlashAttention-style online softmax; handles head_dim 256 and 512 (the FA2 ceiling) + GQA + scale=1.0, runtime per-sequence loop bound (no recompile). CPU/MPS fall back to a left-pad-gather + masked SDPA (BatchedPagedKVCache) as the portable reference. The scatter-append is one vectorized indexed write per layer (de-Pythoned). Net decode: ~20× over the original row-by-row at N=32 (~323 tok/s short).
Sliding-window attention: sliding layers (28/35, window 512) attend only to the last 512 keys, applied on all three paths (prefill SDPA explicit position mask, gather-decode rightmost-W cutoff, kernel window arg). A conditional keeps the exact is_causal path when the window doesn't bite, so short-prompt parity stays byte-identical.
Windowed KV storage: sliding layers also free out-of-window blocks (not just mask them) — PagedKVCache._evict() releases a block once it's fully past the window and sentinels its block-table slot (which the kernel's start_b already skips, so no kernel change). Bounds physical KV on 28/35 layers for long context.
Window-aware admission: sliding pools are sized smaller (capped at the window via CUSTOM_BACKEND_SLIDING_BLOCKS), freeing memory to enlarge the binding full pools. kv_reserve/kv_release reserve per-pool block footprints (sliding capped at window//bs+1) and admit only if every pool fits — so the freed memory becomes more concurrent long-context requests (demonstrated: 1000-token workload to N=32 at ~equal memory vs the baseline's ~N=22 cap).
CUDA-graph decode: decode is dispatch-bound (~1000 tiny kernel launches/step). State lives in fixed GPU buffers (BatchedDecodeState), so the whole decode forward is captured once as a CUDA graph at max_batch_size and replayed each step (pad smaller batches up; one graph, not per-bucket — we're dispatch- not compute-bound). ~2× decode (50→25 ms/step), byte-identical to eager. Disable with CUSTOM_BACKEND_CUDA_GRAPH=0.
Op fusion — ruled out on A10G, REVERSED on A100 (2026-06-14): on A10G decode was memory-bandwidth-bound, so torch.compile-ing the leaf modules (GemmaRMSNorm/GemmaMLP) was a 1.01× no-op under the CUDA graph (those leaves hold no Triton calls → 0 breaks within them). On A100 the bottleneck moved to overhead/occupancy-bound, flipping that premise: whole-model torch.compile (CUSTOM_BACKEND_COMPILE=1, torch.compile(self.model, dynamic=False)) gave 1.47× (re-anchored to 1151 tok/s @ N=32, TPOT 23.8 ms). Measured (CUSTOM_BACKEND_EXPLAIN=1, 2026-06-14): the decode forward compiles to a single break-free graph — 0 graph breaks, 1 graph, 2704 ops. Modern Dynamo traces the @triton.jit attention kernels natively and unrolls the fixed-size paged scatter, so there are no fences — the 1.47× is whole-graph Inductor fusion, not islands. (Corrects an earlier unverified claim that whole-model compile breaks per layer.) CUSTOM_BACKEND_EXPLAIN=1 runs torch._dynamo.explain at capture and prints the break-count-by-reason. Remaining levers are within the single graph: max-autotune-no-cudagraphs (GEMM autotune) and cutting the 2704-op count via more aggressive fusion — NOT custom-op break-killing (no breaks exist). See HANDOFF + tuning_log.
Caches internally via its own PrefixCache + paged pools, not the HF-format CacheManager: set_cache_adapter is a no-op, so the scheduler's KV-pressure gate + release are skipped (admission is bounded by MAX_ACTIVE_KV_TOKENS instead; pool exhaustion → request rejection). Block release happens on row eviction via remove_row_from_cache → PagedKVCache.free_all().

HF-Transformers Torch backend (original)

backends/torch_backend.py. Wraps a HuggingFace causal LM.

Loads the HF model + tokenizer at startup. Uses HF's default sdpa attention (FA2 was ruled out — Gemma 4 head_dim exceeds FA2's 256 ceiling on A10G; H100+FA3 would lift this). Optional torch.compile via COMPILE_MODEL env var; currently off in the Modal image (Gemma 4 + dynamo collide in HF's lm_head slice path).
Implements every continuous-batch primitive listed above.
Holds a reference to the CacheManager so prefill can do lookup/store directly.
Uses HF's DynamicCache as the in-flight KV format. kv_cache/hf_format.py handles the conversion between our paged blocks and HF's flat tensor layout.
Owns a single backend._lock that the scheduler holds for the entire tick — serializes legacy generate/stream calls against the batched path.

MLX backend (Apple-native, limited)

backends/mlx_backend.py. Built on mlx_lm.stream_generate.

Known limitation: mlx_lm.stream_generate owns its own internal KV cache and exposes no hook to plug ours in. The MLX backend accepts a session_id but ignores cache lookup. Continuous-batch primitives are not implemented here. MPS via the torch backend is the primary path on Apple Silicon. A real MLX continuous-batch backend would have to be written directly against mlx.core — listed as a future extension in CLAUDE.md.

See also: Scheduler for who calls these methods · KV Cache for how blocks become DynamicCache and back.