model-as-a-kernel

A llama-family language model as one kernel launch, loadable through kernels. The reference baseline is transformers eager, matched op-bitwise where semantics are exactly defined and within transformers' own execution-path spread on whole-model logits.

An eager greedy generation of a small model is roughly 180 kernel launches per token, and at 135M parameters the launches, not the arithmetic, are the wall. This kernel runs the full forward pass, embedding, norms, projections, rope, attention over the KV cache, MLP, LM head, and greedy argmax, inside one persistent phase-interpreter kernel that crosses a software grid barrier between phases instead of returning to the host, and an in-kernel step loop carries the barrier across token boundaries. An entire generation, prompt consumption included, is a single cudaLaunchKernel, and decode runs an order of magnitude past the eager loop.

SmolLM2-135M decoding live at true relative speed (clock slowed 6x): the eager loop types at 39 tok/s behind thousands of kernel launches; the persistent kernel types at 805 tok/s behind exactly one, with repeated runs token-identical.

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
streams = mm.generate_batch([ids_a, ids_b], max_new=128)   # batched decode

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark.

API

Supported: RMSNorm decoder blocks with rotary attention (GQA, optional Qwen3-style qk RMSNorm), SwiGLU MLP, plain unscaled rope, single device; SmolLM2, TinyLlama, Qwen3 dense, and Llama-class checkpoints without rope scaling. Weights dense bf16 or bitsandbytes nf4.

Method

Model load packs the weights and compiles the architecture into an int64 [n_phases, 16] program naming an op per row 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, with the grid sized from the occupancy API so the barrier cannot deadlock. Arithmetic reproduces transformers eager bf16 semantics op for op, with fixed reduction trees and no floating-point atomics, so outputs are bitwise deterministic. GEMV weights stream through a per-warp cp.async shared-memory ring; attention splits cached positions into 128-token chunks with the softmax combine folded into the O-projection's staging; prompts are consumed in chunks of up to eight tokens, bitwise identical to token-by-token consumption. A phase-per-launch debug mode is bitwise identical to the fused mode.

Measured

Greedy decode, 128 tokens with a 32-token prompt, against transformers eager and a compiled decode loop (torch.compile reduce-overhead, StaticCache, CUDA graphs):

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). Prefill: a 2040-token prompt costs 0.73 s against 2.7 s token by token, bitwise-identical logits either way. Batched decode scales Llama-3.2-1B from 256 tok/s at batch 1 to 1673 at batch 16 (6.5x aggregate), with each sequence's logits bitwise identical to decoding alone.

Correctness

Verified against transformers eager through get_kernel on L4, H200, and RTX PRO 6000, and from source on RTX 6000 Ada. Bitwise (torch.equal): the embedding read, fused rmsnorm+GEMV, rope tables and rotated K, attention at S=1, and the layer-0 KV cache along a 48-token sequence against DynamicCache. Whole-model teacher-forced logits sit within transformers' own execution-path spread (mean 0.12-0.15, argmax agreement 93-96 of 96, matching the spread between transformers' own eager and sdpa paths). Fused, phase-per-launch, and repeated runs are mutually bitwise identical, and the single-launch generation emits the same tokens as one launch per token.

Requirements and limits

• 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; 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.

References

Persistent-kernel programming and software grid barriers (cooperative groups); the launch-overhead analysis of small-model decode; transformers eager semantics as the reference.

License

Apache-2.0.