A long-context mixture-of-experts transformer in PyTorch, built from the Apache-2.0 OpenAI gpt-oss reference: top-k expert routing for sparse compute at dense capacity, rotary scaling configured to 131,072 tokens, grouped-query attention with learned sinks, and deterministic sampling controls.
The forward pass: router logits, a top-k selection, a softmax over only the selected experts, and a batched einsum that touches those experts alone. Grouped-query attention with query heads a multiple of the key-value heads. A learnable per-head attention sink, appended as an extra softmax column and dropped after normalisation, so a head can attend to nothing without renormalising the rest. Sliding-window attention on alternating layers. And an append-only KV cache with snapshot and restore, so one prefill serves many decode sweeps.
Long context is a rotary problem before it is a memory problem. The scaling combines YaRN concentration, 0.1 * log(s) + 1.0, with NTK-by-parts cutpoints and a linear ramp between them: the fastest rotary clocks are left alone, the slowest are stretched, and the band between is blended rather than switched. Sampling is deterministic — logit bias, then repetition, presence and frequency penalties, then top-k, top-p and min-p truncation, then a multinomial draw, with temperature zero collapsing to greedy.
Six invariants hold it together, and breaking any of them fails silently. Rotary embeddings use the global position, cache offset plus local index, and never reset per window. Masks are built against the cache's own width with that same offset. Decode appends exactly one key-value row per layer per step. The query-to-key-value head ratio matches the checkpoint. Top-k ties, dtype and the kernel can each change expert selection. And the seed with the truncation parameters defines a run. The one easiest to get wrong is named: an off-by-one diagonal in the mask at the moment prefill hands over to decode.
Its cost is stated in the open. KV bytes come to roughly 2 * L * B * T * H_kv * d * dtype_bytes; bf16, 24 layers, batch one, 131,072 tokens, eight key-value heads and head dimension 64 land at about 6.0 GiB. It ships the instruments to measure itself: needle-in-a-haystack and passkey evaluations, a Responses-style local server reporting prefill and decode milliseconds, a dispatch microbenchmark, and router telemetry — per-layer expert counts, utilisation, entropy, load balance and routing margin.