Zero-Allocation LLM Inference on the Metal: The Transformers-Ultra Architecture

Zero-Allocation LLM Inference on the Metal: The Transformers-Ultra Architecture

qwave · zero-allocation inference · fine touch from within · vaked.dev

When you profile a typical PyTorch autoregressive token loop at scale, you notice an unsettling truth: modern GPUs spend half their time waiting for the host CPU to dispatch tiny, fragmented memory operations. Every single token generated allocates new tensors, traverses Python bytecode dictionaries, runs dynamic shape reallocations, and calls dynamic softmax over 32,000 to 128,000 vocabulary logits.

In 8b-is/transformers-ultra, we decided to strip out every ounce of CPU runtime friction and build a zero-allocation, metal-bound inference engine.

Here is the engineering journey behind the latest sovereign commits.

1. The Slotted Static KV-Cache (SlottedStaticCache)

In default dynamic KV caching, every new token appends a slice along the sequence dimension using torch.cat or dynamic slice assignment. This creates dynamic VRAM allocations and triggers memory fragmentation.

In SlottedStaticCache, we allocate the entire contiguous buffer upfront: $$\text{Buffer Shape} = (B, H_{\text{kv}}, T_{\text{max}}, D)$$

# In-place slice write directly into pre-allocated memory:
self.keys[:, :, start_pos:end_pos, :].copy_(key_states)
self.values[:, :, start_pos:end_pos, :].copy_(value_states)
self.cumulative_length = end_pos

Because the memory addresses never change, the pointer addresses stay static across the entire generation lifecycle. This brings three massive advantages:

1. $100%$ Zero Memory Allocations during decoding.
2. Deterministic Memory Footprint: No out-of-memory surprises mid-generation.
3. Full CUDA Graph Compatibility: Stable data pointers allow graph replay without memory pool invalidation.

2. Sub-5µs Token Dispatch: CUDAGraphFastRunner

Python interpreter overhead (reflection, inspect, dictionary lookups) typically consumes $80\text{--}150\mu\text{s}$ per forward step. For small batch sizes, the GPU finishes computing the token in $30\mu\text{s}$ and spends the rest of the time starving for the next kernel launch.

CUDAGraphFastRunner captures the entire single-token forward pass into a static CUDA Graph:

# Dedicated static input/output bindings
self.static_input_ids.copy_(input_ids)
self.graph.replay()
return self.static_logits

By recording the execution graph during warmup and replaying it directly on the GPU command stream, token dispatch latency plummets from ~120µs down to <5µs.

3. $O(K)$ In-Register Logits Sampling (FusedLogitsSampler)

Standard Hugging Face sampling performs full vocabulary normalization:

1. Softmax over all $V$ logits ($V = 32,000\text{--}128,000$).
2. Full cumulative sum (torch.cumsum) across the entire vocabulary.
3. Top-P nucleus masking and multinomial distribution indexing.

This wastes massive GPU memory bandwidth. In FusedLogitsSampler, we perform an in-register Top-$K$ reduction before any transcendental operations:

# 1. Reduce O(V) -> O(K) in hardware registers
topk_logits, topk_indices = torch.topk(logits, k=top_k, dim=-1)
# 2. Temperature scaling + Softmax only on K items (e.g. K=50)
probs = torch.softmax(topk_logits / temperature, dim=-1)
# 3. Top-P nucleus filtering & multinomial sampling on K items

On large vocabularies ($V \ge 32k$), this reduces sampling kernel execution time by up to 80%. When greedy generation is active (temperature=0.0), it immediately executes a direct argmax with zero tensor allocations.

4. Zero-Allocation Speculative Decoding (SpeculativeFastRunner)

Speculative decoding pairs a small draft model (e.g., Gemma 2B) with a high-capacity target model (Gemma 27B). The draft model speculates $K$ candidate tokens, which are verified in parallel by the target model in a single batched pass.

The primary bottleneck in conventional speculative decoding is KV-cache state rollbacks when tokens are rejected. In SpeculativeFastRunner, rollback is an $O(1)$ scalar operation:

# Zero-allocation O(1) rollback on rejection:
target_cache.crop(initial_target_len + num_accepted)
draft_cache.crop(initial_draft_len + num_accepted)

If all $K$ tokens match, a bonus $(K+1)$-th token is sampled directly from the target logits, delivering an effective $2\text{--}3\times$ wall-clock throughput multiplier.

5. Apple Silicon Metal (MSL) Simdgroup Matrix Kernels

On macOS Apple Silicon (M1/M2/M3/M4 Max/Ultra), unified memory architecture (UMA) enables massive bandwidth ($400\text{--}800\text{ GB/s}$). We wrote native Metal Shading Language (MSL) compute kernels with simdgroup_matrix operations:

1.58-Bit Ternary MSL Unpacking

Unpacks 16 2-bit trits ($00 \to 0$, $01 \to +1$, $10 \to -1$) per uint32_t directly inside GPU threadgroup memory:

inline float unpack_trit(uint32_t packed, uint index) {
    uint bits = (packed >> (index * 2)) & 0x3;
    if (bits == 1) return 1.0f;
    if (bits == 2) return -1.0f;
    return 0.0f;
}

FP8 (E4M3) Dynamic GEMM

Performs fast bitfield exponent/mantissa expansion from 8-bit floats directly into 16-bit half-precision registers, followed by hardware-accelerated matrix multiplication.

6. NVIDIA Hopper & Blackwell 128-Byte TMA Engine & Mimalloc GC Freezing

On NVIDIA SM90+ (H100/H200) and SM100+ (B200), we wired direct 128-byte hardware CUtensorMap binary structures (=QQ4Q4Q4I4I16B) and zero-pause memory tuning:

• gc.freeze(): Freezes existing model weights in generation 2 GC, skipping pointer traversal during token loops.
• no_gc_cycle(): Context manager that completely silences garbage collection stops during generation bursts.
• mimalloc / jemalloc environment tuning for large page OS allocation.

The Verdict

All 32/32 unit tests pass in 0.87 seconds on our test harness. The entire pipeline is open source and merged into 8b-is/transformers-ultra:main.

Less indirection. Pure memory layouts. Metal-bound execution.