Day 25: the fused read dispatches to the kernel backend, paged goes flash-close to naive, 186 green

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: pjdurden

docs/daily/day-25.md

@@ -0,0 +1,111 @@
+---
+title: "Day 25: the fused read dispatches to the kernel backend"
+parent: Daily log
+nav_order: 25
+---
+
+# Day 25: the fused read dispatches to the kernel backend
+
+Date: 2026-07-16 · Week 6 · Phase 2 Paged memory
+
+## What I added today
+One line moved in `src/nanoserve/cache.py`, and it is the line the whole of Week 6
+was building toward. `PagedKVCache.paged_attention`, the cache's fused read, called
+`paged_attention_reference` directly through Day 24: the byte-exact torch oracle,
+the same SDPA over the assembled slots the tests compare against. Day 25 routes it
+through `paged_attention_dispatch`, the Day-23 entry point that asks
+`select_backend(q.device)` which read this device can actually run. On a CUDA tensor
+with Triton installed that is the `@triton.jit` kernel; on a CPU it is the Day-22
+tlsim model, the same streaming loop written in torch. The engine gets the kernel on
+a card and the correct-and-slow model on a laptop, and the cache never learns which.
+
+Two new tests in `tests/test_cache.py` pin the wiring. One spies on the dispatcher
+symbol the cache imports and asserts the read calls it exactly once, handed this
+layer's own K/V pools and the slot mapping for the whole history, then hands its
+result straight back. The other feeds the cache's pool through both the dispatched
+read and `paged_attention_reference` and pins them to `atol=1e-5`, the same tolerance
+the kernel itself is held to. The consequence rippled into three older equality tests
+that had pinned the cache read `torch.equal` against a contiguous SDPA: they now carry
+the flash tolerance, because the model no longer runs the byte-exact reference. Suite
+**186 green** (5 GPU-gated skips), ruff clean.
+
+## Why it matters
+Until today Week 6 was a kernel in a file. Day 22 wrote the fused CPU model, Day 23
+wrote the real Triton kernel, Day 24 built the instrument to read its scaling, and
+every one of them was reachable only from a test. The model's actual attention still
+ran `paged_attention_reference`, the oracle. This is the wiring step that turns the
+kernel into the engine's real attention: the paged forward now dispatches, so the
+same code path picks up the GPU kernel when a card is present and falls back to the
+model that runs anywhere when it is not. The reference did not disappear, it changed
+jobs. It is the oracle both backends are graded against, not the path either one runs.
+
+The honest cost is that paged attention is no longer bit-identical to the naive
+contiguous path. Both backends stream the online softmax, which reassociates the
+exponent sums, so the output lands a few ulps off a plain SDPA (about 7e-6 on the
+two-layer test model) rather than exactly on it. That is not a regression, it is the
+accuracy trade every flash-attention kernel makes, and it belongs on the model's path
+precisely because that is where the kernel now lives. The token-level tests still hold
+because argmax over these logits does not flip at 1e-5; the logit-level test that
+asked for 1e-6 now asks for 1e-5, the tolerance the streaming softmax actually
+delivers.
+
+## What I learned
+1. **Byte-identity was a property of the oracle, not of paging.** Every "paging
+   changes nothing" test through Day 24 passed `torch.equal` because the cache read
+   *was* the reference, running the identical causal fp32 softmax by exact index. The
+   moment the model runs the streaming kernel instead, the scatter and gather stay
+   exact (same K/V, same slots) but the softmax summation order moves, and equality
+   becomes closeness. The tests that flipped from `equal` to `allclose` were not
+   loosened to hide a bug; they caught exactly the numerics change the dispatch
+   introduced, which is what a good test does.
+2. **The right tolerance is the one the backend is already held to.** The kernel
+   tests pin tlsim and Triton to `paged_attention_reference` at `atol=1e-5, rtol=1e-4`.
+   When the cache read started running that same backend, the correct new tolerance
+   for the cache tests was not a number I picked to make green, it was that same
+   1e-5, propagated up one level. A tolerance that means something is a tolerance
+   borrowed from where the approximation is defined.
+3. **A dispatcher makes the fallback the default and the kernel the exception.** The
+   branch reads `"triton" if device.type == "cuda" and has_triton() else "tlsim"`, so
+   the CPU model is what runs unless both a card and the package are present. That is
+   the right default for a repo that must run on a laptop: the slow-and-correct path
+   is the one you get for free, and the fast path is the one that has to prove its
+   preconditions. The engine is portable first and fast where it can be, not the
+   other way around.
+
+## Diagram
+[paged-read-dispatch.png](../diagrams/paged-read-dispatch.png). The cache read at the
+top, the `select_backend` pill below it, and the two branches: the Triton kernel on
+`cuda & has_triton()`, the tlsim model on `else`. The reference sits dashed in the
+middle, with both backends held to it at `atol=1e-5`, no longer on the path. Both
+converge to the one attention output that flows back to `gqa_attention`. The banner
+carries the trade: Day 24 ran the byte-exact reference, Day 25 runs the streaming
+softmax, so paged is flash-close to naive, not bit-identical.
+
+## Tomorrow
+Measure the dispatch where it matters. The read benchmark and the scaling fit
+(Days 20 and 24) time `gather` and `fused` as standalone functions; wire the runner
+to time the cache's *dispatched* read end to end, so the sweep reports the path the
+engine actually takes rather than a function called in isolation. On CPU that pins the
+tlsim constant against the gather constant under the real slot mapping; on a card it is
+the first number that shows the kernel bending the constant down, which is the whole
+claim Week 6 has been setting up to make. Keep the reference as the correctness gate
+the timed path is checked against before any number is trusted.
+
+## Post angle
+Day 25 of building an LLM inference engine from scratch. Today was one moved line and
+it is the one the whole kernel week was for. My paged cache's attention read was still
+calling the byte-exact torch reference, the oracle. The real Triton kernel and its CPU
+model existed but only the tests ever touched them. Today the cache read dispatches:
+`select_backend(q.device)` picks the Triton kernel on a card and the streaming CPU
+model on a laptop, and the model's actual attention runs whichever the box can run.
+The reference did not go away, it changed jobs, from the path the model runs to the
+oracle both backends are graded against. The honest part: paged attention is no longer
+bit-identical to the naive path. Both backends stream the online softmax, which
+reassociates the exponent sums, so the output lands about 7e-6 off a plain SDPA rather
+than exactly on it. That is not a bug, it is the accuracy trade every flash-attention
+kernel makes, and three of my "paging changes nothing" tests flipped from `torch.equal`
+to `allclose` to carry it. The tolerance I gave them was not invented, it is the same
+1e-5 the kernel is already held to against the reference, propagated up one level. This
+is the shape vLLM and SGLang ship: a dispatcher that runs the fast kernel where the
+hardware allows and a correct fallback everywhere else, with the reference kept as the
+gate, not the path. 186 green.

