model-as-a-kernel

A llama-family language model as one kernel launch, loadable through kernels. A persistent phase-interpreter kernel runs the full forward pass (embedding, norms, QKV projection, rope, attention over the KV cache, SwiGLU MLP, LM head, greedy argmax) crossing a software grid barrier between phases instead of returning to the host, and an in-kernel step loop carries the barrier across token boundaries: each step's argmax feeds the next step's embedding read on the device. An entire greedy generation, prompt consumption included, is a single cudaLaunchKernel.

Usage

from kernels import get_kernel

mak = get_kernel("phanerozoic/model-as-a-kernel", version=1, trust_remote_code=True)

mm = mak.MegaModel.from_pretrained("HuggingFaceTB/SmolLM2-135M")

tokens = mm.generate(prompt_ids, max_new=128)     # one kernel launch, total
logits = mm.decode_step(token_id, pos)            # one launch, fp32 [V] logits
logits = mm.prefill(prompt_ids)                   # chunked prompt consumption
logits, ms = mm.decode_step_timed(token_id, pos)  # per-phase cudaEvent ms

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark. from_pretrained accepts a repo id (requires transformers) or a loaded transformers model.

API

Supported: RMSNorm decoder blocks with rotary attention (GQA, optional Qwen3-style qk RMSNorm), SwiGLU MLP, no attention or MLP biases, plain unscaled rope, batch size 1, single device. Covers SmolLM2, TinyLlama, Qwen3 dense, and Llama-class checkpoints without rope scaling. max_seq bounds the KV cache, max_gen the device-side token output buffer.

How it works

Model load packs the weights (bf16, contiguous) and compiles the architecture into a program: an int64 [n_phases, 16] tensor whose rows name an op (embed, fused rmsnorm+GEMV, qk-norm+rope+cache-append, attention, swiglu+GEMV, argmax) with its pointers and sizes. One persistent kernel interprets the program; all resident blocks execute each phase cooperatively and cross a sense-reversing global barrier. The grid is sized from the occupancy API, so every block is co-resident and the barrier cannot deadlock. A phase-per-launch mode runs the same device code one phase per kernel for debugging and cudaEvent timing, and is bitwise identical to the fused mode.

Arithmetic reproduces transformers eager bf16 semantics op for op: fp32 accumulation with one rounding per op boundary, RMSNorm's normalize/round/scale/round sequence, fp32 softmax with probabilities cast to bf16 before the PV product, bf16 rope multiplies with bf16 cos/sin tables, scores rounded to bf16 before and after the 1/sqrt(D) scale. Reductions use fixed trees and no atomics touch floating-point data, so outputs are bitwise deterministic.

GEMV weights stream through a per-warp cp.async shared-memory ring whose depth is sized to the device at launch, so tile loads stay in flight continuously across row boundaries instead of serializing on register hazards; input staging is 16-byte vectorized and begins after the first weight tiles are already flying, and while the grid barrier settles each block prefetches the next projection's first tiles into L2. The ring is pure scheduling: each lane consumes its tiles in order and the reduction trees are unchanged.

Attention splits cached positions into 128-token chunks with the softmax combine folded into the O-projection's input staging, where the per-chunk combine weights are computed once per block and consumed with vectorized reads; rope and the KV append are fused into the attention phase, and each chunk's K/V is copied through cp.async into the dynamic shared region (idle during attention) so scoring reads shared memory rather than a chain of serial global rows. Prompts are consumed in chunks of up to eight tokens with each weight tile feeding all chunk rows, which amortizes weight traffic across the chunk; chunked consumption is bitwise identical to token-by-token consumption.

Correctness

Verified against transformers eager through get_kernel on L4 (sm89), H200 (sm90), and RTX PRO 6000 (sm120), and on RTX 6000 Ada (sm89) from source. SmolLM2-135M carries the full gate set below; Qwen3-0.6B (qk-norm, D=128) and TinyLlama-1.1B (GQA 32/4) are certified on the parity band, layer-0 KV cache, and mode-equivalence gates, and gemma-4-E2B-it (gemma norms, per-layer embeddings, shared KV, sliding windows, gelu gating, softcapped head) on the parity band (31/32 argmax agreement on H200 and RTX PRO 6000) and mode-equivalence gates.

Bitwise (torch.equal): the embedding read, the fused rmsnorm+GEMV, rope (cos/sin tables and rotated K against apply_rotary_pos_emb), and attention at S=1 on every card; the layer-0 K/V cache along a 48-token sequence against DynamicCache (bitwise on sm120 and sm89, within one bf16 ulp on a 0.03% element fraction on H200). The fused mode, the phase-per-launch mode, and repeated runs are mutually bitwise identical, and the single-launch generation emits the same token stream as one launch per token.

Whole model, absolute logit deviation from cache-stepped eager over 96 teacher-forced positions (bf16 outputs are order-sensitive, so any two implementations disagree; deviations grow with depth):

transformers' own alternative execution paths (batched forward, sdpa attention), measured identically on the same cards, land at mean 0.06 to 0.12, q99 0.25 to 0.50, max 0.44 to 0.94, argmax agreement 93 to 96 of 96. Greedy sequences diverge at any argmax disagreement, as they do between transformers' own paths.

Measured

Greedy decode, 128 tokens with a 32-token prompt, against transformers eager generate and a compiled decode loop (torch.compile reduce-overhead, StaticCache, CUDA graphs, device-side greedy feedback) on the same device.

SmolLM2-135M bf16:

Llama-3.2-1B bf16:

Steady-state decode at position 512 on Llama-3.2-1B: 1.85 ms/token on H200 (1.62 TB/s effective weight bandwidth), 2.34 ms on RTX PRO 6000, 3.8 ms on RTX 6000 Ada (about 80 percent of the card).

Prefill: a 2040-token prompt costs 0.73 s against 2.7 s consumed token by token (RTX 6000 Ada, SmolLM2), bitwise-identical logits either way. An eager greedy generation of SmolLM2 is roughly 180 kernel launches per token; here the entire generation is one.

Batched decode

generate_batch runs many sequences through the step loop at once, with each step's per-sequence argmax feeding that sequence's next embedding, so a whole batched generation is again one launch. Every reduction is per-sequence, so a sequence's logits are bitwise identical whether it decodes alone or in a batch, and the batch is independent of how sequences are ordered. Up to eight sequences stage their input in shared memory; wider batches (to batch_max(), 16 on the llama-family path) route the large projections' input through a global scratch. Steady-state decode on Llama-3.2-1B (RTX 6000 Ada) scales from 256 tokens per second at batch 1 to 1673 at batch 16 (6.5x aggregate).

Scope

One CUDA device, greedy sampling in-kernel (route the fp32 logits to an external sampler for anything else). Llama-family architectures without rope attention scaling or attention biases. Weights are dense bf16 or bitsandbytes nf4 (nf4 weights are stored packed and dequantized in the kernel). Prompts prefill in chunks of at most eight tokens. The program bakes device pointers, so the MegaModel owns its weights and buffers for its lifetime.

License

Apache-2.0.