Kernels:

phanerozoic
/

model-as-a-kernel

Kernel card Files Files and versions

xet

Community

phanerozoic commited on Jul 13

Commit

95a427c

verified ·

1 Parent(s): f4ba638

model-as-a-kernel v1 source

Browse files

Files changed (2) hide show

README.md +122 -0
build.toml +2 -2

README.md ADDED Viewed

	@@ -0,0 +1,122 @@

+---
+library_name: kernels
+license: apache-2.0
+---
+# model-as-a-kernel
+An entire llama-family decode step as one kernel launch, loadable through
+`kernels`. A persistent phase-interpreter kernel executes the whole forward
+pass (embedding read, every layer's RMSNorms, QKV projection, rope, attention
+over the KV cache, SwiGLU MLP, the LM head, and the greedy argmax) inside a
+single launch, crossing a software grid barrier between phases instead of
+returning to the host. Greedy generation is closed-loop on the device: the
+argmax of one launch feeds the embedding read of the next, so a token costs
+exactly one kernel launch, no framework code, no host synchronization, and
+nothing on the CPU between tokens.
+## How it works
+At model load the Python side 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) and carry its pointers and sizes. One
+persistent kernel interprets the program; all resident blocks execute each
+phase cooperatively, then cross a sense-reversing global barrier. The launch
+geometry comes from the occupancy API, so every block is co-resident and the
+barrier cannot deadlock. A phase-per-launch debug mode runs the identical
+device code one phase at a time; it is bitwise identical to the fused mode by
+construction and gives per-phase cudaEvent timings.
+Arithmetic reproduces transformers eager bf16 semantics op for op: bf16
+storage with fp32 accumulation and one rounding per op boundary, the RMSNorm
+double rounding (normalize, round, weight multiply, round), fp32 softmax with
+the probabilities cast to bf16 before the PV product, bf16 rope multiplies
+with cos/sin rounded to bf16, and scores rounded to bf16 before and after the
+1/sqrt(D) scaling, as eager attention does. Every reduction uses a fixed
+tree, so the kernel is bitwise deterministic run to run: no atomics touch
+floating-point data anywhere.
+## Usage
+```python
+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")   # any plain
+# llama-family checkpoint: SmolLM2, TinyLlama, Qwen3 dense, Llama without
+# rope scaling. Pass a transformers model object to reuse loaded weights.
+tokens = mm.generate(prompt_ids, max_new=128)     # one launch per token
+logits = mm.decode_step(token_id, pos)            # one launch, fp32 [V] logits
+logits = mm.prefill(prompt_ids)                   # sequential decode steps
+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`
+needs `transformers` only if you pass a repo id string.
+## API
+| Symbol | Purpose |
+|---|---|
+| `MegaModel.from_pretrained(model_or_id, device, max_seq, max_gen)` | pack a llama-family checkpoint into a program |
+| `MegaModel(weights, config, ...)` | direct construction from tensors |
+| `decode_step(token, pos, phased=False)` | one decode step, returns live fp32 logits `[V]` |
+| `prefill(ids)` | feed a prompt token by token |
+| `generate(prompt_ids, max_new)` | greedy generation, closed loop on device |
+| `decode_step_timed(token, pos)` | phase-per-launch mode with per-phase ms |
+| `ops.mak_run / mak_run_steps / mak_run_phased / mak_num_blocks` | raw program launches |
+Supported architecture: RMSNorm decoder blocks with rotary attention (GQA,
+optional Qwen3-style per-head qk RMSNorm), SwiGLU MLP, no attention or MLP
+biases, plain unscaled rope, batch size 1, single device. `max_seq` bounds
+the KV cache, `max_gen` the device-side token output buffer.
+## Correctness
+bf16 transformers are order-sensitive: transformers itself does not
+reproduce its own logits between a batched forward and cache-stepped decode
+(the two differ by whole bf16 ulps that amplify layer over layer), so network
+outputs are only defined up to that reproduction band, and the correctness
+contract has two parts.
+Op-level, checked bitwise against eager-semantics replication in torch:
+the embedding read, the fused rmsnorm+GEMV, rope against
+`apply_rotary_pos_emb` (cos/sin tables and rotated K exact), attention at
+S=1, and the layer-0 V cache across all positions are `torch.equal` to the
+reference.
+Whole-model, on SmolLM2-135M over a 96-position teacher-forced sequence:
+the deviation from cache-stepped transformers eager must sit inside twice
+transformers' own cache-vs-batch band with comparable argmax agreement.
+Measured on RTX 6000 Ada (sm89): transformers self-band 1.94 max-abs with
+95/96 argmax agreement; this kernel vs cache-stepped eager 2.50 max-abs with
+90/96 agreement. The fused single-launch mode, the phase-per-launch mode, and
+repeated runs are mutually bitwise identical.
+## Measured
+Greedy decode, SmolLM2-135M bf16, 128 tokens with a 32-token prompt,
+against `transformers` eager `generate` on the same device:
+| GPU | this kernel | transformers eager | speedup |
+|---|---|---|---|
+| RTX 6000 Ada (sm89) | 378 tok/s | 37 tok/s | 10.3x |
+Per-phase cudaEvent brackets at position 512 attribute the remaining time to
+the attention phases (76% in phase-per-launch mode), which are latency-bound
+at one block per head; the LM-head GEMV already streams at ~700 GB/s. The
+speedup comes from removing all inter-op framework overhead: at batch 1 a
+135M-parameter step is ~180 launches of microseconds-scale kernels under
+eager execution, and exactly one launch here.
+## Limits
+Batch size 1, decode only (prefill runs as sequential decode steps), greedy
+sampling in-kernel (route the logits to an external sampler for anything
+else), llama-family architectures without rope scaling or attention biases,
+one CUDA device at a time. The program bakes device pointers, so the
+`MegaModel` owns its weights and buffers for its lifetime.

build.toml CHANGED Viewed

@@ -1,5 +1,5 @@
 [general]
-name = "model_as_a_kernel"
 version = 1
 edition = 5
 license = "Apache-2.0"
@@ -16,7 +16,7 @@ src = [
   "torch-ext/torch_binding.h",
 ]
-[kernel.model_as_a_kernel]
 backend = "cuda"
 cuda-capabilities = ["8.0", "8.6", "8.9", "9.0", "10.0", "12.0"]
 depends = ["torch"]

 [general]
+name = "model-as-a-kernel"
 version = 1
 edition = 5
 license = "Apache-2.0"
   "torch-ext/torch_binding.h",
 ]
+[kernel.model-as-a-kernel]
 backend = "cuda"
 cuda-capabilities = ["8.0", "8.6", "8.9", "9.0", "10.0", "12.0"]
 depends = ["torch"]