kernel-fun 0.2.0.dev1__tar.gz

kernel_fun-0.2.0.dev1/.github/workflows/release.yml

@@ -0,0 +1,95 @@
+# Publishing kernel-fun to PyPI.
+#
+# Auth is PyPI Trusted Publishing (OIDC): no API token is stored in this repo or in GitHub
+# secrets. PyPI matches four claims out of the token GitHub mints for this job -- owner,
+# repository, THIS FILENAME, and the environment named below. Renaming this file, or
+# renaming an environment, breaks publishing with a claims mismatch until the publisher
+# entry on PyPI is edited to match. See README.md "Releasing".
+#
+#   TestPyPI   workflow_dispatch, target=testpypi   -> environment `testpypi`
+#   PyPI       push a `v*` tag (or dispatch)        -> environment `pypi`
+#
+# A version on PyPI can never be reused, so the build job refuses a tag that disagrees with
+# `__version__`; the tag is the thing humans read, and the metadata is what actually ships.
+name: Release
+
+on:
+  push:
+    tags: ["v*"]
+  workflow_dispatch:
+    inputs:
+      target:
+        description: "Index to publish to"
+        type: choice
+        options: [testpypi, pypi]
+        default: testpypi
+
+permissions: {}
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v7
+      - uses: astral-sh/setup-uv@v10.1.0
+        with:
+          python-version: "3.12"
+
+      - name: Build sdist and wheel
+        run: uv build
+
+      # Metadata that fails here fails AFTER upload otherwise, and by then the version is
+      # burned. Cheap gate: PyPI's own renderer/validator, run before anything is sent.
+      - name: Check metadata
+        run: uvx twine check --strict dist/*
+
+      # The wheel's version comes from src/kernel_fun/__init__.py. If someone tags v0.3.0
+      # without bumping it, the upload succeeds and ships 0.2.0 under a 0.3.0 tag -- an
+      # inconsistency that cannot be corrected in place. Refuse instead.
+      - name: Tag agrees with __version__
+        if: startsWith(github.ref, 'refs/tags/v')
+        run: |
+          tag="${GITHUB_REF_NAME#v}"
+          built=$(ls dist/*.whl | sed -E 's/.*kernel_fun-(.+)-py3-none-any\.whl/\1/')
+          echo "tag=$tag  built=$built"
+          [ "$tag" = "$built" ] || { echo "::error::tag v$tag != __version__ $built"; exit 1; }
+
+      - uses: actions/upload-artifact@v7
+        with:
+          name: dist
+          path: dist/
+          if-no-files-found: error
+
+  testpypi:
+    needs: build
+    if: github.event_name == 'workflow_dispatch' && inputs.target == 'testpypi'
+    runs-on: ubuntu-latest
+    environment:
+      name: testpypi
+      url: https://test.pypi.org/p/kernel-fun
+    permissions:
+      id-token: write  # mint the OIDC token PyPI trades for a short-lived upload token
+    steps:
+      - uses: actions/download-artifact@v8
+        with:
+          name: dist
+          path: dist/
+      - uses: pypa/gh-action-pypi-publish@v1.14.2
+        with:
+          repository-url: https://test.pypi.org/legacy/
+
+  pypi:
+    needs: build
+    if: startsWith(github.ref, 'refs/tags/v') || (github.event_name == 'workflow_dispatch' && inputs.target == 'pypi')
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/kernel-fun
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/download-artifact@v8
+        with:
+          name: dist
+          path: dist/
+      - uses: pypa/gh-action-pypi-publish@v1.14.2

kernel_fun-0.2.0.dev1/.gitignore

@@ -0,0 +1,31 @@
+# Build artifacts. This repo, unlike the research ladder, IS installable, so these are
+# reachable here in a way they were not there.
+dist/
+build/
+*.egg-info/
+
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.ruff_cache/
+
+# Generated Triton artifacts. The training image points TRITON_CACHE_DIR outside the tree;
+# a native run on a workstation would land here.
+.triton-cache/
+
+# A .venv inside the tree is how the old repo lost a day: it shadowed the container's conda
+# python and an image that definitely had torch reported "No module named torch". Install
+# this package INTO the image's python (`pip install .`), or put its src on PYTHONPATH.
+.venv/
+.python-version
+
+# macOS Finder droppings.
+.DS_Store
+*.swp
+
+# Profiler captures: 5-50MB each and git history is forever. tools/prodtime.py and
+# tools/graphprobe.py print their numbers; the captures leave a session by rsync.
+*.ncu-rep
+*.nsys-rep
+*.qdrep
+*.sqlite

kernel_fun-0.2.0.dev1/ALGORITHM.md

@@ -0,0 +1,229 @@
+# KDA — the math these kernels implement
+
+Everything below is fla 0.5.2's own decomposition (`fla/ops/kda/{chunk_fwd.py,chunk_bwd.py,
+chunk_intra.py,wy_fast.py}` and `fla/ops/common/chunk_delta_h.py`), stated precisely so each
+port has one written contract. All logs are base-2 (fla multiplies the raw log-space gate by
+`RCP_LN2 = 1/ln2` in the cumsum), all gates are ≤ 0 and decreasing down a chunk, and `BT` is
+the chunk size (64), `BC = 16` the sub-chunk. Per (batch b, value-head hv); q/k live at head
+`h = hv // (HV//H)`. `scale = K^-0.5`.
+
+This package ports the forward scan+readout and four of the backward's seven stages out of
+fla. Everything else is fla's, called stage by stage from `kda/chain.py` — which is the
+authoritative table of who owns what; this file is the math each stage implements.
+
+| where | module | replaces |
+|---|---|---|
+| fwd, Aqk zero-fill | `kda/_kernels/fwd_intra_triton.py` | a `masked_fill` over the whole tile |
+| fwd, scan + o | `kda/_kernels/fwd_state.py` | `chunk_gated_delta_rule_fwd_h` + the o stage |
+| bwd 2, re-scan (B1) | `kda/_kernels/bwd_scan.py` | `chunk_gated_delta_rule_fwd_h` |
+| bwd 4, dhu (B2a) | `kda/_kernels/bwd_dhu.py` | `chunk_gated_delta_rule_bwd_dhu` |
+| bwd 5, wy_dqkg (B2b) | `kda/_kernels/bwd_wy.py` | `chunk_kda_bwd_wy_dqkg_fused` |
+| bwd 6, intra | `kda/_kernels/bwd_intra.py` (+`bwd_intra_triton.py`) | `chunk_kda_bwd_intra` |
+| bwd 7, dg cumsum | — | folded into stage 6's epilogue |
+
+The gate activation (`-exp(A_log)·softplus(g + dt_bias)`, when `use_gate_in_kernel`) is
+fused into stage 1's cumsum via fla's `kda_gate_chunk_cumsum`, not run as eager fp32 torch
+ops. At prod8192 the eager form is four passes over a [B,T,HV,K] fp32 tensor — ~2ms and
+~2GiB of saved activations *per layer*, against 1.9ms for the entire fused stage.
+
+---
+
+# Forward
+
+## Stage 1 — cumsum (fla, stays Triton)
+
+    g2[t, d] = (1/ln2) * sum_{s <= t, s in chunk(t)} g[s, d]        fp32 [B, T, HV, K]
+
+Chunk-local inclusive cumsum, per dimension. `G[c, d] = g2[last row of chunk c, d]` is the
+per-chunk total decay — a K-vector, where a scalar-gate op (gdn) would have a scalar.
+
+## Stage 2 — intra / WY (fla, stays Triton)
+
+`chunk_kda_fwd_intra` produces, per chunk (row indices i, j within the chunk):
+
+    Aqk[i, j] = scale * sum_d q[i,d] * k[j,d] * exp2(g2[i,d] - g2[j,d])     for i >= j, else 0
+    Akk       = the solved lower-triangular WY inverse, beta folded in
+    w[i, :]   = (Akk @ (beta * exp2(g2) * k))[i, :]                         bf16 [B, T, HV, K]
+    u[i, :]   = (Akk @ (beta * v))[i, :]                                    bf16 [B, T, HV, V]
+    kg[i, d]  = k[i,d] * exp2(G[d] - g2[i,d])                               bf16 [B, T, HV, K]
+    qg        = q * exp2(g2)   (only materialized when disable_recompute; None on our path)
+
+Two things to internalize:
+
+- **`scale` ships inside Aqk** (the intra kernels multiply it in), while the `q @ h` readout
+  term gets `scale` applied in the o stage. Don't apply it twice.
+- **The per-dim gate sits INSIDE Aqk's dot product.** `exp2(g2_i - g2_j)` would be a scalar
+  under a per-head gate, so it could be factored out of one MMA and applied afterwards in
+  SIMT. Per-dim, that factorization needs `exp2(+g2_j)` on one operand — unbounded — so fla
+  computes Aqk in sub-chunks (BC=16) where relative gates stay bounded. That is Triton work
+  worth keeping: the CuTe kernel LOADS Aqk instead of computing it.
+
+## Stage 3+4 — scan + o (`kda/_kernels/fwd_state.py`, fused)
+
+Serial over chunks c, state h [K, V] fp32:
+
+    v'_c = u_c - w_c @ h_c                                  (MMA "WH" + subtract)
+    o_c  = scale * (q_c * exp2(g2_c)) @ h_c + Aqk_c @ v'_c  (MMAs "OH" + "OI")
+    h_{c+1}[d, :] = exp2(G_c[d]) * h_c[d, :] + (kg_c^T @ v'_c)[d, :]    (MMA "DH" + decay)
+
+after the last chunk, `ht = h` (fp32, TMA store). fla stores v' (`v_new`) and per-chunk h to
+HBM between its two kernels; the fusion keeps both in smem/tmem — that traffic is the win.
+
+Notes that cost time to rediscover: the state decay multiply indexes the tmem fragment's row
+(Ld16x256b lanes span rows r and r+8), and OH's q operand must be pre-gated in SIMT
+(`qg = q·exp2(g2)`, per-dim) from the q tile plus a g2 tile (fp32 [BT,K], 32KB/stage smem).
+
+---
+
+# Backward
+
+Residuals from the forward: `(q, k, v, g2, beta, Aqk, Akk, h0)` — exactly what fla's own
+autograd Function saves when `disable_recompute=False`. The forward must preserve that
+rounding contract: Aqk/Akk stored bf16 by intra, g2 the fla cumsum, v_new/h recomputed by
+the backward itself.
+
+One deliberate deviation, and it is fla's own behaviour rather than ours: under
+`use_gate_in_kernel`, **g2 is not saved**. The backward recomputes it from the raw (and
+half-size) `g` with one kernel launch, which is what fla's recompute path does — saving it
+would cost 1GiB per layer at prod8192 for something a single launch reproduces.
+
+## The seven launches
+
+1. **recompute** (`recompute_w_u_fwd`): w = Akk @ (β·exp2(g2)·k), u = Akk @ (β·v),
+   qg = q·exp2(g2), kg = k·exp2(G−g2). All bf16 [B,T,HV,K/V].
+2. **re-scan** (`chunk_gated_delta_rule_fwd_h`): the forward recurrence again, this time
+   materializing every per-chunk state: h_c [K,V] checkpoints (bf16, [B,NT,HV,K,V]) and
+   v'_c = u_c − w_c @ h_c (`v_new`). Serial over chunks.
+3. **dAv** (`chunk_kda_bwd_dAv`): dAqk[r,s] = scale·Σ_v do[r,v]·v'[s,v] (lower tri, fp32
+   [B,T,HV,BT]); dv[r,:] = Σ_{s≥r} Aqk[s,r]·do[s,:] (Aqk^T @ do, upper-masked).
+4. **dhu** (`chunk_gated_delta_rule_bwd_dhu`): the reverse scan. State dh [K,V] fp32; fla
+   folds the per-chunk decay on the K axis, updates dv (dv2 = dv + w-side term) per chunk,
+   and emits dh checkpoints (bf16 [B,NT,HV,K,V]) + dh0. Serial, reverse.
+5. **wy_dqkg** (`chunk_kda_bwd_wy_dqkg_fused`), per chunk, consuming h_c and dh_c:
+   - dq[r,d] = scale·exp2(g2[r,d])·Σ_v do[r,v]·h_c[d,v]                    (h consumer)
+   - dk[r,d] = exp2(G[d]−g2[r,d])·Σ_v v'[r,v]·dh_c[d,v] + WY corrections   (dh consumer)
+   - dw[r,d] = −Σ_v dv[r,v]·h_c[d,v]; folded via Akk into dk/db/dAkk       (h consumer)
+   - dgk[d] += exp2(G[d])·Σ_v h_c[d,v]·dh_c[d,v] (last row of chunk)       (both)
+   - dv2 = β·(Akk^T @ dv); db += rowsum terms; dAkk (strict lower, fp32)
+   - dg[r,d] = q·dq − k·dk + last-row dgk + kg·(Akk@dw)·β  (the inter-chunk part)
+6. **intra** (`chunk_kda_bwd_intra`): the intra-chunk dAqk/dAkk consumers — **the dominant
+   stage, half of fla's backward**; see below.
+7. **dg cumsum**: dg = chunk-local REVERSE cumsum of dg (fp32).
+
+Ours are 2, 4, 5 and 6, and 7 disappears into 6's epilogue. Stages 2/4/5 are the "B1/B2"
+fusions listed as standing headroom in the first version of this file: B1 re-runs the
+forward scan with h checkpointed via TMA store off the critical path, B2a keeps dh^T
+tmem-resident through the reverse scan, and B2b is a full-K restructure of fla's fused
+kernel — the same math over one K slab instead of NK, so v_new/do/dv are read once each
+instead of four times.
+
+GVA (HV > H): dq/dk are produced at HV and reduced to H after stage 6.
+
+## Where the time and bytes go (prod8192 = B16 × T8192 × H16, K128/V256, b300)
+
+Stage times measured 2026-08-18: 25.3ms total — intra 12.62, wy_dqkg 6.15, dhu 2.22,
+fwd_h 1.88, recompute 1.18, dAv 0.77, dg cumsum 0.51. Inter-stage HBM: h 2.1GB + dh 2.1GB
++ v_new 1.1GB + w/u/qg/kg 2.7GB + dv 2×1.1GB + fp32 dq/dk/dg 3.2GB + dAqk/dAkk 1.1GB.
+
+## Stage 6, bwd_intra — half the backward
+
+### What it computes
+
+Adds the intra-chunk gradient parts, given dAqk and dAkk (both fp32, strictly-masked). With
+r, s chunk-local rows, all per (b, hv, chunk), and every gate factor per-dim d:
+
+    dq[r,d]  += Σ_{s ≤ r} dAqk[r,s] · k[s,d] · exp2(g2[r,d] − g2[s,d])
+    dwk[r,d]  = Σ_{s ≤ r} dAkk[r,s] · k[s,d] · exp2(g2[r,d] − g2[s,d])
+    db[r]    += Σ_d dwk[r,d] · k[r,d]
+    dk[s,d]  += Σ_{r ≥ s} (dAqk[r,s]·q[r,d] + dAkk[r,s]·β_r·k[r,d]) · exp2(g2[r,d] − g2[s,d])
+                + β_s · dwk[s,d]
+    dg[r,d]  += q[r,d]·dq_intra[r,d] + β_r·dwk[r,d]·k[r,d] − dkt[r,d]·k[r,d]
+                (dkt = the Σ_{r ≥ s} term above, i.e. the column-side accumulation)
+
+### The numerics law (do not relitigate)
+
+Every exp2 argument above is `g2[r,d] − g2[s,d]` with r ≥ s — one-sided, ≤ 0, safe at any
+gate magnitude. KDA's real init makes decay ~16 log2 units PER STEP (`exp(A_log) ∈ [1,16]`),
+so any factorization through a reference row m — `exp2(g_r − g_m)·exp2(g_m − g_s)` — has one
+factor of the pair unbounded unless m sits between r and s. On a diagonal block no single m
+does, which is why fla's diagonals are scalar j-loops and why the SAFE_GATE midpoint trick
+carries a bounded-gate contract this op does not satisfy. Cross-sub-chunk blocks ARE safe
+(any boundary between the blocks separates all (r,s) pairs) — fla already exploits that with
+MMAs there. Falsified before: two-sided midpoint/binary-tree diagonal forms — they overflow
+(a production NaN, 2026-08-17) or don't win. `test_kimi_delta_attention_cute_extreme_decay`
+in `src/test/nn/attention/kda_test.py` is the guard for exactly this class.
+
+### Why fla's version costs 12.6ms
+
+Grid (NK·NC, NT, B·HV) = (16, 128, 256) → **524k CTAs of [BC=16, BK=32] work**:
+
+- **Stream multiplicity.** Each of the NK=4 K-slabs re-reads the same [BC,BC] dAqk/dAkk
+  tiles (fp32 → 4× the 1.1GB), and each of the NC=4 sub-chunk owners re-reads its j-loop
+  partners' q/k/g tiles. db is written as an NK-slab and reduced afterwards.
+- **exp2 volume.** ~1.1e10 exp2 at prod8192 → ~2.5ms of pure SFU even at perfect
+  utilization, and utilization is poor at [16,32] granularity. The *inherent* one-sided
+  minimum is ~4× smaller (~2.3e9, ~0.5–0.6ms) IF each factor is computed once and shared
+  between the dAqk/dAkk products AND between the row-side (dq) and column-side (dk/dg)
+  passes. fla shares the first pair but recomputes across passes — they live in different
+  CTAs.
+- **Latency.** 16-iteration serial scalar loops of tiny loads with nothing to overlap them.
+
+Floors: traffic ~7.5GB at one-read-one-write → ~1.1ms; exp2 ~0.6ms (overlappable). The
+honest target is ~2.5–3ms; fla is 5× above it.
+
+### What `kda/_kernels/bwd_intra.py` does instead (6.34ms SIMT; ~5.0ms with the MMA path)
+
+One CTA per (chunk, b·hv), 512 threads as (32 d-lanes) × (16 row-lanes) — one warp spans a
+full row. Each thread walks one row per 16-block with mirror pairing (even blocks r0+rlane,
+odd blocks r0+15−rlane) so diagonal work is uniform per thread, and owns 4 **strided**
+columns d = lane + 32c (consecutive-column ownership was 4-way bank-conflicted on every
+scalar load — 27ms). q/k/g2/dq/dk/dg and the [64,64] dA tiles are read once, dq2/dk2/dg2
+written once, db completed in-CTA: the 4× multiplicities are gone.
+
+Cross-block pairs factor through the **s-block end** boundary e(s) = 16(j+1):
+`exp2(g_r − g_s) = exp2(g_r − g_e)·exp2(g_e − g_s)`, both exponents ≤ 0 at any gate
+magnitude (r ≥ e > s on a decreasing g) — the same safety class as fla's own `kg` operand.
+Diagonal pairs (same 16-block) keep exactly one one-sided exp2 per (r,s,d): nothing inside a
+diagonal block is ever factorized. All three prescale arrays (kb, qb, kbb) are built in one
+phase up front, so the two sweeps fuse per block with no inter-sweep barrier.
+
+Outputs land ~1e-7..2e-6 abs of a fp64 reference, where fla's tf32 `tl.dot`s land 3e-4..1.3e-3.
+
+The shipped kernel additionally runs the **cross-sub-chunk** blocks on tcgen05 MMAs rather
+than SIMT (the diagonals stay scalar j-loops — see the numerics law above, which is not
+negotiable), and folds two things into its epilogue that fla spends separate work on: the
+chunk-local reverse cumsum of dg (deleting stage 7 entirely) and the bf16 cast of dq/dk,
+which makes the wrapper's casts no-ops. The dg fold is unconditional; the bf16 emit is
+gated on `HV == H`, because the GVA reduction that follows must sum in fp32.
+
+Constraints: K=128, BT=64, fixed-length, no `safe_gate`, and a grid of at least 1024 CTAs
+(smaller grids underfill a 148-SM box and the per-call marshaling isn't amortized — T512
+regressed to 0.90× without the gate). Off that box the wrapper falls back to
+`bwd_intra_triton.py` (the same math restructured in Triton, 9.26ms, K ≤ 128), which falls
+back to fla's kernel for varlen/`safe_gate`. **That floor is per stage, not per call**: a
+shape can clear `is_supported` and still run this one stage on Triton, which is why the
+test arms in `src/test/nn/attention/kda_test.py` are sized to clear it explicitly.
+
+### Register-pressure notes for whoever edits the CuTe kernel
+
+ptxas left alone targets 64 registers (dynamic smem hides the 1-CTA/SM cap from it) and
+spills 1–2KB/thread. `min_blocks_per_mp=1` plus `--maxrregcount=128` (frozen as `_MAXREG`)
+and *un*-unrolling the inner 16-iteration loops (`cutlass.range(unroll=4..8)`; full unroll
+lets ptxas hoist whole load batches) gets `LOCAL_SIZE_BYTES=0`. The incoming-grad gmem reads
+cost 3.25ms read-at-use; prefetching them a compute-section ahead is worth ~3ms. 1024
+threads was falsified twice (0.45×, 64-reg spills), as was holding P1/P2 in smem (27.5ms).
+
+## Standing headroom
+
+The B1/B2 fusions listed here originally are now stages 2, 4 and 5. What is left, in the
+order the ladder is working on it:
+
+- **recompute_w_u** (stage 1, 2 × 1.17ms) and the forward cumsum (0.32ms) — both fla's,
+  both counter-blocked on the profiling box.
+- **wy_dqkg** at 4.99ms is now the largest single stage.
+- Remaining intra headroom is ~4ms (latency-bound loads/prescale/epilogue at 16 warps/SM).
+
+None of these are edited here. This directory is a vendored snapshot; kernel work happens
+in the `kernel-fun-2` ladder, where a bench row can say whether an idea was worth keeping —
+three of idea 004's did not survive that test.
+- `recompute_w_u`/`dAv`/cumsum (2.4ms combined) are near their traffic floor — leave them.

kernel_fun-0.2.0.dev1/INTEGRATION.md

@@ -0,0 +1,222 @@
+# Landing this in OLMo-core
+
+The last time kernels from this repo landed in the training repo, the ladder reported −3%
+tokens/sec and it took two days to establish that there was no steady-state regression at
+all (the ladder repo's `kernels/gnorm/LADDER-REGRESSION.md`). Almost none of that time went
+into the kernels.
+This checklist exists so that does not happen twice.
+
+## 1. The call site
+
+**The unit of sharing is this package** (state as of 2026-09-08; the repo split below is
+2026-09-09). OLMo-core gets it one of two ways, on two branches that are otherwise
+identical:
+
+- `caleb/cute-kda-vendored` — the package tree copied to `src/olmo_core/kernel_fun/`,
+  byte-identical to this repo's `src/kernel_fun/` except the top-level `__init__.py`, which
+  carries `VENDORED_FROM`. No private dependency, so no tokens. scaling-ladders'
+  `caleb/cute-kda` points here.
+- `caleb/cute-kda` (PR #837) — installs the package as a `kernel-fun` extra pinned to a sha.
+  Needs a GitHub token everywhere the lock resolves.
+
+> **The shas in both places changed meaning on 2026-09-09**, when this package moved out of
+> the ladder repo — which held the name `allenai/kernel-fun` until that day and is now
+> `allenai/kernel-fun-dev`, and where this tree was `packages/kernel-fun/` — into the repo
+> that took the name, **`allenai/kernel-fun`** (this one). `VENDORED_FROM` and the pinned
+> extra now name a commit HERE, and the extra's URL drops
+> `#subdirectory=packages/kernel-fun`. The last pre-split vendoring was ladder sha
+> `7a6983b`, which exists only in `kernel-fun-dev`; the ladder sha a family's *kernels*
+> were cut from lives on in `src/kernel_fun/<family>/_provenance.py::SOURCE_COMMIT`
+> (`6fc6309` today) and is a different number. When re-vendoring, bump `VENDORED_FROM` to a
+> full kernel-fun sha and say so in the commit message — a bare sha is now ambiguous.
+
+The extra's URL, for reference:
+
+```
+kernel-fun @ git+https://github.com/allenai/kernel-fun.git@<sha>
+```
+
+Three rules for that extra, from reading OLMo-core's build (2026-09-09):
+
+- **Bare `kernel-fun`, no `[cu13]`/`[cu12]`.** OLMo-core's images already carry the CuTe
+  DSL through the `fa4` extra (flash-attn-4 -> nvidia-cutlass-dsl), and the CUDA major is
+  decided once, in the Dockerfile (`CUDA_VERSION_PATH`, and `FLASH_ATTN_4_EXTRAS='[cu13]'`
+  on the B300 image). One extra then works on both the cu128 and cu130 images; the
+  package's own CUDA extras are for environments that have no DSL at all. Nothing in the
+  bare install replaces the image's torch/triton/fla — the floors are below both images.
+- **The git URL cannot sit in the extra** if `ai2-olmo-core` is to keep publishing: PyPI
+  rejects metadata with a direct-URL requirement (the pyproject says so at `dion`). Put
+  the source in `[tool.uv.sources]` as `dion` does, or publish kernel-fun to PyPI and pin
+  a version.
+- **fla must be 0.5.2 on the branch that merges.** The package requires
+  `fla-core>=0.5.2`; a branch still on flash-linear-attention 0.4.1 (the B300 image
+  branch, today) would end up with a mixed `fla/` tree.
+
+If a cu130 image is ever built with FA4's bare requirement, the package logs a
+`CUDA 12 one (nvidia-cutlass-dsl-libs-cu12 without -cu13)` warning at first use — see
+`support.cute_cuda_mismatch`. It warns rather than falls back, because the wrong-build DSL
+may still compile; treat the line as a build-arg bug.
+
+Both route `flash_linear_attn_api.py::dispatch_chunk_kda` and `dispatch_causal_conv1d` to
+the package when `KimiDeltaAttentionConfig.use_cute_kernel=True`; that one flag drives BOTH
+families (it is plumbed to the three Q/K/V `CausalConv1d`s). There is no other copy of
+these kernels in OLMo-core — the old frozen `nn/attention/kda_cute/` tree went three ideas
+stale before it was deleted, which is why the package exists. Do not reintroduce one.
+
+**To ship a release:** commit/push **kernel-fun** (this repo — tag it if the version
+moved); on OLMo-core's vendored branch re-copy the changed files, bump `VENDORED_FROM` to
+the full kernel-fun sha, bump the scaling-ladders submodule; on the package branch bump
+the sha in the extra and re-lock. Then §2–§3 below. Nothing in this flow touches the ladder
+repo — a ladder commit only ever enters through `tools/vendor.py` (README, "Two repos").
+
+The kda call site, for reference:
+
+```python
+# flash_linear_attn_api.py
+if use_cute_kernel:
+    from kernel_fun.kda import chunk_kda as cute_chunk_kda
+
+    return cute_chunk_kda(
+        q=q, k=k, v=v, g=g, beta=beta, A_log=A_log, dt_bias=dt_bias,
+        scale=scale, initial_state=initial_state,
+        output_final_state=output_final_state,
+        use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+        use_gate_in_kernel=use_gate_in_kernel,
+        cu_seqlens=cu_seqlens,
+    )
+```
+
+The package decides internally and falls back to fla itself, so a packed-document batch or
+a non-Blackwell node needs no branch in the caller (the old `cute_kda_supported` is gone).
+
+**The conv call site**, same file, is the second family. `CausalConv1d.forward`
+(`src/olmo_core/nn/convolution.py`) reaches the package through `dispatch_causal_conv1d`:
+
+```python
+# flash_linear_attn_api.py
+def dispatch_causal_conv1d(x, weight, bias, activation, backend="triton", cu_seqlens=None):
+    from kernel_fun.cconv import causal_conv1d
+
+    return causal_conv1d(
+        x=x, weight=weight, bias=bias, activation=activation, backend=backend,
+        cu_seqlens=cu_seqlens,
+    )
+```
+
+No flag: the package forwards anything off its box (bias, cu_seqlens, `activation=None`,
+`backend="cuda"`) to fla itself, and `KERNEL_FUN_CCONV_DISABLE=1` is the whole-family
+switch. Like the kda entry point, this one IS `torch.compiler.disable`d — as of
+2026-09-04. It was not, on the theory that fla's conv is plain Python around an
+autograd.Function and so is ours, so Dynamo would treat the two alike. The 30M mainline
+ladder falsified that: Dynamo speculated OUR Function's backward (two tensor arguments is
+the shape it agrees to trace; fla's eleven-argument one is not) and asserted inside
+`cconv_bwd` on a symbolic `dy.stride(0)`. So this entry point costs ONE graph break that
+fla does not — expect §3's break diff to show it, at three convs per KDA layer, and judge
+it against the ~36 ms/step the kernels return rather than against zero.
+
+Three things the caller should still do:
+
+```python
+kernel_fun.kda.warmup(K=head_dim, V=head_v_dim, HV=n_v_heads)   # before step 1
+kernel_fun.cconv.warmup(B=microbatch, T=seq_len, D=(key_dim, value_dim))   # ditto
+```
+
+Do NOT call `versions()` from the model's `forward` to log it. Both entry points log it
+themselves, once, the first time either family engages — and they are
+`torch.compiler.disable`d, so the line costs no graph break. The caller-side version of
+this cost two breaks in the KDA layer's compiled block and raised
+`TypeError: unhashable type: 'dict'` through olmo-core's `lru_cache`-based `log_once`
+(2026-09-04). `versions()` stays public for a startup log outside any compiled region.
+
+The cconv autotune key is `(D, W, B, T)`, so its warmup must see the REAL microbatch and
+sequence length and every channel count the layer convolves (q/k at `n_heads*head_dim`,
+v at `expand_v` times that). A shape it has not seen re-autotunes at first use: a few
+seconds, once, per shape — fine for training, but it is what a step-1 "regression" looks
+like.
+
+The old wrapper computed the gate activation in eager torch. Do not reintroduce that: pass
+`use_gate_in_kernel=True` and let the package fuse it into the cumsum. At prod8192 the
+eager form costs ~2 ms and ~2 GiB of saved activations per layer, which is a third of what
+these kernels win.
+
+## 2. Before the first real run
+
+- [ ] `pytest -q` from this repo on a B300, green (`-m slow` adds the prod-shaped rows;
+      the cconv tests also run on an H100).
+- [ ] `PYTHONPATH=src python tools/prodtime.py` on that B300 — the package's own numbers at
+      the production shapes, next to fla's, from the installed copy. (The ladder's
+      `python -m loop bench kda --impl 000,005 --case prod8192_gate` and
+      `... bench cconv --impl 000,001 --case prod8192_qk` are the recorded originals.)
+- [ ] Confirm fla's Triton autotune cache directory is persistent and pre-warmed in the
+      training image. Otherwise every rank re-autotunes at step 1 and the startup cost you
+      just moved with `warmup()` comes back through another door.
+- [ ] Grep the startup log for the `kernel-fun kda: engaged` line (and `cconv: engaged`).
+      If either says `falling back` instead, the reason is on the same line — read it
+      before looking at anything else. The `kernel-fun {...}` versions line appears just
+      above the first of them; if it is missing, neither family ever engaged.
+
+## 3. The A/B that settles it
+
+Three arms, same GPU, same node, first step discarded, enough metric windows to be real
+(the 1.4b gnorm run was killed after 7 windows, and half the "−3%" came from reading the
+average anyway):
+
+1. fla (`use_cute_kernel=False`, conv routed to fla)
+2. kernel-fun, both families
+3. kernel-fun with `KERNEL_FUN_DISABLE=1`
+
+Arm 3 is the control that separates *the wrapper* from *the kernels*: it runs the same
+call path and the same fallback logic but fla's kernels underneath. If arm 3 is not within
+noise of arm 1, the problem is integration, not numerics, and no amount of kernel work will
+fix it. If arm 2 disappoints, `KERNEL_FUN_KDA_DISABLE=1` / `KERNEL_FUN_CCONV_DISABLE=1`
+split it by family without a config change.
+
+A fourth arm exists only for a model whose KDA grid is below the chain-level floor —
+`B * HV * (V//64) < 256`, which the small OLMoE3 candidate's B4/HV8/V256 is, at 128.
+`KERNEL_FUN_KDA_MIN_CTAS=128` moves that floor so those calls dispatch here instead of
+reporting `grid too small` and going to fla. It is a measured opt-in for one shape on one
+box, not a better default: the b1 scan and dhu backwards keep their own 256-CTA floors and
+stay on fla at that grid, so the arm is worth running only against arm 1, on the model you
+intend to train, and only after the log stops saying `grid too small`. Set it for
+`warmup()` too — it warms one grid per floor in play, and the stages this arm adds are
+autotuned in step 1 otherwise.
+
+Then:
+
+- [ ] **`TORCH_LOGS=recompiles,graph_breaks`, diff arm 1 against arm 2.** Free — the Beaker
+      launcher sets it already. At the kda call the counts must be EQUAL: fla's own
+      `chunk_kda` is `@torch.compiler.disable`d too, so that one MOVES a break rather than
+      adding one. At the conv, expect arm 2 to show exactly ONE more break per conv call
+      (three per KDA layer) — see §1; more than that, or a recompile loop, is a bug.
+      Do this first; it is the cheapest way to catch the failure mode that cost two days.
+- [ ] **torch-profile ~20 steps, compare CPU gap time between GPU kernels** across arms.
+      That is where a host-side regression shows up and no benchmark in this repo can see
+      it. While you are in the profile, check that no `softplus` / `mul` / `to_copy` at
+      `[B,T,HV,K]` survives — if they do, the gate is not fused and you are paying twice.
+      And check the conv: `cconv_fwd_strip` / `cconv_bwd_strip` should be the only conv
+      kernels, and `causal_conv1d_fwd_kernel` should appear NOWHERE in the backward (fla
+      re-ran the forward there to rematerialize the pre-activation; ours does not).
+- [ ] **`torch.cuda.memory_summary()` after ~5 steps**, both arms. Active memory should be
+      within a few hundred MiB. The call caches used to pin ~24 GiB/rank at these shapes;
+      that is fixed and tested, and this is the check that keeps it fixed.
+- [ ] Report `throughput/device/TPS (actual avg)` with step 1 discarded.
+
+## 4. What to expect
+
+**kda.** At the production call the op measures **1.540x** (23.74 ms vs fla's 36.56 at
+B=16, T=8192, H=HV=16, K=128, V=256; 2026-09-02). The 1.288x port was worth 1.037x
+end-to-end on the 1.4b mainline ladder, which puts KDA at ~16% of step time; the same
+arithmetic gives **~1.06x** here, i.e. about +6% tokens/sec. Relative to the branch's
+current 1.545x chain the new wy stage is +0.03x on the op — real, small, and the same
+arithmetic says ~+0.3 points of TPS. Do not expect to see it in a noisy A/B.
+
+**cconv.** The 810m/B300 trace (the ladder's `profiles/810m-b300-20260901.md`) put the three conv calls
+at ~50 ms of a ~520 ms step: bwd 35.0 + fwd 7.7 + the backward's own forward re-run 6.9.
+At the measured 4.65x / 1.62x that is ~13 ms — **~36 ms/step back, about 7% of the step,
+~1.075x tokens/sec at 810m**. The conv channel count is set by head count, which the
+ladder pins from 810m up, so the fraction survives scale. This is the larger of the two
+predictions and the easier one to see.
+
+Both predictions assume the ops' step-time fractions are unchanged, and they are the
+numbers to check against the A/B rather than to quote from. Judge the port on its fraction
+of real step time — that is the one durable lesson from last time.