ThunderKittens kernels for hopper

README.md

@@ -15,6 +15,18 @@ This kernel achieves around a 3.7x speedup over an XLA optimized kernel, with li
 uv add flash-hog
 ```
 
+## Optional: ThunderKittens kernels (Hopper)
+Opt-in TK kernels for the double-backward, ~1.65x faster end-to-end on H200
+(causal, head_dim 64, seq % 128 == 0). Off by default; unsupported shapes fall back to Pallas.
+```sh
+uv add 'flash-hog[tk]'   # optional: CUDA build tools from PyPI (needs a host C++ compiler)
+```
+```python
+from flash_hog.jax import _tk_gpu as tk
+tk.enable()   # JIT-builds the plugin on first use (cached); then TK kernels are live
+```
+Set `THUNDERKITTENS_PATH` to use a local ThunderKittens checkout instead of the auto-fetched one.
+
 ## Method
 Flash Hog does 4 recomputation passes to avoid any atomics or saving any intermediary tensors of shape `(N_Q, N_K)`.
 This shakes out to be thread-wise tiling across Q in 3 passes first, once to compute `dd`, then once for `b`, then once for both `dQ'` and `ddO`.

flash_hog/csrc/tk_bwdbwd/BENCHMARKS.md

@@ -0,0 +1,78 @@
+# ThunderKittens double-backward: benchmarks vs the Pallas path
+
+End-to-end timing of the full causal-attention **HVP** (forward + backward +
+double-backward, one jitted graph) through `flash_hog.jax.attention.dot_product_attention`,
+comparing the stock Pallas double-backward against the opt-in ThunderKittens path
+(`flash_hog.jax._tk_gpu.enable()`). Only the double-backward differs between the two
+columns — the cuDNN forward and first backward are identical.
+
+Setup: NVIDIA H200, CUDA 12.8.1, `jax[cuda13]==0.10.1`, 12 heads, head_dim 64, causal,
+bf16 kernels with fp32 inputs/outputs. ThunderKittens pinned at `34b15f7e`. Times are
+means over 10 iterations (5 at seq ≥ 16k, 3 at ≥ 64k) after warmup; expect ~±5%
+machine-to-machine variance.
+
+## B = 1 (latency)
+
+| batch | seq | Pallas (ms) | TK (ms) | speedup | faithfulness |
+|--:|--:|--:|--:|--:|:--|
+| 1 | 512 | 0.18 | 0.22 | 0.82x | cos 1.0000 vs fp32 (Pallas: 1.0000) |
+| 1 | 1024 | 0.33 | 0.28 | 1.17x | cos 1.0000 vs fp32 (Pallas: 1.0000) |
+| 1 | 2048 | 0.88 | 0.66 | 1.34x | cos 1.0000 vs fp32 (Pallas: 1.0000) |
+| 1 | 4096 | 2.72 | 1.76 | 1.55x | cos 1.0000 vs Pallas |
+| 1 | 8192 | 9.46 | 5.56 | 1.70x | cos 1.0000 vs Pallas |
+| 1 | 16384 | 34.94 | 19.21 | 1.82x | cos 1.0000 vs Pallas |
+| 1 | 32768 | 137.54 | 73.63 | 1.87x | cos 1.0000 vs Pallas |
+| 1 | 65536 | 555.30 | 289.03 | 1.92x | cos 1.0000 vs Pallas |
+| 1 | 131072 | 2282.74 | 1143.72 | **2.00x** | cos 1.0000 vs Pallas |
+
+## Constant token budget (B × S = 262,144)
+
+GPU saturated at every row; the throughput view.
+
+| batch | seq | Pallas (ms) | TK (ms) | speedup | faithfulness |
+|--:|--:|--:|--:|--:|:--|
+| 512 | 512 | 27.72 | 18.34 | 1.51x | cos 1.0000 vs Pallas |
+| 256 | 1024 | 44.73 | 27.16 | 1.65x | cos 1.0000 vs Pallas |
+| 128 | 2048 | 78.34 | 44.83 | 1.75x | cos 1.0000 vs Pallas |
+| 64 | 4096 | 146.21 | 80.18 | 1.82x | cos 1.0000 vs Pallas |
+| 32 | 8192 | 281.33 | 150.76 | 1.87x | cos 1.0000 vs Pallas |
+| 16 | 16384 | 557.66 | 291.53 | 1.91x | cos 1.0000 vs Pallas |
+| 8 | 32768 | 1104.21 | 577.58 | 1.91x | cos 1.0000 vs Pallas |
+| 4 | 65536 | 2227.75 | 1137.58 | 1.96x | cos 1.0000 vs Pallas |
+| 2 | 131072 | 4602.62 | 2287.29 | **2.01x** | cos 1.0000 vs Pallas |
+
+The speedup grows with sequence length because the double-backward is the dominant
+O(S²) term of the HVP: as it takes over the runtime, the kernel-level advantage
+(stage1 ~1.9x, stage2 ~1.3x over the Pallas kernels) shows through fully. The only
+regression is tiny single-sequence shapes (1×512), where launch overhead dominates —
+exactly the regime `enable()`'s per-call fallback leaves available to Pallas anyway
+for unsupported shapes.
+
+Faithfulness: where the dense fp32 reference fits in memory (B=1, seq ≤ 2048), both
+paths give cos = 1.0000 against it; the TK path's relative error (~4e-3) is slightly
+tighter than Pallas (~5e-3). At larger shapes the two paths agree with each other to
+cos = 1.0000 (~2e-3 rel).
+
+## Reproducing
+
+```python
+import jax, jax.numpy as jnp
+import flash_hog.jax.attention as fa
+from flash_hog.jax import _tk_gpu as tk
+
+def attn(q, k, v, scale):
+    qb, kb, vb = (x.astype(jnp.bfloat16) for x in (q, k, v))
+    return fa.dot_product_attention(qb, kb, vb, is_causal=True, scale=scale).astype(jnp.float32)
+
+def tree_dot(a, b):
+    return sum(jnp.vdot(x, y) for x, y in zip(jax.tree.leaves(a), jax.tree.leaves(b)))
+
+def make_hvp(cot, tan, scale):           # grad of <grad(loss), tan> == HVP
+    def loss(x):
+        return jnp.vdot(attn(*x, scale), cot)
+    return jax.jit(lambda x: jax.grad(lambda y: tree_dot(jax.grad(loss)(y), tan))(x))
+
+# time make_hvp(...)(qkv) with tk.disable() vs tk.enable()
+```
+
+Install `flash-hog[tk]`; the plugin JIT-builds (cached) on first `tk.enable()`.

