models/llama-3.1-8b-instructREADME.md11.0 kB
md
# Llama 3.1 8B Instruct
An 8.03B-parameter decoder with the Llama architecture, instruction-tuned:
grouped-query attention (32 query heads, 8 key/value heads, head dimension
128), rotary positions with base 500000 and `llama3`-style rope scaling
(context stretched from an 8192-token pretraining length to 131072), RMS
normalization, a SwiGLU MLP, and an untied output head, with a KV cache for
token-by-token decoding. The vocabulary is 128256 tokens.
The Linnet source is the same Llama package the other decoders in this zoo
use (`tinyllama-1.1b-chat`, `smollm2-1.7b-instruct`, `mistral-7b-instruct-v0.3`),
with `THETA` set to this model's rotary base and the generics set to its
width, head counts, and depth. What is new here is `crate.rope`'s rope
scaling: Llama 3.1 rescales each rotary frequency's *wavelength* in three
bands before computing angles, rather than using the raw base frequency
directly (`rope_scaling` in `config.json`, type `llama3`). Per frequency,
independent of position:
- a wavelength longer than `original_max_position_embeddings /
low_freq_factor` (8192) is divided by `factor` (8);
- a wavelength shorter than `original_max_position_embeddings /
high_freq_factor` (2048) is left alone;
- in between, the frequency is linearly blended between the unscaled and
the divided-by-`factor` value.
This is ordinary elementwise scalar arithmetic on the `[D / 2]` frequency
table (two nested `select`s on the wavelength comparisons, mirroring
`transformers`' `torch.where` chain in `_compute_llama3_parameters`); the
language's existing arithmetic and `select` were enough, with no new index
arithmetic involved, since the rescaling depends only on which frequency
`i` is being computed, never on the sequence position.
## Loading
```python
from linnet import nest
model = nest.load("llama-3.1-8b-instruct", backend="torch", numerics="fast")
logits = model(tokens) # Tensor[1, S; i32] -> Tensor[1, S, 128256; bf16]
```
The weights are the published, sharded `model-0000N-of-00004.safetensors`
(bf16); `bindings.json` maps Linnet parameter paths to their tensor names.
Every name, shape, and dtype is checked before anything runs.
Every `DecoderLayer`'s KV cache (its `state` members) is allocated eagerly
at load time, sized by `Batch` and `MaxSeq`, whether or not `decode` is
ever called -- so `MaxSeq` sets a memory cost you pay just to build the
model, not only to decode at that length. The card's default is
`MaxSeq = 8192`, Llama 3.1's own pretraining context before rope scaling
stretches it further: at that default the caches add about 1 GiB in bf16
(32 layers x 2 caches x 8 key/value heads x 8192 positions x 128 head
dimension) on top of the ~15 GiB of bf16 weights, so the model fits a
24 GB GPU with room for activations. Raise it for a longer decoding
context with `generics={"MaxSeq": 131072}` for the full rope-scaled
context -- but that alone is about 16 GiB of cache in bf16 (twice the
model's own weight size), so plan for roughly 31 GiB total and a GPU
larger than 24 GB. Cache size scales linearly with `MaxSeq`, so a value
in between (`16384`, `32768`, ...) trades context length for memory
directly.
## Entries
| Entry | |
| --- | --- |
| `forward<B, S>(tokens)` | logits for a whole sequence |
| `next_token<B, S>(tokens)` | logits for the last position |
| `decode(token, pos)` | one token through the KV caches (`Batch = 1`, `MaxSeq = 8192` by default) |
| `prefill<S>(tokens, pos)` | a whole prompt through the KV caches in one pass; logits after its last token, `decode` continues at `pos + S` |
| `prefill_slots<M, S>(tokens, slots, lengths)` | `M` requests' prompts into rows `slots` of the caches in one pass (each padded to `S`, its first `lengths[m]` tokens real), as they join a batch being served |
| `prefill_packed<P>(tokens, rows, positions, segments, last)` | several requests' prompts packed end to end into one pass of `P` tokens, each token given its cache row, its position, and its prompt: no padding between prompts, and each prompt sees only itself |
| `decode_rows(tokens, positions)` | one token for every row of the caches, each at its own position: the step `linnet.serve` takes for continuous batching |
| `prefill_paged<P, Rows, Pages>(tokens, positions, rows, slots, last, table)` | for serving from pages: prompts packed end to end, each token written at its place in the pool and attending over its row's pages up to its position (chunks, shared prefixes) |
| `decode_paged<Rows, Pages>(tokens, positions, table)` | `decode_rows` for serving from pages: loaded with `Batch = 1`, the caches' one row is a pool of `MaxSeq` positions in pages of `PageSize` (64), and row `b`'s positions lie in the pages `table[b]` lists |
| `step_paged<P, Rows, Pages>(...)` | `prefill_paged` and `decode_paged` in one pass |
| `generate<Steps>(token, pos)` | greedy decoding in the graph |
| `sample<Steps>(token, pos, key, temperature)` | sampling with `std.random` |
| `generate_until<MaxNew>(token, pos, eos)` | decoding until an end token |
## Shards
`Shards` (1 unless bound) splits the model across devices: one shard holds
`Heads / Shards` query heads, `KvHeads / Shards` key and value heads and
`Inner / Shards` hidden units, with the KV caches to match, and the
attention's and the MLP's output projections end in
`std.nn.parallel::all_reduce`, the sum over the shards. `lm_head` holds
`Vocab / Shards` rows of the vocabulary, and `std.nn.parallel::all_gather`
sets the shards' slices of the logits side by side. On one device the sum
and the gather are the value itself.
`linnet.torch.load(..., tensor_parallel=mesh)` binds `Shards` to the mesh
size and gives each process its part of the checkpoint.
The training entries train split as well: each split computation reads its
input through `std.nn.parallel::shared`, and `loss_packed` and
`log_probs_packed` run over the vocabulary's parts
(`std.nn.loss::split_cross_entropy`, `split_token_log_probs`).
## Provenance
- Weights: [NousResearch/Meta-Llama-3.1-8B-Instruct](https://huggingface.co/NousResearch/Meta-Llama-3.1-8B-Instruct).
Meta's own [meta-llama/Llama-3.1-8B-Instruct](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct)
is the original release of the same weights, but it sits behind Meta's
manual gated-access approval; this card points at NousResearch's ungated
mirror instead so `check` can read the checkpoint headers and anyone can
load it without requesting access. The license is Meta's own regardless
of which repository serves the bytes: the Llama 3.1 Community License
(SPDX-style identifier `llama3.1`), not a license NousResearch grants.
- Code: [meta-llama/llama-models](https://github.com/meta-llama/llama-models).
- Paper: [The Llama 3 Herd of Models](https://arxiv.org/abs/2407.21783).
## Validation
Checked against `transformers` (`LlamaForCausalLM`) on CPU, `torch.no_grad()`,
under the project's memory discipline for multi-billion-parameter models.
- **Rope-scaling table, directly, no model (done locally).** Before
trusting a model-level comparison, `crate.rope::tables`' scaled
`inv_freq` (all 64 frequencies, the entire wavelength range from a
single-token to a 20-million-token period) was compared elementwise
against `transformers.models.llama.modeling_llama.LlamaRotaryEmbedding(config).inv_freq`
for this checkpoint's real `rope_scaling` -- no weights, no model, just
the frequency table, so it costs nothing to check at every position
band. This caught a real bug: the smooth-interpolation band (frequency
indices 29-34 of 64, wavelengths 2048-8192) had its two blend weights
swapped, `smooth[i]` multiplying the scaled term instead of
`(1 - smooth[i])` -- confirmed by a boundary check (at
`wavelen == high_freq_wavelen`, `smooth = 1`, and the correct blend must
equal the *unscaled* frequency for continuity with the band just beyond
it; the swapped version gave the scaled one instead). Fixed in
`src/rope.linnet`; re-checked, `inv_freq` now matches to **5.96e-8**
(float32's own precision floor) at every one of the 64 frequencies. The
cos/sin tables built from it also match to 4.69e-7 at short positions,
widening to 7.8e-3 at position 131071 -- but by then the worst-matching
dimension is a high-frequency one the scaling never touches, so that
residual is ordinary f32 rounding on a large `position x frequency`
angle, present in any rope implementation at that length, not a second
scaling bug.
- **Truncated, f32, tight tolerance (done locally).** Both sides built from
the real checkpoint with only the first 2 of 32 layers
(`generics={"Layers": 2, "MaxSeq": 16, "T": "f32"}` on the Linnet side,
`num_hidden_layers=2` on the reference), prompt `"The capital of France
is"`. `MaxSeq` is overridden down from the card's default -- the KV
cache `state` members are allocated at load time regardless of `Layers`,
so leaving `MaxSeq` at its default would still have cost several GiB of
zeroed cache for no reason. Run under `flock` and a `systemd-run --user
--scope -p MemoryMax=12G -p MemorySwapMax=0` cap, peak RSS 11.58 GiB
(both sides, measured separately, well inside the cap). Before the rope
fix above this gave 4.97e-3, about a thousand times larger than the
sibling Llama-family cards' numbers on the same comparison (Mistral 7B
4.3e-6, Qwen3-8B 5.25e-6, Phi-3 7.15e-6) -- a real difference, not f32
noise, and the rope-table check above found exactly why. After the fix,
max absolute difference on the logits: **6.68e-6**, in line with those
siblings; top-1 token agrees exactly (token 105690, "Paris" continues
correctly). This proves the non-positional parts of the architecture
(GQA, SwiGLU, RMSNorm, the untied head) independent of rope scaling.
- **Whole model, bf16, loose tolerance, including a long-enough prompt for
rope scaling to engage (deferred).** Not run on this machine: holding a
32-layer, 8B-parameter model in bf16 (about 16 GiB) is exactly the kind
of run the project's memory discipline asks multi-billion-parameter
models to avoid on this shared 48 GiB CPU box, and the long prompt this
pass needs to push `position x frequency` past where a wrong `llama3`
scaling would diverge from a correct one makes it heavier still. This
card's whole-model and long-position numbers are deferred to a batched
GPU run (RunPod) that covers every model in the zoo together. The rope
scaling itself does not have to wait for that run, though: the direct
table check above already verifies it, including at long positions,
without loading the model at all.
The piecewise rule itself is `crate.rope::tables`' two nested `select`s
described above, over the wavelength comparisons `wavelen[i] >
low_freq_wavelen` and `wavelen[i] < high_freq_wavelen`, with the smooth
blend `(1.0 - smooth[i]) * (base_inv_freq[i] / FACTOR) + smooth[i] *
base_inv_freq[i]` in between -- so the GPU run can check it directly
against `transformers`' `_compute_llama3_parameters` at whatever positions
it covers, not just re-derive it from the constants.
Numbers are recorded in the change that added this model rather than
restated here, since they come from one specific run on one specific
machine; rerun the comparison above to reproduce them.bench.json20.7 kB
json
{
"date": "2026-09-29T15:13:03+00:00",
"environment": {
"python": "3.12.3",
"machine": "x86_64",
"torch": "2.14.0",
"jax": "0.11.2",
"transformers": "5.17.0",
"diffusers": "0.40.0",
"onnxruntime": "1.30.0",
"linnet": "0.1.0",
"gpu": "NVIDIA H100 80GB HBM3",
"gpus": "1",
"driver": "580.126.09"
},
"workload": {
"device": "cuda",
"prompt_tokens": 512,
"new_tokens": 128,
"batch": 1,
"warmup": 3,
"iters": 10,
"seed": 0,
"offload_gib": 8.0,
"serve_requests": 256,
"serve_concurrency": 64,
"serve_prompt_min": 128,
"serve_prompt_max": 512,
"serve_new": 128
},
"methods": [
{
"method": "transformers (eager)",
"kind": "reference",
"metrics": {
"ttft_ms": 22.157019935548306,
"decode_tok_s": 51.15306021144812,
"load_s": 5.2900608368217945,
"peak_vram_mib": 16430.9
},
"max_abs_diff": 0.0,
"notes": "",
"error": null,
"key": "transformers-eager"
},
{
"method": "transformers (torch.compile, static cache)",
"kind": "reference",
"metrics": {
"ttft_ms": 18.614680506289005,
"decode_tok_s": 110.13551214711164,
"load_s": 4.29446985758841,
"peak_vram_mib": 16476.0
},
"max_abs_diff": 0.0625,
"notes": "",
"error": null,
"key": "transformers-compile"
},
{
"method": "vLLM",
"kind": "reference",
"metrics": {
"ttft_ms": 16.496555879712105,
"decode_tok_s": 152.04963713106298,
"load_s": 74.5860153734684,
"first_token": 198.0,
"peak_vram_mib": 69060.9
},
"max_abs_diff": null,
"notes": "reserves a KV-cache pool up front (gpu_memory_utilization 0.85), so its memory is a setting, not a need",
"error": null,
"key": "vllm"
},
{
"method": "Linnet torch (generated source)",
"kind": "linnet",
"metrics": {
"ttft_ms": 17.403587698936462,
"decode_tok_s": 79.03901979336037,
"load_s": 3.589412832632661,
"peak_vram_mib": 16244.0
},
"max_abs_diff": 0.09375,
"notes": "KV cache compiled for 768 positions",
"error": null,
"key": "linnet-torch"
},
{
"method": "Linnet torch (CUDA graphs)",
"kind": "linnet",
"metrics": {
"ttft_ms": 13.043941929936409,
"decode_tok_s": 161.30111336854645,
"load_s": 3.1000381652265787,
"peak_vram_mib": 23314.9
},
"max_abs_diff": 0.0625,
"notes": "KV cache compiled for 768 positions",
"error": null,
"key": "linnet-cudagraphs"
},
{
"method": "Linnet torch (torch.compile, inductor)",
"kind": "linnet",
"metrics": {
"ttft_ms": 13.116896152496338,
"decode_tok_s": 140.95978163439494,
"load_s": 3.346582494676113,
"peak_vram_mib": 23310.9
},
"max_abs_diff": 0.0625,
"notes": "KV cache compiled for 768 positions",
"error": null,
"key": "linnet-inductor"
},
{
"method": "Linnet JAX (XLA, StableHLO)",
"kind": "linnet",
"metrics": {
"ttft_ms": 16.598870046436787,
"decode_tok_s": 166.6193060975316,
"load_s": 10.672614000737667,
"peak_vram_mib": 16703.0,
"driver_vram_mib": 33408.9
},
"max_abs_diff": 0.09375,
"notes": "KV cache compiled for 768 positions; peak memory is JAX's allocator peak; the driver shows its pool, which grows in whole regions",
"error": null,
"key": "linnet-jax"
},
{
"method": "Linnet JAX (XLA, generated source)",
"kind": "linnet",
"metrics": {
"ttft_ms": 15.713012777268887,
"decode_tok_s": 166.30785696783045,
"load_s": 10.5814537871629,
"peak_vram_mib": 16703.0,
"driver_vram_mib": 33406.9
},
"max_abs_diff": 0.09619140625,
"notes": "KV cache compiled for 768 positions; peak memory is JAX's allocator peak; the driver shows its pool, which grows in whole regions",
"error": null,
"key": "linnet-jax-source"
},
{
"method": "KerasHub (JAX)",
"kind": "reference",
"metrics": {
"ttft_ms": 27.488000690937042,
"decode_tok_s": 159.9121890574667,
"load_s": 31.004931453615427,
"peak_vram_mib": 17764.5,
"driver_vram_mib": 33404.0
},
"max_abs_diff": 0.09375,
"notes": "peak memory is JAX's allocator peak; the driver shows its pool, which grows in whole regions",
"error": null,
"key": "keras-hub"
},
{
"method": "Linnet ONNX f32 -> ONNX Runtime (CUDA)",
"kind": "linnet",
"metrics": {
"ttft_ms": 42.025355622172356,
"decode_tok_s": 53.66807237612467,
"load_s": 0.5022364966571331,
"peak_vram_mib": 36012.0
},
"max_abs_diff": 0.06537771224975586,
"notes": "f32; ONNX Runtime's CUDA execution provider; first calls build the sessions; logits copied to the host each step for the argmax",
"error": null,
"key": "linnet-onnx-f32"
},
{
"method": "linnet-onnx-f32-trt",
"kind": "linnet",
"metrics": {
"peak_vram_mib": 67034.0
},
"max_abs_diff": null,
"notes": "",
"error": "Fail: [ONNXRuntimeError] : 1 : FAIL : TensorRT EP failed to create engine from network for fused node: TensorrtExecutionProvider_TRTKernel_graph_main_3492707504992015285_0_0",
"key": "linnet-onnx-f32-trt"
},
{
"method": "Linnet ONNX f16 -> ONNX Runtime (CUDA)",
"kind": "linnet",
"metrics": {
"ttft_ms": 30.475120060145855,
"decode_tok_s": 70.99978329232621,
"load_s": 0.514266612008214,
"peak_vram_mib": 17930.0
},
"max_abs_diff": 0.0703125,
"notes": "f16; ONNX Runtime's CUDA execution provider; first calls build the sessions; logits copied to the host each step for the argmax",
"error": null,
"key": "linnet-onnx-f16"
},
{
"method": "Linnet ONNX f16 -> ONNX Runtime (TensorRT)",
"kind": "linnet",
"metrics": {
"ttft_ms": 28.629434294998646,
"decode_tok_s": 47.46078403900455,
"load_s": 0.4854988493025303,
"peak_vram_mib": 42964.0
},
"max_abs_diff": 9.7734375,
"notes": "f16; ONNX Runtime's TensorRT execution provider; first calls build the sessions; logits copied to the host each step for the argmax",
"error": null,
"key": "linnet-onnx-f16-trt"
},
{
"method": "Linnet ONNX bf16 -> ONNX Runtime (CUDA)",
"kind": "linnet",
"metrics": {
"ttft_ms": 31.128734350204468,
"decode_tok_s": 71.78795364072812,
"load_s": 0.4809684809297323,
"peak_vram_mib": 16920.0
},
"max_abs_diff": 0.09375,
"notes": "bf16; ONNX Runtime's CUDA execution provider; first calls build the sessions; logits copied to the host each step for the argmax",
"error": null,
"key": "linnet-onnx-bf16"
},
{
"method": "Linnet ONNX bf16 -> ONNX Runtime (TensorRT)",
"kind": "linnet",
"metrics": {
"ttft_ms": 28.49352080374956,
"decode_tok_s": 49.48714017482276,
"load_s": 0.488997345790267,
"peak_vram_mib": 40818.0
},
"max_abs_diff": 0.125,
"notes": "bf16; ONNX Runtime's TensorRT execution provider; first calls build the sessions; logits copied to the host each step for the argmax",
"error": null,
"key": "linnet-onnx-bf16-trt"
},
{
"method": "vLLM (offline, continuous batching)",
"kind": "reference",
"metrics": {
"serve_tok_s": 5449.406424204707,
"serve_s": 6.013131972402334,
"load_s": 134.7541569545865,
"peak_vram_mib": 70326.2
},
"max_abs_diff": null,
"notes": "256 requests of 128-512 prompt tokens and 128 new tokens, at most 64 in flight; KV-cache pool at gpu_memory_utilization 0.85",
"error": null,
"key": "serve-vllm"
},
{
"method": "transformers (generate_batch, continuous batching)",
"kind": "reference",
"metrics": {
"serve_tok_s": 881.3750614021327,
"serve_s": 37.17826999537647,
"load_s": 3.9765127412974834,
"peak_vram_mib": 81028.0
},
"max_abs_diff": null,
"notes": "256 requests of 128-512 prompt tokens and 128 new tokens, at most 64 in flight; paged|sdpa attention; reserves a paged KV-cache pool up front, so its memory is a setting, not a need; no per-request timestamps",
"error": null,
"key": "serve-transformers"
},
{
"method": "KerasHub (JAX, static batches)",
"kind": "reference",
"metrics": {
"serve_tok_s": 390.9685955480599,
"serve_s": 83.81235826388001,
"load_s": 30.097288263961673,
"peak_vram_mib": 26425.0,
"driver_vram_mib": 33396.0
},
"max_abs_diff": null,
"notes": "256 requests of 128-512 prompt tokens and 128 new tokens, at most 64 in flight; every batch runs to the longest possible prompt plus the new tokens; peak memory is JAX's allocator peak; the driver shows its pool, which grows in whole regions",
"error": null,
"key": "serve-keras-hub"
},
{
"method": "Triton Inference Server (vLLM backend)",
"kind": "reference",
"metrics": {
"serve_tok_s": 2997.927961265269,
"serve_s": 10.930215943604708,
"load_s": 61.096618136391044,
"serve_ttft_ms": 4545.441899448633,
"peak_vram_mib": 69134.4
},
"max_abs_diff": null,
"notes": "256 requests of 128-512 prompt tokens and 128 new tokens, at most 64 in flight; over HTTP, streamed, first tokens timed at the client",
"error": null,
"key": "serve-triton-vllm"
},
{
"method": "Linnet torch (linnet.serve, CUDA graphs)",
"kind": "linnet",
"metrics": {
"serve_tok_s": 5599.671676188662,
"serve_s": 5.851771656423807,
"load_s": 6.9916987195611,
"serve_ttft_ms": 2508.469292894006,
"peak_vram_mib": 28440.2
},
"max_abs_diff": null,
"notes": "256 requests of 128-512 prompt tokens and 128 new tokens, at most 64 in flight; 64 fixed cache rows of 640 positions",
"error": null,
"key": "serve-linnet-torch"
},
{
"method": "Linnet jax (linnet.serve, XLA)",
"kind": "linnet",
"metrics": {
"serve_tok_s": 4734.384229905995,
"serve_s": 6.921280235983431,
"load_s": 12.016882566735148,
"serve_ttft_ms": 3088.8994042761624,
"peak_vram_mib": 22580.7,
"driver_vram_mib": 33466.0
},
"max_abs_diff": null,
"notes": "256 requests of 128-512 prompt tokens and 128 new tokens, at most 64 in flight; 64 fixed cache rows of 640 positions; peak memory is JAX's allocator peak; the driver shows its pool, which grows in whole regions",
"error": null,
"key": "serve-linnet-jax"
},
{
"method": "Linnet onnx (linnet.serve, ONNX Runtime, f16)",
"kind": "linnet",
"metrics": {
"serve_tok_s": 3557.667502413826,
"serve_s": 9.210529083386064,
"load_s": 0.2594957249239087,
"serve_ttft_ms": 4340.559789445251,
"peak_vram_mib": 38456.0
},
"max_abs_diff": null,
"notes": "256 requests of 128-512 prompt tokens and 128 new tokens, at most 64 in flight; 64 fixed cache rows of 640 positions",
"error": null,
"key": "serve-linnet-onnx"
},
{
"method": "Triton Inference Server (Python backend, linnet.serve)",
"kind": "linnet",
"metrics": {
"serve_tok_s": 5376.75503814025,
"serve_s": 6.0943821631371975,
"load_s": 645.4065256789327,
"serve_ttft_ms": 2558.723919093609,
"peak_vram_mib": 29020.4
},
"max_abs_diff": null,
"notes": "256 requests of 128-512 prompt tokens and 128 new tokens, at most 64 in flight; over HTTP, streamed, first tokens timed at the client",
"error": null,
"key": "serve-triton-linnet"
},
{
"method": "Linnet -> vLLM",
"kind": "linnet",
"metrics": {
"ttft_ms": 12.238213792443275,
"decode_tok_s": 157.06471770669012,
"load_s": 39.27447564713657,
"first_token": 198.0,
"peak_vram_mib": 69060.8,
"driver_vram_mib": 69060.8
},
"max_abs_diff": null,
"notes": "linnet.hf.export (16 s), then vLLM on the exported checkpoint; reserves a KV-cache pool up front (gpu_memory_utilization 0.85), so its memory is a setting, not a need",
"error": null,
"key": "linnet-vllm"
},
{
"method": "Linnet -> vLLM (offline, continuous batching)",
"kind": "linnet",
"metrics": {
"serve_tok_s": 5607.039905193021,
"serve_s": 5.8440818246454,
"load_s": 34.699564065784216,
"peak_vram_mib": 70316.8,
"driver_vram_mib": 70316.8
},
"max_abs_diff": null,
"notes": "linnet.hf.export (0 s), then vLLM on the exported checkpoint; 256 requests of 128-512 prompt tokens and 128 new tokens, at most 64 in flight; KV-cache pool at gpu_memory_utilization 0.85",
"error": null,
"key": "serve-linnet-vllm"
},
{
"method": "Linnet -> llama.cpp (GGUF, f16)",
"kind": "linnet",
"metrics": {
"ttft_ms": 21.954108851338923,
"decode_tok_s": 170.921139,
"load_s": 75.52920866943896,
"peak_vram_mib": 15352.8
},
"max_abs_diff": null,
"notes": "linnet.gguf.export, then llama-bench with every layer on the GPU; the first token is the prompt at llama-bench's prompt rate",
"error": null,
"key": "linnet-llamacpp"
},
{
"method": "Linnet torch (offloaded)",
"kind": "linnet",
"metrics": {
"ttft_ms": 179.64914068579674,
"decode_tok_s": 5.7258174125537575,
"load_s": 19.60615223273635,
"peak_vram_mib": 9150.0
},
"max_abs_diff": 0.09375,
"notes": "KV cache compiled for 768 positions; cuda:0: embedding, layers.0-13; host, streamed in: layers.14-31, norm, lm_head",
"error": null,
"key": "linnet-offload"
},
{
"method": "Linnet torch (all GPUs)",
"kind": "linnet",
"metrics": {
"ttft_ms": 22.51227991655469,
"decode_tok_s": 87.83253758244847,
"load_s": 6.9974130764603615,
"peak_vram_mib": 23878.0
},
"max_abs_diff": 0.09375,
"notes": "KV cache compiled for 768 positions",
"error": null,
"key": "linnet-gpus"
},
{
"method": "Linnet torch (tensor parallel on 2 GPUs, shards)",
"kind": "linnet",
"metrics": {
"ttft_ms": 10.849429294466972,
"decode_tok_s": 238.08858552870754,
"load_s": 4.768562152981758,
"peak_vram_mib": 27757.0
},
"max_abs_diff": null,
"notes": "one process per GPU under torchrun, NCCL, CUDA graphs; KV cache compiled for 768 positions",
"error": null,
"key": "linnet-tp-torch"
},
{
"method": "Linnet JAX (XLA, tensor parallel on 2 GPUs)",
"kind": "linnet",
"metrics": {
"ttft_ms": 16.219401732087135,
"decode_tok_s": 180.8528170455239,
"load_s": 16.072626769542694,
"peak_vram_mib": 18801.2,
"driver_vram_mib": 38733.8
},
"max_abs_diff": 0.125,
"notes": "KV cache compiled for 768 positions; peak memory is JAX's allocator peak; the driver shows its pool, which grows in whole regions",
"error": null,
"key": "linnet-tp-jax"
},
{
"method": "vLLM (tensor parallel on 2 GPUs)",
"kind": "reference",
"metrics": {
"ttft_ms": 11.868055909872055,
"decode_tok_s": 232.93155993715257,
"load_s": 285.9878695830703,
"first_token": 198.0,
"peak_vram_mib": 142713.8
},
"max_abs_diff": null,
"notes": "reserves a KV-cache pool up front (gpu_memory_utilization 0.85), so its memory is a setting, not a need",
"error": null,
"key": "vllm-tp"
},
{
"method": "SGLang",
"kind": "reference",
"metrics": {
"ttft_ms": 21.038793958723545,
"decode_tok_s": 157.94771258807424,
"load_s": 55.7341874409467,
"first_token": 198.0,
"peak_vram_mib": 69639.6
},
"max_abs_diff": null,
"notes": "reserves a KV-cache pool up front (mem_fraction_static 0.85), so its memory is a setting, not a need",
"error": null,
"key": "sglang"
},
{
"method": "Linnet -> SGLang",
"kind": "linnet",
"metrics": {
"ttft_ms": 19.887838512659073,
"decode_tok_s": 157.86890601209936,
"load_s": 38.44177514873445,
"first_token": 198.0,
"peak_vram_mib": 69575.6,
"driver_vram_mib": 69575.6
},
"max_abs_diff": null,
"notes": "linnet.hf.export (18 s), then SGLang on the exported checkpoint; reserves a KV-cache pool up front (mem_fraction_static 0.85), so its memory is a setting, not a need",
"error": null,
"key": "linnet-sglang"
},
{
"method": "SGLang (offline, continuous batching)",
"kind": "reference",
"metrics": {
"serve_tok_s": 4775.564298230856,
"serve_s": 6.861597489565611,
"load_s": 55.798140831291676,
"peak_vram_mib": 71689.6
},
"max_abs_diff": null,
"notes": "256 requests of 128-512 prompt tokens and 128 new tokens, at most 64 in flight; KV-cache pool at mem_fraction_static 0.85",
"error": null,
"key": "serve-sglang"
},
{
"method": "Linnet -> SGLang (offline, continuous batching)",
"kind": "linnet",
"metrics": {
"serve_tok_s": 4756.81316730149,
"serve_s": 6.888645580038428,
"load_s": 44.909736623987556,
"peak_vram_mib": 71625.6,
"driver_vram_mib": 71625.6
},
"max_abs_diff": null,
"notes": "linnet.hf.export (0 s), then SGLang (offline, continuous batching) on the exported checkpoint; 256 requests of 128-512 prompt tokens and 128 new tokens, at most 64 in flight; KV-cache pool at mem_fraction_static 0.85",
"error": null,
"key": "serve-linnet-sglang"
},
{
"method": "Text Generation Inference",
"kind": "reference",
"metrics": {
"ttft_ms": 27.898590080440044,
"decode_tok_s": 112.42300659927503,
"load_s": 44.11180199868977,
"peak_vram_mib": 60612.9
},
"max_abs_diff": null,
"notes": "over HTTP, streamed; the prompt is the token ids decoded and tokenized again; reserves a KV-cache pool up front (cuda-memory-fraction 0.85), so its memory is a setting, not a need",
"error": null,
"key": "tgi"
},
{
"method": "Linnet -> Text Generation Inference",
"kind": "linnet",
"metrics": {
"ttft_ms": 25.694426149129868,
"decode_tok_s": 114.79205057180067,
"load_s": 40.10608994215727,
"peak_vram_mib": 60612.9,
"driver_vram_mib": 60612.9
},
"max_abs_diff": null,
"notes": "linnet.hf.export (19 s), then Text Generation Inference on the exported checkpoint; over HTTP, streamed; the prompt is the token ids decoded and tokenized again; reserves a KV-cache pool up front (cuda-memory-fraction 0.85), so its memory is a setting, not a need",
"error": null,
"key": "linnet-tgi"
},
{
"method": "Text Generation Inference (continuous batching)",
"kind": "reference",
"metrics": {
"serve_tok_s": 2621.3774893461728,
"serve_s": 11.87047654390335,
"load_s": 42.13999486528337,
"serve_ttft_ms": 3609.404830262065,
"peak_vram_mib": 62156.9
},
"max_abs_diff": null,
"notes": "256 requests of 128-512 prompt tokens and 128 new tokens, at most 64 in flight; over HTTP, streamed; stops at end-of-sequence (TGI cannot ignore it), so throughput counts the tokens produced; KV-cache pool at cuda-memory-fraction 0.85",
"error": null,
"key": "serve-tgi"
},
{
"method": "Linnet -> Text Generation Inference (continuous batching)",
"kind": "linnet",
"metrics": {
"serve_tok_s": 2565.842909470642,
"serve_s": 12.127008978277445,
"load_s": 40.06648236885667,
"serve_ttft_ms": 4157.670503482223,
"peak_vram_mib": 62706.9,
"driver_vram_mib": 62706.9
},
"max_abs_diff": null,
"notes": "linnet.hf.export (0 s), then Text Generation Inference (continuous batching) on the exported checkpoint; 256 requests of 128-512 prompt tokens and 128 new tokens, at most 64 in flight; over HTTP, streamed; stops at end-of-sequence (TGI cannot ignore it), so throughput counts the tokens produced; KV-cache pool at cuda-memory-fraction 0.85",
"error": null,
"key": "serve-linnet-tgi"
}
],
"reference": "transformers-eager"
}bindings.json21.8 kB
json
{
"embedding.weight": "model.embed_tokens.weight",
"norm.weight": "model.norm.weight",
"lm_head.weight": "lm_head.weight",
"layers.0.attention_norm.weight": "model.layers.0.input_layernorm.weight",
"layers.0.mlp_norm.weight": "model.layers.0.post_attention_layernorm.weight",
"layers.0.attention.q_proj.weight": "model.layers.0.self_attn.q_proj.weight",
"layers.0.attention.k_proj.weight": "model.layers.0.self_attn.k_proj.weight",
"layers.0.attention.v_proj.weight": "model.layers.0.self_attn.v_proj.weight",
"layers.0.attention.o_proj.weight": "model.layers.0.self_attn.o_proj.weight",
"layers.0.mlp.gate.weight": "model.layers.0.mlp.gate_proj.weight",
"layers.0.mlp.up.weight": "model.layers.0.mlp.up_proj.weight",
"layers.0.mlp.down.weight": "model.layers.0.mlp.down_proj.weight",
"layers.1.attention_norm.weight": "model.layers.1.input_layernorm.weight",
"layers.1.mlp_norm.weight": "model.layers.1.post_attention_layernorm.weight",
"layers.1.attention.q_proj.weight": "model.layers.1.self_attn.q_proj.weight",
"layers.1.attention.k_proj.weight": "model.layers.1.self_attn.k_proj.weight",
"layers.1.attention.v_proj.weight": "model.layers.1.self_attn.v_proj.weight",
"layers.1.attention.o_proj.weight": "model.layers.1.self_attn.o_proj.weight",
"layers.1.mlp.gate.weight": "model.layers.1.mlp.gate_proj.weight",
"layers.1.mlp.up.weight": "model.layers.1.mlp.up_proj.weight",
"layers.1.mlp.down.weight": "model.layers.1.mlp.down_proj.weight",
"layers.2.attention_norm.weight": "model.layers.2.input_layernorm.weight",
"layers.2.mlp_norm.weight": "model.layers.2.post_attention_layernorm.weight",
"layers.2.attention.q_proj.weight": "model.layers.2.self_attn.q_proj.weight",
"layers.2.attention.k_proj.weight": "model.layers.2.self_attn.k_proj.weight",
"layers.2.attention.v_proj.weight": "model.layers.2.self_attn.v_proj.weight",
"layers.2.attention.o_proj.weight": "model.layers.2.self_attn.o_proj.weight",
"layers.2.mlp.gate.weight": "model.layers.2.mlp.gate_proj.weight",
"layers.2.mlp.up.weight": "model.layers.2.mlp.up_proj.weight",
"layers.2.mlp.down.weight": "model.layers.2.mlp.down_proj.weight",
"layers.3.attention_norm.weight": "model.layers.3.input_layernorm.weight",
"layers.3.mlp_norm.weight": "model.layers.3.post_attention_layernorm.weight",
"layers.3.attention.q_proj.weight": "model.layers.3.self_attn.q_proj.weight",
"layers.3.attention.k_proj.weight": "model.layers.3.self_attn.k_proj.weight",
"layers.3.attention.v_proj.weight": "model.layers.3.self_attn.v_proj.weight",
"layers.3.attention.o_proj.weight": "model.layers.3.self_attn.o_proj.weight",
"layers.3.mlp.gate.weight": "model.layers.3.mlp.gate_proj.weight",
"layers.3.mlp.up.weight": "model.layers.3.mlp.up_proj.weight",
"layers.3.mlp.down.weight": "model.layers.3.mlp.down_proj.weight",
"layers.4.attention_norm.weight": "model.layers.4.input_layernorm.weight",
"layers.4.mlp_norm.weight": "model.layers.4.post_attention_layernorm.weight",
"layers.4.attention.q_proj.weight": "model.layers.4.self_attn.q_proj.weight",
"layers.4.attention.k_proj.weight": "model.layers.4.self_attn.k_proj.weight",
"layers.4.attention.v_proj.weight": "model.layers.4.self_attn.v_proj.weight",
"layers.4.attention.o_proj.weight": "model.layers.4.self_attn.o_proj.weight",
"layers.4.mlp.gate.weight": "model.layers.4.mlp.gate_proj.weight",
"layers.4.mlp.up.weight": "model.layers.4.mlp.up_proj.weight",
"layers.4.mlp.down.weight": "model.layers.4.mlp.down_proj.weight",
"layers.5.attention_norm.weight": "model.layers.5.input_layernorm.weight",
"layers.5.mlp_norm.weight": "model.layers.5.post_attention_layernorm.weight",
"layers.5.attention.q_proj.weight": "model.layers.5.self_attn.q_proj.weight",
"layers.5.attention.k_proj.weight": "model.layers.5.self_attn.k_proj.weight",
"layers.5.attention.v_proj.weight": "model.layers.5.self_attn.v_proj.weight",
"layers.5.attention.o_proj.weight": "model.layers.5.self_attn.o_proj.weight",
"layers.5.mlp.gate.weight": "model.layers.5.mlp.gate_proj.weight",
"layers.5.mlp.up.weight": "model.layers.5.mlp.up_proj.weight",
"layers.5.mlp.down.weight": "model.layers.5.mlp.down_proj.weight",
"layers.6.attention_norm.weight": "model.layers.6.input_layernorm.weight",
"layers.6.mlp_norm.weight": "model.layers.6.post_attention_layernorm.weight",
"layers.6.attention.q_proj.weight": "model.layers.6.self_attn.q_proj.weight",
"layers.6.attention.k_proj.weight": "model.layers.6.self_attn.k_proj.weight",
"layers.6.attention.v_proj.weight": "model.layers.6.self_attn.v_proj.weight",
"layers.6.attention.o_proj.weight": "model.layers.6.self_attn.o_proj.weight",
"layers.6.mlp.gate.weight": "model.layers.6.mlp.gate_proj.weight",
"layers.6.mlp.up.weight": "model.layers.6.mlp.up_proj.weight",
"layers.6.mlp.down.weight": "model.layers.6.mlp.down_proj.weight",
"layers.7.attention_norm.weight": "model.layers.7.input_layernorm.weight",
"layers.7.mlp_norm.weight": "model.layers.7.post_attention_layernorm.weight",
"layers.7.attention.q_proj.weight": "model.layers.7.self_attn.q_proj.weight",
"layers.7.attention.k_proj.weight": "model.layers.7.self_attn.k_proj.weight",
"layers.7.attention.v_proj.weight": "model.layers.7.self_attn.v_proj.weight",
"layers.7.attention.o_proj.weight": "model.layers.7.self_attn.o_proj.weight",
"layers.7.mlp.gate.weight": "model.layers.7.mlp.gate_proj.weight",
"layers.7.mlp.up.weight": "model.layers.7.mlp.up_proj.weight",
"layers.7.mlp.down.weight": "model.layers.7.mlp.down_proj.weight",
"layers.8.attention_norm.weight": "model.layers.8.input_layernorm.weight",
"layers.8.mlp_norm.weight": "model.layers.8.post_attention_layernorm.weight",
"layers.8.attention.q_proj.weight": "model.layers.8.self_attn.q_proj.weight",
"layers.8.attention.k_proj.weight": "model.layers.8.self_attn.k_proj.weight",
"layers.8.attention.v_proj.weight": "model.layers.8.self_attn.v_proj.weight",
"layers.8.attention.o_proj.weight": "model.layers.8.self_attn.o_proj.weight",
"layers.8.mlp.gate.weight": "model.layers.8.mlp.gate_proj.weight",
"layers.8.mlp.up.weight": "model.layers.8.mlp.up_proj.weight",
"layers.8.mlp.down.weight": "model.layers.8.mlp.down_proj.weight",
"layers.9.attention_norm.weight": "model.layers.9.input_layernorm.weight",
"layers.9.mlp_norm.weight": "model.layers.9.post_attention_layernorm.weight",
"layers.9.attention.q_proj.weight": "model.layers.9.self_attn.q_proj.weight",
"layers.9.attention.k_proj.weight": "model.layers.9.self_attn.k_proj.weight",
"layers.9.attention.v_proj.weight": "model.layers.9.self_attn.v_proj.weight",
"layers.9.attention.o_proj.weight": "model.layers.9.self_attn.o_proj.weight",
"layers.9.mlp.gate.weight": "model.layers.9.mlp.gate_proj.weight",
"layers.9.mlp.up.weight": "model.layers.9.mlp.up_proj.weight",
"layers.9.mlp.down.weight": "model.layers.9.mlp.down_proj.weight",
"layers.10.attention_norm.weight": "model.layers.10.input_layernorm.weight",
"layers.10.mlp_norm.weight": "model.layers.10.post_attention_layernorm.weight",
"layers.10.attention.q_proj.weight": "model.layers.10.self_attn.q_proj.weight",
"layers.10.attention.k_proj.weight": "model.layers.10.self_attn.k_proj.weight",
"layers.10.attention.v_proj.weight": "model.layers.10.self_attn.v_proj.weight",
"layers.10.attention.o_proj.weight": "model.layers.10.self_attn.o_proj.weight",
"layers.10.mlp.gate.weight": "model.layers.10.mlp.gate_proj.weight",
"layers.10.mlp.up.weight": "model.layers.10.mlp.up_proj.weight",
"layers.10.mlp.down.weight": "model.layers.10.mlp.down_proj.weight",
"layers.11.attention_norm.weight": "model.layers.11.input_layernorm.weight",
"layers.11.mlp_norm.weight": "model.layers.11.post_attention_layernorm.weight",
"layers.11.attention.q_proj.weight": "model.layers.11.self_attn.q_proj.weight",
"layers.11.attention.k_proj.weight": "model.layers.11.self_attn.k_proj.weight",
"layers.11.attention.v_proj.weight": "model.layers.11.self_attn.v_proj.weight",
"layers.11.attention.o_proj.weight": "model.layers.11.self_attn.o_proj.weight",
"layers.11.mlp.gate.weight": "model.layers.11.mlp.gate_proj.weight",
"layers.11.mlp.up.weight": "model.layers.11.mlp.up_proj.weight",
"layers.11.mlp.down.weight": "model.layers.11.mlp.down_proj.weight",
"layers.12.attention_norm.weight": "model.layers.12.input_layernorm.weight",
"layers.12.mlp_norm.weight": "model.layers.12.post_attention_layernorm.weight",
"layers.12.attention.q_proj.weight": "model.layers.12.self_attn.q_proj.weight",
"layers.12.attention.k_proj.weight": "model.layers.12.self_attn.k_proj.weight",
"layers.12.attention.v_proj.weight": "model.layers.12.self_attn.v_proj.weight",
"layers.12.attention.o_proj.weight": "model.layers.12.self_attn.o_proj.weight",
"layers.12.mlp.gate.weight": "model.layers.12.mlp.gate_proj.weight",
"layers.12.mlp.up.weight": "model.layers.12.mlp.up_proj.weight",
"layers.12.mlp.down.weight": "model.layers.12.mlp.down_proj.weight",
"layers.13.attention_norm.weight": "model.layers.13.input_layernorm.weight",
"layers.13.mlp_norm.weight": "model.layers.13.post_attention_layernorm.weight",
"layers.13.attention.q_proj.weight": "model.layers.13.self_attn.q_proj.weight",
"layers.13.attention.k_proj.weight": "model.layers.13.self_attn.k_proj.weight",
"layers.13.attention.v_proj.weight": "model.layers.13.self_attn.v_proj.weight",
"layers.13.attention.o_proj.weight": "model.layers.13.self_attn.o_proj.weight",
"layers.13.mlp.gate.weight": "model.layers.13.mlp.gate_proj.weight",
"layers.13.mlp.up.weight": "model.layers.13.mlp.up_proj.weight",
"layers.13.mlp.down.weight": "model.layers.13.mlp.down_proj.weight",
"layers.14.attention_norm.weight": "model.layers.14.input_layernorm.weight",
"layers.14.mlp_norm.weight": "model.layers.14.post_attention_layernorm.weight",
"layers.14.attention.q_proj.weight": "model.layers.14.self_attn.q_proj.weight",
"layers.14.attention.k_proj.weight": "model.layers.14.self_attn.k_proj.weight",
"layers.14.attention.v_proj.weight": "model.layers.14.self_attn.v_proj.weight",
"layers.14.attention.o_proj.weight": "model.layers.14.self_attn.o_proj.weight",
"layers.14.mlp.gate.weight": "model.layers.14.mlp.gate_proj.weight",
"layers.14.mlp.up.weight": "model.layers.14.mlp.up_proj.weight",
"layers.14.mlp.down.weight": "model.layers.14.mlp.down_proj.weight",
"layers.15.attention_norm.weight": "model.layers.15.input_layernorm.weight",
"layers.15.mlp_norm.weight": "model.layers.15.post_attention_layernorm.weight",
"layers.15.attention.q_proj.weight": "model.layers.15.self_attn.q_proj.weight",
"layers.15.attention.k_proj.weight": "model.layers.15.self_attn.k_proj.weight",
"layers.15.attention.v_proj.weight": "model.layers.15.self_attn.v_proj.weight",
"layers.15.attention.o_proj.weight": "model.layers.15.self_attn.o_proj.weight",
"layers.15.mlp.gate.weight": "model.layers.15.mlp.gate_proj.weight",
"layers.15.mlp.up.weight": "model.layers.15.mlp.up_proj.weight",
"layers.15.mlp.down.weight": "model.layers.15.mlp.down_proj.weight",
"layers.16.attention_norm.weight": "model.layers.16.input_layernorm.weight",
"layers.16.mlp_norm.weight": "model.layers.16.post_attention_layernorm.weight",
"layers.16.attention.q_proj.weight": "model.layers.16.self_attn.q_proj.weight",
"layers.16.attention.k_proj.weight": "model.layers.16.self_attn.k_proj.weight",
"layers.16.attention.v_proj.weight": "model.layers.16.self_attn.v_proj.weight",
"layers.16.attention.o_proj.weight": "model.layers.16.self_attn.o_proj.weight",
"layers.16.mlp.gate.weight": "model.layers.16.mlp.gate_proj.weight",
"layers.16.mlp.up.weight": "model.layers.16.mlp.up_proj.weight",
"layers.16.mlp.down.weight": "model.layers.16.mlp.down_proj.weight",
"layers.17.attention_norm.weight": "model.layers.17.input_layernorm.weight",
"layers.17.mlp_norm.weight": "model.layers.17.post_attention_layernorm.weight",
"layers.17.attention.q_proj.weight": "model.layers.17.self_attn.q_proj.weight",
"layers.17.attention.k_proj.weight": "model.layers.17.self_attn.k_proj.weight",
"layers.17.attention.v_proj.weight": "model.layers.17.self_attn.v_proj.weight",
"layers.17.attention.o_proj.weight": "model.layers.17.self_attn.o_proj.weight",
"layers.17.mlp.gate.weight": "model.layers.17.mlp.gate_proj.weight",
"layers.17.mlp.up.weight": "model.layers.17.mlp.up_proj.weight",
"layers.17.mlp.down.weight": "model.layers.17.mlp.down_proj.weight",
"layers.18.attention_norm.weight": "model.layers.18.input_layernorm.weight",
"layers.18.mlp_norm.weight": "model.layers.18.post_attention_layernorm.weight",
"layers.18.attention.q_proj.weight": "model.layers.18.self_attn.q_proj.weight",
"layers.18.attention.k_proj.weight": "model.layers.18.self_attn.k_proj.weight",
"layers.18.attention.v_proj.weight": "model.layers.18.self_attn.v_proj.weight",
"layers.18.attention.o_proj.weight": "model.layers.18.self_attn.o_proj.weight",
"layers.18.mlp.gate.weight": "model.layers.18.mlp.gate_proj.weight",
"layers.18.mlp.up.weight": "model.layers.18.mlp.up_proj.weight",
"layers.18.mlp.down.weight": "model.layers.18.mlp.down_proj.weight",
"layers.19.attention_norm.weight": "model.layers.19.input_layernorm.weight",
"layers.19.mlp_norm.weight": "model.layers.19.post_attention_layernorm.weight",
"layers.19.attention.q_proj.weight": "model.layers.19.self_attn.q_proj.weight",
"layers.19.attention.k_proj.weight": "model.layers.19.self_attn.k_proj.weight",
"layers.19.attention.v_proj.weight": "model.layers.19.self_attn.v_proj.weight",
"layers.19.attention.o_proj.weight": "model.layers.19.self_attn.o_proj.weight",
"layers.19.mlp.gate.weight": "model.layers.19.mlp.gate_proj.weight",
"layers.19.mlp.up.weight": "model.layers.19.mlp.up_proj.weight",
"layers.19.mlp.down.weight": "model.layers.19.mlp.down_proj.weight",
"layers.20.attention_norm.weight": "model.layers.20.input_layernorm.weight",
"layers.20.mlp_norm.weight": "model.layers.20.post_attention_layernorm.weight",
"layers.20.attention.q_proj.weight": "model.layers.20.self_attn.q_proj.weight",
"layers.20.attention.k_proj.weight": "model.layers.20.self_attn.k_proj.weight",
"layers.20.attention.v_proj.weight": "model.layers.20.self_attn.v_proj.weight",
"layers.20.attention.o_proj.weight": "model.layers.20.self_attn.o_proj.weight",
"layers.20.mlp.gate.weight": "model.layers.20.mlp.gate_proj.weight",
"layers.20.mlp.up.weight": "model.layers.20.mlp.up_proj.weight",
"layers.20.mlp.down.weight": "model.layers.20.mlp.down_proj.weight",
"layers.21.attention_norm.weight": "model.layers.21.input_layernorm.weight",
"layers.21.mlp_norm.weight": "model.layers.21.post_attention_layernorm.weight",
"layers.21.attention.q_proj.weight": "model.layers.21.self_attn.q_proj.weight",
"layers.21.attention.k_proj.weight": "model.layers.21.self_attn.k_proj.weight",
"layers.21.attention.v_proj.weight": "model.layers.21.self_attn.v_proj.weight",
"layers.21.attention.o_proj.weight": "model.layers.21.self_attn.o_proj.weight",
"layers.21.mlp.gate.weight": "model.layers.21.mlp.gate_proj.weight",
"layers.21.mlp.up.weight": "model.layers.21.mlp.up_proj.weight",
"layers.21.mlp.down.weight": "model.layers.21.mlp.down_proj.weight",
"layers.22.attention_norm.weight": "model.layers.22.input_layernorm.weight",
"layers.22.mlp_norm.weight": "model.layers.22.post_attention_layernorm.weight",
"layers.22.attention.q_proj.weight": "model.layers.22.self_attn.q_proj.weight",
"layers.22.attention.k_proj.weight": "model.layers.22.self_attn.k_proj.weight",
"layers.22.attention.v_proj.weight": "model.layers.22.self_attn.v_proj.weight",
"layers.22.attention.o_proj.weight": "model.layers.22.self_attn.o_proj.weight",
"layers.22.mlp.gate.weight": "model.layers.22.mlp.gate_proj.weight",
"layers.22.mlp.up.weight": "model.layers.22.mlp.up_proj.weight",
"layers.22.mlp.down.weight": "model.layers.22.mlp.down_proj.weight",
"layers.23.attention_norm.weight": "model.layers.23.input_layernorm.weight",
"layers.23.mlp_norm.weight": "model.layers.23.post_attention_layernorm.weight",
"layers.23.attention.q_proj.weight": "model.layers.23.self_attn.q_proj.weight",
"layers.23.attention.k_proj.weight": "model.layers.23.self_attn.k_proj.weight",
"layers.23.attention.v_proj.weight": "model.layers.23.self_attn.v_proj.weight",
"layers.23.attention.o_proj.weight": "model.layers.23.self_attn.o_proj.weight",
"layers.23.mlp.gate.weight": "model.layers.23.mlp.gate_proj.weight",
"layers.23.mlp.up.weight": "model.layers.23.mlp.up_proj.weight",
"layers.23.mlp.down.weight": "model.layers.23.mlp.down_proj.weight",
"layers.24.attention_norm.weight": "model.layers.24.input_layernorm.weight",
"layers.24.mlp_norm.weight": "model.layers.24.post_attention_layernorm.weight",
"layers.24.attention.q_proj.weight": "model.layers.24.self_attn.q_proj.weight",
"layers.24.attention.k_proj.weight": "model.layers.24.self_attn.k_proj.weight",
"layers.24.attention.v_proj.weight": "model.layers.24.self_attn.v_proj.weight",
"layers.24.attention.o_proj.weight": "model.layers.24.self_attn.o_proj.weight",
"layers.24.mlp.gate.weight": "model.layers.24.mlp.gate_proj.weight",
"layers.24.mlp.up.weight": "model.layers.24.mlp.up_proj.weight",
"layers.24.mlp.down.weight": "model.layers.24.mlp.down_proj.weight",
"layers.25.attention_norm.weight": "model.layers.25.input_layernorm.weight",
"layers.25.mlp_norm.weight": "model.layers.25.post_attention_layernorm.weight",
"layers.25.attention.q_proj.weight": "model.layers.25.self_attn.q_proj.weight",
"layers.25.attention.k_proj.weight": "model.layers.25.self_attn.k_proj.weight",
"layers.25.attention.v_proj.weight": "model.layers.25.self_attn.v_proj.weight",
"layers.25.attention.o_proj.weight": "model.layers.25.self_attn.o_proj.weight",
"layers.25.mlp.gate.weight": "model.layers.25.mlp.gate_proj.weight",
"layers.25.mlp.up.weight": "model.layers.25.mlp.up_proj.weight",
"layers.25.mlp.down.weight": "model.layers.25.mlp.down_proj.weight",
"layers.26.attention_norm.weight": "model.layers.26.input_layernorm.weight",
"layers.26.mlp_norm.weight": "model.layers.26.post_attention_layernorm.weight",
"layers.26.attention.q_proj.weight": "model.layers.26.self_attn.q_proj.weight",
"layers.26.attention.k_proj.weight": "model.layers.26.self_attn.k_proj.weight",
"layers.26.attention.v_proj.weight": "model.layers.26.self_attn.v_proj.weight",
"layers.26.attention.o_proj.weight": "model.layers.26.self_attn.o_proj.weight",
"layers.26.mlp.gate.weight": "model.layers.26.mlp.gate_proj.weight",
"layers.26.mlp.up.weight": "model.layers.26.mlp.up_proj.weight",
"layers.26.mlp.down.weight": "model.layers.26.mlp.down_proj.weight",
"layers.27.attention_norm.weight": "model.layers.27.input_layernorm.weight",
"layers.27.mlp_norm.weight": "model.layers.27.post_attention_layernorm.weight",
"layers.27.attention.q_proj.weight": "model.layers.27.self_attn.q_proj.weight",
"layers.27.attention.k_proj.weight": "model.layers.27.self_attn.k_proj.weight",
"layers.27.attention.v_proj.weight": "model.layers.27.self_attn.v_proj.weight",
"layers.27.attention.o_proj.weight": "model.layers.27.self_attn.o_proj.weight",
"layers.27.mlp.gate.weight": "model.layers.27.mlp.gate_proj.weight",
"layers.27.mlp.up.weight": "model.layers.27.mlp.up_proj.weight",
"layers.27.mlp.down.weight": "model.layers.27.mlp.down_proj.weight",
"layers.28.attention_norm.weight": "model.layers.28.input_layernorm.weight",
"layers.28.mlp_norm.weight": "model.layers.28.post_attention_layernorm.weight",
"layers.28.attention.q_proj.weight": "model.layers.28.self_attn.q_proj.weight",
"layers.28.attention.k_proj.weight": "model.layers.28.self_attn.k_proj.weight",
"layers.28.attention.v_proj.weight": "model.layers.28.self_attn.v_proj.weight",
"layers.28.attention.o_proj.weight": "model.layers.28.self_attn.o_proj.weight",
"layers.28.mlp.gate.weight": "model.layers.28.mlp.gate_proj.weight",
"layers.28.mlp.up.weight": "model.layers.28.mlp.up_proj.weight",
"layers.28.mlp.down.weight": "model.layers.28.mlp.down_proj.weight",
"layers.29.attention_norm.weight": "model.layers.29.input_layernorm.weight",
"layers.29.mlp_norm.weight": "model.layers.29.post_attention_layernorm.weight",
"layers.29.attention.q_proj.weight": "model.layers.29.self_attn.q_proj.weight",
"layers.29.attention.k_proj.weight": "model.layers.29.self_attn.k_proj.weight",
"layers.29.attention.v_proj.weight": "model.layers.29.self_attn.v_proj.weight",
"layers.29.attention.o_proj.weight": "model.layers.29.self_attn.o_proj.weight",
"layers.29.mlp.gate.weight": "model.layers.29.mlp.gate_proj.weight",
"layers.29.mlp.up.weight": "model.layers.29.mlp.up_proj.weight",
"layers.29.mlp.down.weight": "model.layers.29.mlp.down_proj.weight",
"layers.30.attention_norm.weight": "model.layers.30.input_layernorm.weight",
"layers.30.mlp_norm.weight": "model.layers.30.post_attention_layernorm.weight",
"layers.30.attention.q_proj.weight": "model.layers.30.self_attn.q_proj.weight",
"layers.30.attention.k_proj.weight": "model.layers.30.self_attn.k_proj.weight",
"layers.30.attention.v_proj.weight": "model.layers.30.self_attn.v_proj.weight",
"layers.30.attention.o_proj.weight": "model.layers.30.self_attn.o_proj.weight",
"layers.30.mlp.gate.weight": "model.layers.30.mlp.gate_proj.weight",
"layers.30.mlp.up.weight": "model.layers.30.mlp.up_proj.weight",
"layers.30.mlp.down.weight": "model.layers.30.mlp.down_proj.weight",
"layers.31.attention_norm.weight": "model.layers.31.input_layernorm.weight",
"layers.31.mlp_norm.weight": "model.layers.31.post_attention_layernorm.weight",
"layers.31.attention.q_proj.weight": "model.layers.31.self_attn.q_proj.weight",
"layers.31.attention.k_proj.weight": "model.layers.31.self_attn.k_proj.weight",
"layers.31.attention.v_proj.weight": "model.layers.31.self_attn.v_proj.weight",
"layers.31.attention.o_proj.weight": "model.layers.31.self_attn.o_proj.weight",
"layers.31.mlp.gate.weight": "model.layers.31.mlp.gate_proj.weight",
"layers.31.mlp.up.weight": "model.layers.31.mlp.up_proj.weight",
"layers.31.mlp.down.weight": "model.layers.31.mlp.down_proj.weight"
}linnet.toml60 B
toml
[package]
name = "llama"
version = "0.1.0"
language = "0.1"nest.toml1.1 kB
toml
[model]
name = "llama-3.1-8b-instruct"
title = "Llama 3.1 8B Instruct"
summary = "An 8B-parameter Llama-architecture decoder with grouped-query attention, llama3 rope scaling for a 131072-token context, a 128256-token vocabulary, and an untied output head, instruction-tuned."
license = "llama3.1"
family = "llama"
tags = ["text-generation", "decoder-only", "grouped-query-attention", "chat"]
[links]
huggingface = "https://huggingface.co/NousResearch/Meta-Llama-3.1-8B-Instruct"
github = "https://github.com/meta-llama/llama-models"
arxiv = "https://arxiv.org/abs/2407.21783"
[source]
path = "src/lib.linnet"
root = "Model"
entry = "forward"
[generics]
Vocab = 128256
H = 4096
Heads = 32
KvHeads = 8
Inner = 14336
Layers = 32
Batch = 1
MaxSeq = 8192
T = "bf16"
[check]
B = 1
S = 8
[weights]
repo = "NousResearch/Meta-Llama-3.1-8B-Instruct"
revision = "d10aef7999a2b5ba950ab3974312feeedbfe0b77"
files = [
"model-00001-of-00004.safetensors",
"model-00002-of-00004.safetensors",
"model-00003-of-00004.safetensors",
"model-00004-of-00004.safetensors",
]
bindings = "bindings.json"samples.json1.0 kB
json
{
"kind": "text",
"prompt": "Explain in two sentences why the sky is blue.",
"chat": true,
"output": "Here is a two-sentence explanation:\n\nThe sky appears blue because of a phenomenon called Rayleigh scattering, in which shorter (blue) wavelengths of light are scattered more than longer (red) wavelengths by the tiny molecules of gases in the Earth's atmosphere. As a result, the blue light is dispersed in all directions, reaching our eyes from every part of the sky and giving it its blue color.",
"reference": "Here is a two-sentence explanation:\n\nThe sky appears blue because of a phenomenon called Rayleigh scattering, in which shorter (blue) wavelengths of light are scattered more than longer (red) wavelengths by the tiny molecules of gases in the Earth's atmosphere. As a result, the blue light is dispersed in all directions, reaching our eyes from every part of the sky and giving it its blue color.",
"reference_stack": "transformers (KV cache, argmax, bf16)",
"tokens": 81,
"agreeing_prefix": 81
}src
attention.linnet20.0 kB
linnet
// Grouped-query attention: `KvHeads` key/value heads shared by `Heads`
// query heads, through `std.nn.attention::grouped_attention`. The `where`
// clause states the divisibility the reshapes depend on; the checker proves
// every shape from it.
//
// `forward` attends over a whole sequence. `decode` takes one position and
// keeps the keys and values seen so far in two `state` members, sized by
// the block's `Batch` and `MaxSeq`: the block assigns them, the runtime
// keeps them between calls, and graph exports thread them in and out.
// `prefill` is `decode` for a whole prompt: every position in one pass.
// `decode_rows` and `prefill_slots` serve many requests at once: each row of
// the caches belongs to one request, at its own length.
module llama.attention
use crate.rope::{THETA, tables}
use std.nn.attention::{
causal_mask,
grouped_attention,
grouped_attention_rows,
paged_attention,
paged_prefill_attention,
}
use std.nn.cache::{page_slots, write_at, write_rows, write_slots, write_span, write_tokens}
use std.nn.parallel::{all_reduce}
use std.nn.linear::{Linear}
use std.nn.rope::{rope, rope_rows}
pub block GroupedQueryAttention<
H: Dim,
Heads: Dim,
KvHeads: Dim,
Batch: Dim,
MaxSeq: Dim,
T: Float,
Shards: Dim = 1,
PageSize: Dim = 64,
>
where
Shards > 0,
Heads % Shards == 0,
KvHeads % Shards == 0,
KvHeads / Shards > 0,
(Heads / Shards) % (KvHeads / Shards) == 0,
Heads > 0,
KvHeads > 0,
H % Heads == 0,
Heads % KvHeads == 0,
(H / Heads) % 2 == 0,
MaxSeq > 0,
Batch > 0,
PageSize > 0
{
sub q_proj: Linear<H, Heads / Shards * (H / Heads), T>
sub k_proj: Linear<H, KvHeads / Shards * (H / Heads), T>
sub v_proj: Linear<H, KvHeads / Shards * (H / Heads), T>
sub o_proj: Linear<Heads / Shards * (H / Heads), H, T>
state cache_k: Tensor[Batch, KvHeads / Shards, MaxSeq, H / Heads; T]
state cache_v: Tensor[Batch, KvHeads / Shards, MaxSeq, H / Heads; T]
pub fn forward<B: Dim, S: Dim>(x: Tensor[B, S, H; T]) -> Tensor[B, S, H; T] {
let (cos_table, sin_table) = tables<S, H / Heads, T>(THETA)
let q = rope(
split_heads<B, S, Heads / Shards, H / Heads, T>(q_proj.forward(x)),
cos_table,
sin_table,
)
let k = rope(
split_heads<B, S, KvHeads / Shards, H / Heads, T>(k_proj.forward(x)),
cos_table,
sin_table,
)
let v = split_heads<B, S, KvHeads / Shards, H / Heads, T>(v_proj.forward(x))
let mixed = grouped_attention(
q,
k,
v,
rsqrt(cast<f32>(H / Heads)),
some(causal_mask<S, S>()),
)
return all_reduce(o_proj.forward(merge_heads<B, S, Heads / Shards, H / Heads, T>(mixed)))
}
// One new position `pos`: its key and value are written into the caches
// at that position (a masked write, no scatter), and the query attends
// over every cached position up to and including it.
pub fn decode(x: Tensor[Batch, 1, H; T], pos: i32) -> Tensor[Batch, 1, H; T] {
let (cos_table, sin_table) = tables<MaxSeq, H / Heads, T>(THETA)
let cos_at[d] = cos_table[cast<i64>(pos), d]
let sin_at[d] = sin_table[cast<i64>(pos), d]
let cos_row = reshape(cos_at, [1, H / Heads])
let sin_row = reshape(sin_at, [1, H / Heads])
let q = rope(
split_heads<Batch, 1, Heads / Shards, H / Heads, T>(q_proj.forward(x)),
cos_row,
sin_row,
)
let k = rope(
split_heads<Batch, 1, KvHeads / Shards, H / Heads, T>(k_proj.forward(x)),
cos_row,
sin_row,
)
let v = split_heads<Batch, 1, KvHeads / Shards, H / Heads, T>(v_proj.forward(x))
cache_k = write_at(cache_k, k, pos)
cache_v = write_at(cache_v, v, pos)
let positions = iota<i32>(MaxSeq)
let seen[s] = positions[s] <= pos
let mixed = grouped_attention(
q,
cache_k,
cache_v,
rsqrt(cast<f32>(H / Heads)),
some(reshape(seen, [1, MaxSeq])),
)
return all_reduce(
o_proj.forward(merge_heads<Batch, 1, Heads / Shards, H / Heads, T>(mixed)),
)
}
// `S` new positions starting at `pos`, as a prompt arrives: their keys
// and values go into the caches as one span, and each query attends over
// every cached position up to and including its own. A prompt costs one
// pass instead of one per token, which is most of the time to the first
// generated token. The rope rows come from the same scaled `tables` as
// `decode`, gathered at `pos + iota(S)` rather than recomputed.
pub fn prefill<S: Dim>(x: Tensor[Batch, S, H; T], pos: i32) -> Tensor[Batch, S, H; T]
where S > 0 {
let (cos_table, sin_table) = tables<MaxSeq, H / Heads, T>(THETA)
let rows[i] = cast<i64>(pos) + iota<i64>(S)[i]
let cos_rows[i, d] = cos_table[rows[i], d]
let sin_rows[i, d] = sin_table[rows[i], d]
let q = rope(
split_heads<Batch, S, Heads / Shards, H / Heads, T>(q_proj.forward(x)),
cos_rows,
sin_rows,
)
let k = rope(
split_heads<Batch, S, KvHeads / Shards, H / Heads, T>(k_proj.forward(x)),
cos_rows,
sin_rows,
)
let v = split_heads<Batch, S, KvHeads / Shards, H / Heads, T>(v_proj.forward(x))
cache_k = write_span(cache_k, k, pos)
cache_v = write_span(cache_v, v, pos)
let positions = iota<i32>(MaxSeq)
let queries = iota<i32>(S)
let seen[i, s] = positions[s] <= pos + queries[i]
let mixed = grouped_attention(q, cache_k, cache_v, rsqrt(cast<f32>(H / Heads)), some(seen))
return all_reduce(
o_proj.forward(merge_heads<Batch, S, Heads / Shards, H / Heads, T>(mixed)),
)
}
// One new position per sequence, each at its own `positions[b]`: the
// step a server takes for every request it is decoding together. Row `b`
// of the caches is written at `positions[b]`, and its query attends over
// that row up to there.
pub fn decode_rows(
x: Tensor[Batch, 1, H; T],
positions: Tensor[Batch; i32],
) -> Tensor[Batch, 1, H; T] {
let (cos_table, sin_table) = tables<MaxSeq, H / Heads, T>(THETA)
let cos_at[b, d] = cos_table[cast<i64>(positions[b]), d]
let sin_at[b, d] = sin_table[cast<i64>(positions[b]), d]
let cos_rows = reshape(cos_at, [Batch, 1, H / Heads])
let sin_rows = reshape(sin_at, [Batch, 1, H / Heads])
let q = rope_rows(
split_heads<Batch, 1, Heads / Shards, H / Heads, T>(q_proj.forward(x)),
cos_rows,
sin_rows,
)
let k = rope_rows(
split_heads<Batch, 1, KvHeads / Shards, H / Heads, T>(k_proj.forward(x)),
cos_rows,
sin_rows,
)
let v = split_heads<Batch, 1, KvHeads / Shards, H / Heads, T>(v_proj.forward(x))
cache_k = write_rows(cache_k, k, positions)
cache_v = write_rows(cache_v, v, positions)
let slots = iota<i32>(MaxSeq)
let seen[b, s] = slots[s] <= positions[b]
let mixed = grouped_attention_rows(
q,
cache_k,
cache_v,
rsqrt(cast<f32>(H / Heads)),
reshape(seen, [Batch, 1, MaxSeq]),
)
return all_reduce(
o_proj.forward(merge_heads<Batch, 1, Heads / Shards, H / Heads, T>(mixed)),
)
}
// `M` requests' prompts of `S` tokens, from their first positions, into
// rows `slots` of the caches: each prompt's keys and values fill its row,
// and its queries attend causally among themselves. The rest of a row
// keeps what an earlier request left there; decoding never reads past its
// own position, so nothing needs clearing.
pub fn prefill_slots<M: Dim, S: Dim>(
x: Tensor[M, S, H; T],
slots: Tensor[M; i32],
) -> Tensor[M, S, H; T]
where
M > 0,
S > 0
{
let (cos_table, sin_table) = tables<S, H / Heads, T>(THETA)
let q = rope(
split_heads<M, S, Heads / Shards, H / Heads, T>(q_proj.forward(x)),
cos_table,
sin_table,
)
let k = rope(
split_heads<M, S, KvHeads / Shards, H / Heads, T>(k_proj.forward(x)),
cos_table,
sin_table,
)
let v = split_heads<M, S, KvHeads / Shards, H / Heads, T>(v_proj.forward(x))
cache_k = write_slots(cache_k, k, slots, 0)
cache_v = write_slots(cache_v, v, slots, 0)
let mixed = grouped_attention(
q,
k,
v,
rsqrt(cast<f32>(H / Heads)),
some(causal_mask<S, S>()),
)
return all_reduce(o_proj.forward(merge_heads<M, S, Heads / Shards, H / Heads, T>(mixed)))
}
// Prompts packed end to end, `P` tokens in all: token `p` is position
// `positions[p]` of prompt `segments[p]`, whose row of the caches is
// `rows[p]`. Each token's key and value go to its row at its position, and
// its query attends to the tokens of its own prompt up to itself -- no
// padding between prompts, and none of the pass's other prompts seen.
pub fn prefill_packed<P: Dim>(
x: Tensor[1, P, H; T],
rows: Tensor[P; i32],
positions: Tensor[P; i32],
segments: Tensor[P; i32],
) -> Tensor[1, P, H; T]
where P > 0 {
let (cos_table, sin_table) = tables<MaxSeq, H / Heads, T>(THETA)
let cos_at[p, d] = cos_table[cast<i64>(positions[p]), d]
let sin_at[p, d] = sin_table[cast<i64>(positions[p]), d]
let q = rope(
split_heads<1, P, Heads / Shards, H / Heads, T>(q_proj.forward(x)),
cos_at,
sin_at,
)
let k = rope(
split_heads<1, P, KvHeads / Shards, H / Heads, T>(k_proj.forward(x)),
cos_at,
sin_at,
)
let v = split_heads<1, P, KvHeads / Shards, H / Heads, T>(v_proj.forward(x))
cache_k = write_tokens(cache_k, k, rows, positions)
cache_v = write_tokens(cache_v, v, rows, positions)
let own[i, j] = segments[i] == segments[j] && positions[j] <= positions[i]
let mixed = grouped_attention(q, k, v, rsqrt(cast<f32>(H / Heads)), some(own))
return all_reduce(o_proj.forward(merge_heads<1, P, Heads / Shards, H / Heads, T>(mixed)))
}
// `prefill_packed` for training: every packed token's output, and no
// cache written.
pub fn packed<P: Dim>(
x: Tensor[1, P, H; T],
positions: Tensor[P; i32],
segments: Tensor[P; i32],
) -> Tensor[1, P, H; T]
where P > 0 {
let (cos_table, sin_table) = tables<MaxSeq, H / Heads, T>(THETA)
let cos_at[p, d] = cos_table[cast<i64>(positions[p]), d]
let sin_at[p, d] = sin_table[cast<i64>(positions[p]), d]
let q = rope(
split_heads<1, P, Heads / Shards, H / Heads, T>(q_proj.forward(x)),
cos_at,
sin_at,
)
let k = rope(
split_heads<1, P, KvHeads / Shards, H / Heads, T>(k_proj.forward(x)),
cos_at,
sin_at,
)
let v = split_heads<1, P, KvHeads / Shards, H / Heads, T>(v_proj.forward(x))
let own[i, j] = segments[i] == segments[j] && positions[j] <= positions[i]
let mixed = grouped_attention(q, k, v, rsqrt(cast<f32>(H / Heads)), some(own))
return all_reduce(o_proj.forward(merge_heads<1, P, Heads / Shards, H / Heads, T>(mixed)))
}
// `prefill_packed`'s prompts, `P` tokens, then one step for each of the
// `Batch` rows, at `step_positions[b]`: every token's key and value go to
// its row at its position, the prompts attend among themselves as in
// `prefill_packed`, and each step over its row as in `decode_rows`.
pub fn step_packed<P: Dim>(
x: Tensor[1, P + Batch, H; T],
rows: Tensor[P; i32],
positions: Tensor[P; i32],
segments: Tensor[P; i32],
step_positions: Tensor[Batch; i32],
) -> Tensor[1, P + Batch, H; T]
where P > 0 {
let every_at = concat(positions, step_positions, axis = 0)
let every_row = concat(rows, iota<i32>(Batch), axis = 0)
let (cos_table, sin_table) = tables<MaxSeq, H / Heads, T>(THETA)
let cos_at[p, d] = cos_table[cast<i64>(every_at[p]), d]
let sin_at[p, d] = sin_table[cast<i64>(every_at[p]), d]
let q = rope(
split_heads<1, P + Batch, Heads / Shards, H / Heads, T>(q_proj.forward(x)),
cos_at,
sin_at,
)
let k = rope(
split_heads<1, P + Batch, KvHeads / Shards, H / Heads, T>(k_proj.forward(x)),
cos_at,
sin_at,
)
let v = split_heads<1, P + Batch, KvHeads / Shards, H / Heads, T>(v_proj.forward(x))
cache_k = write_tokens(cache_k, k, every_row, every_at)
cache_v = write_tokens(cache_v, v, every_row, every_at)
let own[i, j] = segments[i] == segments[j] && positions[j] <= positions[i]
let prompts = grouped_attention(
q[:, :, 0:P, :],
k[:, :, 0:P, :],
v[:, :, 0:P, :],
rsqrt(cast<f32>(H / Heads)),
some(own),
)
let slots = iota<i32>(MaxSeq)
let seen[b, s] = slots[s] <= step_positions[b]
let steps = grouped_attention_rows(
permute(q[:, :, P:P + Batch, :], [2, 1, 0, 3]),
cache_k,
cache_v,
rsqrt(cast<f32>(H / Heads)),
reshape(seen, [Batch, 1, MaxSeq]),
)
let mixed = concat(prompts, permute(steps, [2, 1, 0, 3]), axis = 2)
return all_reduce(
o_proj.forward(merge_heads<1, P + Batch, Heads / Shards, H / Heads, T>(mixed)),
)
}
// Paged serving: with `Batch` 1 the caches' one row is a pool of
// pages, `PageSize` positions each. `decode_paged` is `decode_rows` for
// `Rows` rows whose positions lie in the pages `table` lists: each row's
// key and value go to its position's place in the pool, and its query
// reads its own pages.
pub fn decode_paged<Rows: Dim, Pages: Dim>(
x: Tensor[Rows, 1, H; T],
positions: Tensor[Rows; i32],
table: Tensor[Rows, Pages; i32],
) -> Tensor[Rows, 1, H; T]
where Rows > 0 {
let (cos_table, sin_table) = tables<MaxSeq, H / Heads, T>(THETA)
let cos_at[b, d] = cos_table[cast<i64>(positions[b]), d]
let sin_at[b, d] = sin_table[cast<i64>(positions[b]), d]
let cos_rows = reshape(cos_at, [Rows, 1, H / Heads])
let sin_rows = reshape(sin_at, [Rows, 1, H / Heads])
let q = rope_rows(
split_heads<Rows, 1, Heads / Shards, H / Heads, T>(q_proj.forward(x)),
cos_rows,
sin_rows,
)
let k = rope_rows(
split_heads<Rows, 1, KvHeads / Shards, H / Heads, T>(k_proj.forward(x)),
cos_rows,
sin_rows,
)
let v = split_heads<Rows, 1, KvHeads / Shards, H / Heads, T>(v_proj.forward(x))
let slots = page_slots<Rows, Pages, PageSize>(table, positions)
let pool = fill<i32>([Rows], 0)
cache_k = write_tokens(cache_k, permute(k, [2, 1, 0, 3]), pool, slots)
cache_v = write_tokens(cache_v, permute(v, [2, 1, 0, 3]), pool, slots)
let mixed = paged_attention<Rows, Heads / Shards, KvHeads /
Shards, MaxSeq, Pages, PageSize, H / Heads, T>(
q,
cache_k[0:1, :, :, :],
cache_v[0:1, :, :, :],
table,
positions,
rsqrt(cast<f32>(H / Heads)),
)
return all_reduce(
o_proj.forward(merge_heads<Rows, 1, Heads / Shards, H / Heads, T>(mixed)),
)
}
// Prompts packed end to end over the pool: each token, written at its
// place (`slots`), attends over its row's pages (`table[rows[p]]`) up
// to its position, so a prompt can pass in chunks or start after pages
// it shares.
pub fn prefill_paged<P: Dim, Rows: Dim, Pages: Dim>(
x: Tensor[1, P, H; T],
positions: Tensor[P; i32],
rows: Tensor[P; i32],
slots: Tensor[P; i32],
table: Tensor[Rows, Pages; i32],
) -> Tensor[1, P, H; T]
where P > 0 {
let (cos_table, sin_table) = tables<MaxSeq, H / Heads, T>(THETA)
let cos_at[p, d] = cos_table[cast<i64>(positions[p]), d]
let sin_at[p, d] = sin_table[cast<i64>(positions[p]), d]
let q = rope(
split_heads<1, P, Heads / Shards, H / Heads, T>(q_proj.forward(x)),
cos_at,
sin_at,
)
let k = rope(
split_heads<1, P, KvHeads / Shards, H / Heads, T>(k_proj.forward(x)),
cos_at,
sin_at,
)
let v = split_heads<1, P, KvHeads / Shards, H / Heads, T>(v_proj.forward(x))
let pool = fill<i32>([P], 0)
cache_k = write_tokens(cache_k, k, pool, slots)
cache_v = write_tokens(cache_v, v, pool, slots)
let mixed = paged_prefill_attention<P, Heads / Shards, KvHeads /
Shards, MaxSeq, Rows, Pages, PageSize, H / Heads, T>(
q,
cache_k[0:1, :, :, :],
cache_v[0:1, :, :, :],
table,
rows,
positions,
rsqrt(cast<f32>(H / Heads)),
)
return all_reduce(o_proj.forward(merge_heads<1, P, Heads / Shards, H / Heads, T>(mixed)))
}
// `prefill_paged`'s prompts and `decode_paged`'s steps in one pass.
pub fn step_paged<P: Dim, Rows: Dim, Pages: Dim>(
x: Tensor[1, P + Rows, H; T],
positions: Tensor[P; i32],
rows: Tensor[P; i32],
slots: Tensor[P; i32],
step_positions: Tensor[Rows; i32],
table: Tensor[Rows, Pages; i32],
) -> Tensor[1, P + Rows, H; T]
where
P > 0,
Rows > 0
{
let every_at = concat(positions, step_positions, axis = 0)
let every_slot = concat(
slots,
page_slots<Rows, Pages, PageSize>(table, step_positions),
axis = 0,
)
let (cos_table, sin_table) = tables<MaxSeq, H / Heads, T>(THETA)
let cos_at[p, d] = cos_table[cast<i64>(every_at[p]), d]
let sin_at[p, d] = sin_table[cast<i64>(every_at[p]), d]
let q = rope(
split_heads<1, P + Rows, Heads / Shards, H / Heads, T>(q_proj.forward(x)),
cos_at,
sin_at,
)
let k = rope(
split_heads<1, P + Rows, KvHeads / Shards, H / Heads, T>(k_proj.forward(x)),
cos_at,
sin_at,
)
let v = split_heads<1, P + Rows, KvHeads / Shards, H / Heads, T>(v_proj.forward(x))
let pool = fill<i32>([P + Rows], 0)
cache_k = write_tokens(cache_k, k, pool, every_slot)
cache_v = write_tokens(cache_v, v, pool, every_slot)
let prompts = paged_prefill_attention<P, Heads / Shards, KvHeads /
Shards, MaxSeq, Rows, Pages, PageSize, H / Heads, T>(
q[:, :, 0:P, :],
cache_k[0:1, :, :, :],
cache_v[0:1, :, :, :],
table,
rows,
positions,
rsqrt(cast<f32>(H / Heads)),
)
let steps = paged_attention<Rows, Heads / Shards, KvHeads /
Shards, MaxSeq, Pages, PageSize, H / Heads, T>(
permute(q[:, :, P:P + Rows, :], [2, 1, 0, 3]),
cache_k[0:1, :, :, :],
cache_v[0:1, :, :, :],
table,
step_positions,
rsqrt(cast<f32>(H / Heads)),
)
let mixed = concat(prompts, permute(steps, [2, 1, 0, 3]), axis = 2)
return all_reduce(
o_proj.forward(merge_heads<1, P + Rows, Heads / Shards, H / Heads, T>(mixed)),
)
}
}
// `[B, S, N * D]` -> `[B, N, S, D]`.
fn split_heads<B: Dim, S: Dim, N: Dim, D: Dim, T: Float>(
x: Tensor[B, S, N * D; T],
) -> Tensor[B, N, S, D; T] {
return permute(reshape(x, [B, S, N, D]), [0, 2, 1, 3])
}
fn merge_heads<B: Dim, S: Dim, N: Dim, D: Dim, T: Float>(
x: Tensor[B, N, S, D; T],
) -> Tensor[B, S, N * D; T] {
return reshape(permute(x, [0, 2, 1, 3]), [B, S, N * D])
}lib.linnet20.1 kB
linnet
// A Llama-style decoder: RMSNorm, grouped-query attention with rotary
// positions, SwiGLU, and an untied output head. Widths, head counts, and
// depth are generic parameters; the dtype defaults to bf16 and every
// intermediate accumulation is spelled out in the standard library. The
// `decode` entry generates one token at a time from the KV caches the
// attention blocks own, sized by `Batch` and `MaxSeq`. `crate.rope` adds
// Llama 3.1's `llama3`-style rope scaling on top of the base frequency.
module llama
use crate.attention::{GroupedQueryAttention}
use std.nn.decoding::{argmax}
use std.random::{categorical, split}
use std.nn.embedding::{Embedding}
use std.nn.parallel::{all_gather, all_reduce, shared}
use std.nn.linear::{Linear}
use std.nn.loss::{split_cross_entropy, split_token_log_probs}
use std.nn.mlp::{SwiGlu}
use std.nn.norm::{RmsNorm}
pub block DecoderLayer<
H: Dim,
Heads: Dim,
KvHeads: Dim,
Inner: Dim,
Batch: Dim,
MaxSeq: Dim,
T: Float,
Shards: Dim = 1,
PageSize: Dim = 64,
>
where
Shards > 0,
Heads % Shards == 0,
KvHeads % Shards == 0,
KvHeads / Shards > 0,
(Heads / Shards) % (KvHeads / Shards) == 0,
Inner % Shards == 0,
Heads > 0,
KvHeads > 0,
H % Heads == 0,
Heads % KvHeads == 0,
(H / Heads) % 2 == 0,
MaxSeq > 0,
Batch > 0,
PageSize > 0
{
sub attention_norm: RmsNorm<H, T>
sub attention: GroupedQueryAttention<H, Heads, KvHeads, Batch, MaxSeq, T, Shards, PageSize>
sub mlp_norm: RmsNorm<H, T>
sub mlp: SwiGlu<H, Inner / Shards, T>
pub fn forward<B: Dim, S: Dim>(x: Tensor[B, S, H; T]) -> Tensor[B, S, H; T] {
let attended = x + attention.forward(shared(attention_norm.forward(x)))
return attended + all_reduce(mlp.forward(shared(mlp_norm.forward(attended))))
}
pub fn decode(x: Tensor[Batch, 1, H; T], pos: i32) -> Tensor[Batch, 1, H; T] {
let attended = x + attention.decode(attention_norm.forward(x), pos)
return attended + all_reduce(mlp.forward(mlp_norm.forward(attended)))
}
pub fn prefill<S: Dim>(x: Tensor[Batch, S, H; T], pos: i32) -> Tensor[Batch, S, H; T]
where S > 0 {
let attended = x + attention.prefill(attention_norm.forward(x), pos)
return attended + all_reduce(mlp.forward(mlp_norm.forward(attended)))
}
pub fn decode_rows(
x: Tensor[Batch, 1, H; T],
positions: Tensor[Batch; i32],
) -> Tensor[Batch, 1, H; T] {
let attended = x + attention.decode_rows(attention_norm.forward(x), positions)
return attended + all_reduce(mlp.forward(mlp_norm.forward(attended)))
}
pub fn prefill_slots<M: Dim, S: Dim>(
x: Tensor[M, S, H; T],
slots: Tensor[M; i32],
) -> Tensor[M, S, H; T]
where
M > 0,
S > 0
{
let attended = x + attention.prefill_slots(attention_norm.forward(x), slots)
return attended + all_reduce(mlp.forward(mlp_norm.forward(attended)))
}
pub fn prefill_packed<P: Dim>(
x: Tensor[1, P, H; T],
rows: Tensor[P; i32],
positions: Tensor[P; i32],
segments: Tensor[P; i32],
) -> Tensor[1, P, H; T]
where P > 0 {
let attended =
x + attention.prefill_packed(attention_norm.forward(x), rows, positions, segments)
return attended + all_reduce(mlp.forward(mlp_norm.forward(attended)))
}
pub fn packed<P: Dim>(
x: Tensor[1, P, H; T],
positions: Tensor[P; i32],
segments: Tensor[P; i32],
) -> Tensor[1, P, H; T]
where P > 0 {
let attended = x + attention.packed(shared(attention_norm.forward(x)), positions, segments)
return attended + all_reduce(mlp.forward(shared(mlp_norm.forward(attended))))
}
pub fn step_packed<P: Dim>(
x: Tensor[1, P + Batch, H; T],
rows: Tensor[P; i32],
positions: Tensor[P; i32],
segments: Tensor[P; i32],
step_positions: Tensor[Batch; i32],
) -> Tensor[1, P + Batch, H; T]
where P > 0 {
let attended =
x +
attention.step_packed(
attention_norm.forward(x),
rows,
positions,
segments,
step_positions,
)
return attended + all_reduce(mlp.forward(mlp_norm.forward(attended)))
}
pub fn decode_paged<Rows: Dim, Pages: Dim>(
x: Tensor[Rows, 1, H; T],
positions: Tensor[Rows; i32],
table: Tensor[Rows, Pages; i32],
) -> Tensor[Rows, 1, H; T]
where Rows > 0 {
let attended = x + attention.decode_paged(attention_norm.forward(x), positions, table)
return attended + all_reduce(mlp.forward(mlp_norm.forward(attended)))
}
pub fn prefill_paged<P: Dim, Rows: Dim, Pages: Dim>(
x: Tensor[1, P, H; T],
positions: Tensor[P; i32],
rows: Tensor[P; i32],
slots: Tensor[P; i32],
table: Tensor[Rows, Pages; i32],
) -> Tensor[1, P, H; T]
where P > 0 {
let attended =
x + attention.prefill_paged(attention_norm.forward(x), positions, rows, slots, table)
return attended + all_reduce(mlp.forward(mlp_norm.forward(attended)))
}
pub fn step_paged<P: Dim, Rows: Dim, Pages: Dim>(
x: Tensor[1, P + Rows, H; T],
positions: Tensor[P; i32],
rows: Tensor[P; i32],
slots: Tensor[P; i32],
step_positions: Tensor[Rows; i32],
table: Tensor[Rows, Pages; i32],
) -> Tensor[1, P + Rows, H; T]
where
P > 0,
Rows > 0
{
let attended =
x +
attention.step_paged(
attention_norm.forward(x),
positions,
rows,
slots,
step_positions,
table,
)
return attended + all_reduce(mlp.forward(mlp_norm.forward(attended)))
}
}
pub block Model<
Vocab: Dim,
H: Dim,
Heads: Dim,
KvHeads: Dim,
Inner: Dim,
Layers: Dim,
Batch: Dim,
MaxSeq: Dim,
T: Float = bf16,
Shards: Dim = 1,
PageSize: Dim = 64,
>
where
Shards > 0,
Heads % Shards == 0,
KvHeads % Shards == 0,
KvHeads / Shards > 0,
(Heads / Shards) % (KvHeads / Shards) == 0,
Inner % Shards == 0,
Vocab % Shards == 0,
Heads > 0,
KvHeads > 0,
H % Heads == 0,
Heads % KvHeads == 0,
(H / Heads) % 2 == 0,
MaxSeq > 0,
Batch > 0,
PageSize > 0
{
sub embedding: Embedding<Vocab, H, T>
sub layers: [DecoderLayer<H, Heads, KvHeads, Inner, Batch, MaxSeq, T, Shards, PageSize>; Layers]
sub norm: RmsNorm<H, T>
sub lm_head: Linear<H, Vocab / Shards, T>
// Logits for every position.
pub entry forward<B: Dim, S: Dim>(tokens: Tensor[B, S; i32]) -> Tensor[B, S, Vocab; T] {
let rows = reshape(hidden(tokens), [B * S, H])
let gathered = all_gather<B * S, Vocab / Shards, Shards, T>(lm_head.forward(shared(rows)))
return reshape(gathered, [B, S, Vocab])
}
// Logits for the last position only, as decoding needs: the final norm
// and projection run on one row per sequence.
pub entry next_token<B: Dim, S: Dim>(tokens: Tensor[B, S; i32]) -> Tensor[B, Vocab; T]
where S > 0 {
let last = hidden_states(tokens)[:, S - 1, :]
return logits(last)
}
// Logits for one new token per sequence at position `pos`, attending
// over the positions decoded before it through the layers' KV caches.
// Feed a prompt token by token, then the tokens the model produces.
pub entry decode(token: Tensor[Batch, 1; i32], pos: i32) -> Tensor[Batch, Vocab; T] {
return step(token, pos)
}
// A prompt of `S` tokens at `pos`, written into the KV caches in one
// pass. Returns the logits after its last token, so `decode` continues
// at `pos + S`: the time to the first token is this one call.
pub entry prefill<S: Dim>(tokens: Tensor[Batch, S; i32], pos: i32) -> Tensor[Batch, Vocab; T]
where S > 0 {
var x = embedding.forward(tokens)
static for layer in layers {
x = layer.prefill(x, pos)
}
return logits(x[:, S - 1, :])
}
// Logits for one new token per sequence, each at its own position
// `positions[b]`: the step a server takes for every request it is
// decoding together (continuous batching). A row with no request in it
// computes along and is ignored.
pub entry decode_rows(
tokens: Tensor[Batch, 1; i32],
positions: Tensor[Batch; i32],
) -> Tensor[Batch, Vocab; T] {
var x = embedding.forward(tokens)
static for layer in layers {
x = layer.decode_rows(x, positions)
}
return logits(x[:, 0, :])
}
// `M` requests' prompts into rows `slots` of the caches, as they join a
// batch being decoded: one pass for all of them. Row `m` of `tokens` holds
// `S` tokens of which the first `lengths[m]` are its prompt; the rest pad
// it to one of a few compiled lengths and are never read. Returns the
// logits after each prompt's last token; `decode_rows` continues row
// `slots[m]` at position `lengths[m]`.
pub entry prefill_slots<M: Dim, S: Dim>(
tokens: Tensor[M, S; i32],
slots: Tensor[M; i32],
lengths: Tensor[M; i32],
) -> Tensor[M, Vocab; T]
where
M > 0,
S > 0
{
var x = embedding.forward(tokens)
static for layer in layers {
x = layer.prefill_slots(x, slots)
}
let last[b, h] = x[b, cast<i64>(lengths[b]) - 1, h]
return logits(last)
}
// Training over sequences packed into one row, as `prefill_packed`
// packs prompts: `sum_p weights[p] * -log p(targets[p])` for the token
// after each position (`std.nn.loss::linear_cross_entropy` with the
// output head). A weight of 0 leaves a position out; `mask / count`
// gives the mean. With the model
// split across processes, each holds its rows of the output head.
pub entry loss_packed<P: Dim>(
tokens: Tensor[P; i32],
positions: Tensor[P; i32],
segments: Tensor[P; i32],
targets: Tensor[P; i64],
weights: Tensor[P; f32],
) -> f32
where P > 0 {
let states = packed_states(tokens, positions, segments)
return split_cross_entropy<P, H, Vocab / Shards, Shards, T>(
shared(states),
lm_head.weight,
targets,
weights,
)
}
// Each packed position's log-probability of `targets[p]`, as a
// policy-gradient update weighs them.
pub entry log_probs_packed<P: Dim>(
tokens: Tensor[P; i32],
positions: Tensor[P; i32],
segments: Tensor[P; i32],
targets: Tensor[P; i64],
) -> Tensor[P; f32]
where P > 0 {
let states = packed_states(tokens, positions, segments)
return split_token_log_probs<P, H, Vocab / Shards, Shards, T>(
shared(states),
lm_head.weight,
targets,
)
}
// Each packed position's final hidden state, after the last norm: what
// the output head (or a value or reward head) reads.
pub entry hidden_packed<P: Dim>(
tokens: Tensor[P; i32],
positions: Tensor[P; i32],
segments: Tensor[P; i32],
) -> Tensor[P, H; T]
where P > 0 {
return packed_states(tokens, positions, segments)
}
fn packed_states<P: Dim>(
tokens: Tensor[P; i32],
positions: Tensor[P; i32],
segments: Tensor[P; i32],
) -> Tensor[P, H; T]
where P > 0 {
var x = embedding.forward(reshape(tokens, [1, P]))
static for layer in layers {
x = layer.packed(x, positions, segments)
}
return reshape(norm.forward(x), [P, H])
}
// Prompts of several requests packed end to end into one pass of `P`
// tokens, with no padding between them: token `p` is position
// `positions[p]` of prompt `segments[p]`, and its key and value go to row
// `rows[p]` of the caches. A pack may end in padding that writes where no
// request reads (the engine uses position `MaxSeq - 1`, which decoding
// never reaches) with a segment of its own. Returns the logits after
// prompt `m`'s last token, `last[m]`, for up to `Batch` prompts; rows
// past the pass's prompts repeat whichever token `last` names.
pub entry prefill_packed<P: Dim>(
tokens: Tensor[P; i32],
rows: Tensor[P; i32],
positions: Tensor[P; i32],
segments: Tensor[P; i32],
last: Tensor[Batch; i32],
) -> Tensor[Batch, Vocab; T]
where P > 0 {
var x = embedding.forward(reshape(tokens, [1, P]))
static for layer in layers {
x = layer.prefill_packed(x, rows, positions, segments)
}
let ends[m, h] = x[0, cast<i64>(last[m]), h]
return logits(ends)
}
// `prefill_packed`'s prompts and a step for every row in one pass, as
// `decode_rows` takes it (`step_tokens[b]` at `step_positions[b]`): the
// weights are read once for both. A row the prompts fill steps along,
// written where no request reads (the engine uses `MaxSeq - 1`). Returns
// the logits after each prompt, as `prefill_packed`, then each row's step.
pub entry step_packed<P: Dim>(
tokens: Tensor[P; i32],
rows: Tensor[P; i32],
positions: Tensor[P; i32],
segments: Tensor[P; i32],
last: Tensor[Batch; i32],
step_tokens: Tensor[Batch, 1; i32],
step_positions: Tensor[Batch; i32],
) -> Tensor[2 * Batch, Vocab; T]
where P > 0 {
let every = concat(tokens, reshape(step_tokens, [Batch]), axis = 0)
var x = embedding.forward(reshape(every, [1, P + Batch]))
static for layer in layers {
x = layer.step_packed(x, rows, positions, segments, step_positions)
}
let ends[m, h] = x[0, cast<i64>(last[m]), h]
return logits(concat(ends, x[0, P:P + Batch, :], axis = 0))
}
// Paged serving (`linnet.serve`): loaded with `Batch` 1, the caches'
// one row is a pool of `MaxSeq` positions in pages of `PageSize`, and
// each of `Rows` requests takes pages as it grows. `decode_paged` is
// `decode_rows` for those rows, row `b`'s positions lying in the pages
// `table[b]` lists; a row with no request lists the first page, where
// nobody reads.
pub entry decode_paged<Rows: Dim, Pages: Dim>(
tokens: Tensor[Rows, 1; i32],
positions: Tensor[Rows; i32],
table: Tensor[Rows, Pages; i32],
) -> Tensor[Rows, Vocab; T]
where Rows > 0 {
var x = embedding.forward(tokens)
static for layer in layers {
x = layer.decode_paged(x, positions, table)
}
return logits(x[:, 0, :])
}
// Prompts packed end to end over the pool: token `p`, position
// `positions[p]` of row `rows[p]`, is written at `slots[p]` and attends
// over its row's pages (`table[rows[p]]`) up to its position, so a
// prompt can pass in chunks or start after pages it shares. Padding
// writes at the first page's start, which nobody reads. Returns the
// logits after up to `Rows` prompts.
pub entry prefill_paged<P: Dim, Rows: Dim, Pages: Dim>(
tokens: Tensor[P; i32],
positions: Tensor[P; i32],
rows: Tensor[P; i32],
slots: Tensor[P; i32],
last: Tensor[Rows; i32],
table: Tensor[Rows, Pages; i32],
) -> Tensor[Rows, Vocab; T]
where P > 0 {
var x = embedding.forward(reshape(tokens, [1, P]))
static for layer in layers {
x = layer.prefill_paged(x, positions, rows, slots, table)
}
let ends[m, h] = x[0, cast<i64>(last[m]), h]
return logits(ends)
}
// `prefill_paged`'s prompts and a step for every row as `decode_paged`
// takes it, in one pass.
pub entry step_paged<P: Dim, Rows: Dim, Pages: Dim>(
tokens: Tensor[P; i32],
positions: Tensor[P; i32],
rows: Tensor[P; i32],
slots: Tensor[P; i32],
last: Tensor[Rows; i32],
step_tokens: Tensor[Rows, 1; i32],
step_positions: Tensor[Rows; i32],
table: Tensor[Rows, Pages; i32],
) -> Tensor[2 * Rows, Vocab; T]
where
P > 0,
Rows > 0
{
let every = concat(tokens, reshape(step_tokens, [Rows]), axis = 0)
var x = embedding.forward(reshape(every, [1, P + Rows]))
static for layer in layers {
x = layer.step_paged(x, positions, rows, slots, step_positions, table)
}
let ends[m, h] = x[0, cast<i64>(last[m]), h]
return logits(concat(ends, x[0, P:P + Rows, :], axis = 0))
}
// Greedy generation inside the graph: `Steps` tokens after `token` at
// `pos`, each decode step feeding the caches and its `argmax` the next
// step. The range loop is expanded at compile time, so the whole
// generation exports as one program.
pub entry generate<Steps: Dim>(
token: Tensor[Batch, 1; i32],
pos: i32,
) -> Tensor[Batch, Steps; i32]
where Steps > 0 {
var current = token
var produced = fill<i32>([Batch, Steps], 0)
let slots = iota<i32>(Steps)
static for i in 0..Steps {
let next = argmax(step(current, pos + cast<i32>(i)))
let written[b, s] = select(slots[s] == cast<i32>(i), next[b], produced[b, s])
produced = written
current = reshape(next, [Batch, 1])
}
return produced
}
// Sampling with an explicit key: `Steps` tokens drawn from the softmax
// of the logits divided by `temperature`, one derived key per step, so
// the same key gives the same text on every backend.
pub entry sample<Steps: Dim>(
token: Tensor[Batch, 1; i32],
pos: i32,
key: Tensor[2; i64],
temperature: f32,
) -> Tensor[Batch, Steps; i32]
where Steps > 0 {
var current = token
var produced = fill<i32>([Batch, Steps], 0)
let slots = iota<i32>(Steps)
let keys = split<Steps>(key)
static for i in 0..Steps {
let logits = cast<f32>(step(current, pos + cast<i32>(i))) / temperature
let key_i[j] = keys[i, j]
let next = categorical(key_i, logits)
let written[b, s] = select(slots[s] == cast<i32>(i), next[b], produced[b, s])
produced = written
current = reshape(next, [Batch, 1])
}
return produced
}
// Greedy generation that stops early: at most `MaxNew` tokens, or fewer
// once every sequence has produced `eos`. A runtime `while` loop over
// scalar state; the count of tokens produced is returned with them, and
// positions past it hold zeros.
pub entry generate_until<MaxNew: Dim>(
token: Tensor[Batch, 1; i32],
pos: i32,
eos: i32,
) -> (Tensor[Batch, MaxNew; i32], i32)
where MaxNew > 0 {
var current = token
var produced = fill<i32>([Batch, MaxNew], 0)
var count: i32 = 0
var running = true
let slots = iota<i32>(MaxNew)
while running && count < MaxNew {
let next = argmax(step(current, pos + count))
let written[b, s] = select(slots[s] == count, next[b], produced[b, s])
produced = written
current = reshape(next, [Batch, 1])
count = count + 1
let finished = all[b] (next[b] == eos)
running = !finished
}
return (produced, count)
}
fn step(token: Tensor[Batch, 1; i32], pos: i32) -> Tensor[Batch, Vocab; T] {
var x = embedding.forward(token)
static for layer in layers {
x = layer.decode(x, pos)
}
return logits(x[:, 0, :])
}
// Rows' logits from their final hidden states: each shard's slice of the
// vocabulary, the slices side by side.
fn logits<R: Dim>(x: Tensor[R, H; T]) -> Tensor[R, Vocab; T] {
return all_gather<R, Vocab / Shards, Shards, T>(lm_head.forward(shared(norm.forward(x))))
}
fn hidden<B: Dim, S: Dim>(tokens: Tensor[B, S; i32]) -> Tensor[B, S, H; T] {
return norm.forward(hidden_states(tokens))
}
fn hidden_states<B: Dim, S: Dim>(tokens: Tensor[B, S; i32]) -> Tensor[B, S, H; T] {
var x = embedding.forward(tokens)
static for layer in layers {
x = layer.forward(x)
}
return x
}
}rope.linnet3.5 kB
linnet
// Rotary position tables computed from the compile-time sequence length:
// the positions and frequencies come from `iota`, so a model needs no
// precomputed inputs and the tables have exactly the length of the sequence.
//
// Llama 3.1 adds `rope_scaling` (`config.json`, type `llama3`) on top of
// Llama 3's base frequency: a piecewise rescaling of each frequency's
// wavelength that stretches the 8192-token pretraining context out to
// 131072 without retraining. Per frequency `i` (independent of position),
// with wavelength `wavelen[i] = 2 * pi / inv_freq[i]`:
// - a wavelength longer than `original_max_position_embeddings /
// low_freq_factor` is "low frequency": divide it by `factor`;
// - a wavelength shorter than `original_max_position_embeddings /
// high_freq_factor` is "high frequency": leave it alone;
// - in between, linearly blend the unscaled and the divided-by-`factor`
// frequency by how far `original_max_position_embeddings / wavelen[i]`
// sits between `low_freq_factor` and `high_freq_factor`.
// The three bands are two nested `select`s on those comparisons, mirroring
// `transformers`' `torch.where` chain (`_compute_llama3_parameters`). This
// is ordinary elementwise scalar arithmetic on the `[D / 2]` frequency
// table, the same shape `inv_freq` already had, so it needed nothing beyond
// the arithmetic and `select` the rest of the standard library uses -- no
// index expression does arithmetic, since `i` only ever subscripts plainly.
module llama.rope
// Llama 3 stretched the base frequency from 10000 to 500000.
pub const THETA: f32 = 500000.0
// `rope_scaling` in `config.json`: `{"factor": 8.0, "low_freq_factor": 1.0,
// "high_freq_factor": 4.0, "original_max_position_embeddings": 8192,
// "rope_type": "llama3"}`.
pub const FACTOR: f32 = 8.0
pub const LOW_FREQ_FACTOR: f32 = 1.0
pub const HIGH_FREQ_FACTOR: f32 = 4.0
pub const ORIGINAL_MAX_POSITION_EMBEDDINGS: f32 = 8192.0
pub const PI: f32 = 3.14159265
// `(cos, sin)` tables of shape `[S, D]`, each frequency repeated for both
// halves of the head dimension as `rope` expects.
pub fn tables<S: Dim, D: Dim, T: Float>(theta: f32) -> (Tensor[S, D; T], Tensor[S, D; T])
where D % 2 == 0 {
let base_inv_freq[i] = exp(-(cast<f32>(iota<i64>(D / 2)[i]) * 2.0 / cast<f32>(D)) * log(theta))
let wavelen[i] = (2.0 * PI) / base_inv_freq[i]
let low_freq_wavelen = ORIGINAL_MAX_POSITION_EMBEDDINGS / LOW_FREQ_FACTOR
let high_freq_wavelen = ORIGINAL_MAX_POSITION_EMBEDDINGS / HIGH_FREQ_FACTOR
// The blend used in the medium band; only meaningful there, since the
// `select` chain below only reads it when neither the low- nor the
// high-frequency case applies. `smooth == 1` sits at the
// `high_freq_wavelen` boundary (continuous with the unscaled band just
// beyond it) and `smooth == 0` sits at the `low_freq_wavelen` boundary
// (continuous with the divided-by-`FACTOR` band just beyond that), so
// `(1 - smooth)` -- not `smooth` -- weights the scaled term.
let smooth[i] =
(ORIGINAL_MAX_POSITION_EMBEDDINGS / wavelen[i] - LOW_FREQ_FACTOR) /
(HIGH_FREQ_FACTOR - LOW_FREQ_FACTOR)
let smoothed[i] = (1.0 - smooth[i]) * (base_inv_freq[i] / FACTOR) + smooth[i] * base_inv_freq[i]
let inv_freq[i] = select(
wavelen[i] > low_freq_wavelen,
base_inv_freq[i] / FACTOR,
select(wavelen[i] < high_freq_wavelen, base_inv_freq[i], smoothed[i]),
)
let angles[s, i] = cast<f32>(iota<i64>(S)[s]) * inv_freq[i]
let full = concat(angles, angles, axis = -1)
return (cast<T>(cos(full)), cast<T>(sin(full)))
}model-00001-of-00004.safetensorsHugging Face ↗
model-00002-of-00004.safetensorsHugging Face ↗