docs/diagrams/paged-read-dispatch.png

@@ -18,7 +18,7 @@
 
 import torch
 
-from .kernels.paged_attention import paged_attention_reference
+from .kernels.triton_paged_attention import paged_attention as paged_attention_dispatch
 
 
 class NaiveKVCache:
@@ -301,8 +301,10 @@ class PagedKVCache:
     the tests still compare against. `paged_attention` is the Week-6 fused read: it
     attends directly over the scattered blocks through the Day-18 reference and
     hands `gqa_attention` the attention output, never rebuilding the history. The
-    model runs on the fused read now; a Triton kernel later slots in behind the same
-    signature, held to `paged_attention_reference`.
+    model runs on the fused read now, and since Day 25 that read dispatches through
+    `select_backend`: the Triton kernel on a card, the tlsim CPU model on a laptop,
+    both held to `paged_attention_reference` as the oracle. The reference is the
+    check the dispatch is graded against, no longer the path the model runs.
     """
 
     def __init__(self, config, allocator: BlockAllocator):
@@ -401,24 +403,33 @@ def paged_attention(
 
         This is Week 6's payoff wired into the cache. Rather than hand attention the
         contiguous history to score itself (`append`), the cache does the read: it
-        writes the K/V (see `_write`), then runs `paged_attention_reference` over its
-        own scattered pool, reading each past token through its slot and never
-        rebuilding the contiguous buffer. `gqa_attention` calls this and takes the
-        attention output straight back.
+        writes the K/V (see `_write`), then dispatches the paged read over its own
+        scattered pool, reading each past token through its slot and never rebuilding
+        the contiguous buffer. `gqa_attention` calls this and takes the attention
+        output straight back.
+
+        Since Day 25 the read goes through `paged_attention_dispatch`, the Day-23
+        entry point that asks `select_backend` which read the query's device can run:
+        the Triton kernel on a CUDA tensor, the Day-22 tlsim model on the CPU. The
+        engine gets the kernel on a card and the correct-and-slow model on a laptop
+        without the cache knowing which, and `paged_attention_reference` steps back to
+        being the oracle both backends are graded against, not the path either runs.
 
         k, v: [1, num_kv_heads, new_seq, head_dim] this step's compact K/V.
         q:    [1, num_attention_heads, new_seq, head_dim] this step's rotated query.
         n_rep: GQA repeat factor (`config.num_kv_groups`).
         scale: softmax scale; defaults to `head_dim ** -0.5`, matching `gqa_attention`.
 
         Returns [1, num_attention_heads, new_seq, head_dim], the attention output
-        before o_proj, byte-identical to a contiguous SDPA over the same K/V. The
-        pool and table are left exactly as `append` would leave them; only the
-        return value differs.
+        before o_proj. It matches a contiguous SDPA over the same K/V to a few ulps
+        rather than bit for bit: both backends stream the online softmax, which
+        reassociates the exponent sums, the accuracy trade every flash-attention
+        kernel makes. The pool and table are left exactly as `append` would leave
+        them; only the return value differs.
         """
         self._write(layer, k, v)
         slot_mapping = self._slots_for(range(self.table.num_tokens), q.device)
-        return paged_attention_reference(
+        return paged_attention_dispatch(
             q, self.k_pool[layer], self.v_pool[layer], slot_mapping, n_rep, scale
         )