flash_hog/csrc/tk_bwdbwd/ffi.cu

@@ -0,0 +1,98 @@
+// XLA FFI plugin: ThunderKittens double-backward (stage1 + stage2) for causal attention.
+//
+// "TkBwdBwd": launches tk_stage1 then tk_stage2 on the XLA stream (stage2
+// reads the dD/B vectors stage1 writes; same-stream ordering suffices).
+//
+//   inputs : Q, K, V, dO, ddQ, ddK, ddV   bf16 (B, H, T, hd)   [BHTD, contiguous]
+//            L, D                         f32  (B, H, T)
+//   attr   : scale (f32)
+//   outputs: dQ2, ddO, dK2, dV2           bf16 (B, H, T, hd)
+//            dD, B                        f32  (B, H, T)       [stage1->stage2 scratch]
+//
+// Constraints: SM90 (Hopper), head_dim == 64, T % 128 == 0, causal, q_heads == kv_heads.
+// Built at runtime by flash_hog/jax/_tk_build.py (pip CUDA tools, cached).
+
+#include "stage1.cuh"
+#include "stage2.cuh"
+
+#include "xla/ffi/api/ffi.h"
+namespace ffi = xla::ffi;
+
+namespace s1 = flash_hog_tk::stage1;
+namespace s2 = flash_hog_tk::stage2;
+
+static constexpr int HEAD_DIM = 64;
+
+static ffi::Error TkBwdBwdImpl(
+    cudaStream_t stream,
+    ffi::AnyBuffer q, ffi::AnyBuffer k, ffi::AnyBuffer v, ffi::AnyBuffer dO,
+    ffi::AnyBuffer ddq, ffi::AnyBuffer ddk, ffi::AnyBuffer ddv,
+    ffi::AnyBuffer l, ffi::AnyBuffer d,
+    float scale,
+    ffi::Result<ffi::AnyBuffer> dq2, ffi::Result<ffi::AnyBuffer> ddo,
+    ffi::Result<ffi::AnyBuffer> dk2, ffi::Result<ffi::AnyBuffer> dv2,
+    ffi::Result<ffi::AnyBuffer> dd, ffi::Result<ffi::AnyBuffer> b) {
+    auto dims = q.dimensions();
+    if (dims.size() != 4) return ffi::Error::InvalidArgument("Q must be (B,H,T,hd)");
+    const unsigned B = dims[0], H = dims[1], T = dims[2], hd = dims[3];
+    if (hd != HEAD_DIM) return ffi::Error::InvalidArgument("head_dim must be 64");
+    if (T % 128 != 0) return ffi::Error::InvalidArgument("T must be divisible by 128");
+
+    static bool attrs_set = false;
+    if (!attrs_set) {
+        cudaFuncSetAttribute(s1::tk_stage1, cudaFuncAttributeMaxDynamicSharedMemorySize, s1::SMEM_BYTES);
+        cudaFuncSetAttribute(s2::tk_stage2, cudaFuncAttributeMaxDynamicSharedMemorySize, s2::SMEM_BYTES);
+        attrs_set = true;
+    }
+
+    using kittens::bf16;
+    auto bfp = [](ffi::AnyBuffer& x) { return reinterpret_cast<bf16*>(x.untyped_data()); };
+    auto bfr = [](ffi::Result<ffi::AnyBuffer>& x) { return reinterpret_cast<bf16*>(x->untyped_data()); };
+    auto flp = [](ffi::AnyBuffer& x) { return reinterpret_cast<float*>(x.untyped_data()); };
+    auto flr = [](ffi::Result<ffi::AnyBuffer>& x) { return reinterpret_cast<float*>(x->untyped_data()); };
+
+    s1::tk_globals G1{
+        s1::tk_qgl{bfp(q), B, H, T, HEAD_DIM},   s1::tk_qgl{bfp(dO), B, H, T, HEAD_DIM},
+        s1::tk_qgl{bfp(ddq), B, H, T, HEAD_DIM}, s1::tk_qgl{bfr(dq2), B, H, T, HEAD_DIM},
+        s1::tk_qgl{bfr(ddo), B, H, T, HEAD_DIM},
+        s1::tk_kgl{bfp(k), B, H, T, HEAD_DIM},   s1::tk_kgl{bfp(v), B, H, T, HEAD_DIM},
+        s1::tk_kgl{bfp(ddk), B, H, T, HEAD_DIM}, s1::tk_kgl{bfp(ddv), B, H, T, HEAD_DIM},
+        flp(l), flp(d), flr(dd), flr(b), (int)T, scale};
+    s2::tk_globals G2{
+        s2::tk_qgl{bfp(q), B, H, T, HEAD_DIM},   s2::tk_qgl{bfp(dO), B, H, T, HEAD_DIM},
+        s2::tk_qgl{bfp(ddq), B, H, T, HEAD_DIM},
+        s2::tk_kgl{bfp(k), B, H, T, HEAD_DIM},   s2::tk_kgl{bfp(v), B, H, T, HEAD_DIM},
+        s2::tk_kgl{bfp(ddk), B, H, T, HEAD_DIM}, s2::tk_kgl{bfp(ddv), B, H, T, HEAD_DIM},
+        s2::tk_kgl{bfr(dk2), B, H, T, HEAD_DIM}, s2::tk_kgl{bfr(dv2), B, H, T, HEAD_DIM},
+        s2::tk_vgl{flp(l), B, H, 1, T},  s2::tk_vgl{flp(d), B, H, 1, T},
+        s2::tk_vgl{flr(dd), B, H, 1, T}, s2::tk_vgl{flr(b), B, H, 1, T},
+        (int)T, scale};
+
+    dim3 grid(T / 128, H, B);
+    s1::tk_stage1<<<grid, s1::TK_WORKERS * 32, s1::SMEM_BYTES, stream>>>(G1);
+    s2::tk_stage2<<<grid, s2::TK_WORKERS * 32, s2::SMEM_BYTES, stream>>>(G2);
+    if (cudaError_t e = cudaGetLastError(); e != cudaSuccess)
+        return ffi::Error::Internal(cudaGetErrorString(e));
+    return ffi::Error::Success();
+}
+
+XLA_FFI_DEFINE_HANDLER_SYMBOL(
+    TkBwdBwd, TkBwdBwdImpl,
+    ffi::Ffi::Bind()
+        .Ctx<ffi::PlatformStream<cudaStream_t>>()
+        .Arg<ffi::AnyBuffer>()   // Q
+        .Arg<ffi::AnyBuffer>()   // K
+        .Arg<ffi::AnyBuffer>()   // V
+        .Arg<ffi::AnyBuffer>()   // dO
+        .Arg<ffi::AnyBuffer>()   // ddQ
+        .Arg<ffi::AnyBuffer>()   // ddK
+        .Arg<ffi::AnyBuffer>()   // ddV
+        .Arg<ffi::AnyBuffer>()   // L
+        .Arg<ffi::AnyBuffer>()   // D
+        .Attr<float>("scale")
+        .Ret<ffi::AnyBuffer>()   // dQ2
+        .Ret<ffi::AnyBuffer>()   // ddO
+        .Ret<ffi::AnyBuffer>()   // dK2
+        .Ret<ffi::AnyBuffer>()   // dV2
+        .Ret<ffi::AnyBuffer>()   // dD
+        .Ret<ffi::AnyBuffer>());  // B