diff --git a/README.md b/README.md index 4004d37..b3baf73 100644 --- a/README.md +++ b/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`. diff --git a/flash_hog/csrc/tk_bwdbwd/BENCHMARKS.md b/flash_hog/csrc/tk_bwdbwd/BENCHMARKS.md new file mode 100644 index 0000000..e7db558 --- /dev/null +++ b/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 == 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()`. diff --git a/flash_hog/csrc/tk_bwdbwd/ffi.cu b/flash_hog/csrc/tk_bwdbwd/ffi.cu new file mode 100644 index 0000000..97abf96 --- /dev/null +++ b/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 dq2, ffi::Result ddo, + ffi::Result dk2, ffi::Result dv2, + ffi::Result dd, ffi::Result 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(x.untyped_data()); }; + auto bfr = [](ffi::Result& x) { return reinterpret_cast(x->untyped_data()); }; + auto flp = [](ffi::AnyBuffer& x) { return reinterpret_cast(x.untyped_data()); }; + auto flr = [](ffi::Result& x) { return reinterpret_cast(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<<>>(G1); + s2::tk_stage2<<>>(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>() + .Arg() // Q + .Arg() // K + .Arg() // V + .Arg() // dO + .Arg() // ddQ + .Arg() // ddK + .Arg() // ddV + .Arg() // L + .Arg() // D + .Attr("scale") + .Ret() // dQ2 + .Ret() // ddO + .Ret() // dK2 + .Ret() // dV2 + .Ret() // dD + .Ret()); // B \ No newline at end of file diff --git a/flash_hog/csrc/tk_bwdbwd/stage1.cuh b/flash_hog/csrc/tk_bwdbwd/stage1.cuh new file mode 100644 index 0000000..5f92b15 --- /dev/null +++ b/flash_hog/csrc/tk_bwdbwd/stage1.cuh @@ -0,0 +1,260 @@ +// ThunderKittens kernel: stage 1 of the causal-attention double-backward (SM90 only). + +#pragma once + +#include "kittens.cuh" + +namespace flash_hog_tk { +namespace stage1 { + +constexpr int HEAD_DIM = 64; + +template struct bwd_over_bwd_attend_ker_tile_dims {}; +template<> struct bwd_over_bwd_attend_ker_tile_dims<64> { + constexpr static int tile_width = (64); // head_dim + constexpr static int qo_height = (4*16); // wgmma M = 64 (4 warps x 16 rows) + constexpr static int kv_height = (2*16); // 32: 4 streamed tensors + width-32 fp32 + // register tiles fit the 224-reg budget + constexpr static int stages = (4); // ring: 4 stages x 4 tensors x 4KB = 64KB +}; + +template struct bwd_over_bwd_globals { + using dims = bwd_over_bwd_attend_ker_tile_dims; + using q_tile = kittens::st_bf; + using k_tile = kittens::st_bf; + using per_example_vec = kittens::sv_fl; + using q_gl = kittens::gl; + using k_gl = kittens::gl; +}; + +constexpr int CONSUMER_WARPGROUPS = (2); // block = 128 queries (2 WGs x 64 rows) +constexpr int PRODUCER_WARPGROUPS = (1); +constexpr int NUM_WARPGROUPS = (CONSUMER_WARPGROUPS + PRODUCER_WARPGROUPS); +constexpr int NUM_WORKERS = (NUM_WARPGROUPS * kittens::WARPGROUP_WARPS); + +constexpr int TK_TQ = bwd_over_bwd_attend_ker_tile_dims::qo_height; +constexpr int TK_TK = bwd_over_bwd_attend_ker_tile_dims::kv_height; +constexpr int KV_STAGES = bwd_over_bwd_attend_ker_tile_dims::stages; +constexpr int TK_CWG = CONSUMER_WARPGROUPS; +constexpr int TK_WORKERS = NUM_WORKERS; +constexpr int SMEM_BYTES = 160 * 1024; + +using tk_qtile = bwd_over_bwd_globals::q_tile; +using tk_ktile = bwd_over_bwd_globals::k_tile; +using tk_qgl = bwd_over_bwd_globals::q_gl; +using tk_kgl = bwd_over_bwd_globals::k_gl; + +struct tk_globals { // field order is load-bearing (aggregate init) + tk_qgl Q, dO, ddQ, dQ2, ddO; // query-side (64-row tiles) + tk_kgl K, V, ddK, ddV; // key-side (32-row tiles) + const float* Lp; // (B,H,T) logsumexp (natural log) + const float* Dp; // (B,H,T) D = rowsum(dO*O) + float* dDo; // (B,H,T) out: dD (stage2 input) + float* Bo; // (B,H,T) out: B (stage2 input) + int N; + float scale; +}; + +// helper for causal mask +__device__ static inline void tk_cmask(kittens::rt_fl<16, TK_TK>& S, int q16, int kt) { + #pragma unroll + for (int j = 0; j < TK_TK / 16; j++) { + int k16 = kt * (TK_TK / 16) + j; + auto& sub = reinterpret_cast&>(S.tiles[0][j]); + if (k16 > q16) kittens::warp::neg_infty(sub); + else if (k16 == q16) + kittens::warp::make_causal(sub, sub, kittens::base_types::constants::neg_infty()); + } +} + +__global__ __launch_bounds__(TK_WORKERS * 32, 1) +void tk_stage1(const __grid_constant__ tk_globals g) { + using namespace kittens; + extern __shared__ int __shm[]; + tma_swizzle_allocator al((int*)&__shm[0]); + const int wid = warpid(), wgid = wid / 4; // wg 0,1 = consumers; wg 2 = producer + + using G = bwd_over_bwd_globals; + using q_tile = G::q_tile; + using k_tile = G::k_tile; + using per_example_vec = G::per_example_vec; + + const int qblk = blockIdx.x, head = blockIdx.y, batch = blockIdx.z; + constexpr float LOG2E = 1.44269504089f; + + q_tile (&q_s)[CONSUMER_WARPGROUPS] = al.allocate(); + q_tile (&do_s)[CONSUMER_WARPGROUPS] = al.allocate(); + q_tile (&dq_s)[CONSUMER_WARPGROUPS] = al.allocate(); // ddQ + k_tile (&k_s)[KV_STAGES] = al.allocate(); + k_tile (&v_s)[KV_STAGES] = al.allocate(); + k_tile (&dk_s)[KV_STAGES] = al.allocate(); // ddK + k_tile (&dv_s)[KV_STAGES] = al.allocate(); // ddV + per_example_vec (&l_s)[CONSUMER_WARPGROUPS] = al.allocate(); + per_example_vec (&d_s)[CONSUMER_WARPGROUPS] = al.allocate(); + + const int k_tiles = 4 * (qblk + 1); // causal: key tiles this block needs + const int total_k_tiles = 2 * k_tiles; // the key sequence is streamed twice + + __shared__ kittens::semaphore q_semaphore, k_arrived[KV_STAGES], compute_done[KV_STAGES]; + + if (threadIdx.x == 0) { + init_semaphore(q_semaphore, 0, 1); + for (int s = 0; s < KV_STAGES; s++) { + init_semaphore(k_arrived[s], 0, 1); + init_semaphore(compute_done[s], TK_CWG, 0); + } + + tma::expect_bytes(q_semaphore, 3 * CONSUMER_WARPGROUPS * sizeof(q_tile)); + for (int i = 0; i < CONSUMER_WARPGROUPS; i++) { + coord qidx = {batch, head, qblk * CONSUMER_WARPGROUPS + i, 0}; + tma::load_async(q_s[i], g.Q, qidx, q_semaphore); + tma::load_async(do_s[i], g.dO, qidx, q_semaphore); + tma::load_async(dq_s[i], g.ddQ, qidx, q_semaphore); + } + + for (int k_iter = 0; k_iter < KV_STAGES - 1 && k_iter < total_k_tiles; k_iter++) { + coord kv_idx = {batch, head, k_iter, 0}; + tma::expect_bytes(k_arrived[k_iter], 4 * sizeof(k_tile)); + tma::load_async(k_s[k_iter], g.K, kv_idx, k_arrived[k_iter]); + tma::load_async(v_s[k_iter], g.V, kv_idx, k_arrived[k_iter]); + tma::load_async(dk_s[k_iter], g.ddK, kv_idx, k_arrived[k_iter]); + tma::load_async(dv_s[k_iter], g.ddV, kv_idx, k_arrived[k_iter]); + } + } + + const int vbase = ((batch * gridDim.y + head) * g.N) + qblk * CONSUMER_WARPGROUPS * TK_TQ; + for (int i = threadIdx.x; i < TK_CWG * TK_TQ; i += TK_WORKERS * 32) { + l_s[i / TK_TQ][i % TK_TQ] = g.Lp[vbase + i] * LOG2E; + d_s[i / TK_TQ][i % TK_TQ] = g.Dp[vbase + i]; + } + + __syncthreads(); + + if (wgid == NUM_WARPGROUPS - 1) { + // producer wg + warpgroup::decrease_registers<32>(); + if (wid == CONSUMER_WARPGROUPS * 4) { + for (int kv_idx = KV_STAGES - 2; kv_idx <= total_k_tiles - 2; kv_idx++) { + int it = kv_idx + 1, s = it % KV_STAGES; + int kt = (it < k_tiles) ? it : it - k_tiles; + coord kv_tile_idx = {batch, head, kt, 0}; + warp::tma::expect_bytes(k_arrived[s], 4 * sizeof(k_tile)); + warp::tma::load_async(k_s[s], g.K, kv_tile_idx, k_arrived[s]); + warp::tma::load_async(v_s[s], g.V, kv_tile_idx, k_arrived[s]); + warp::tma::load_async(dk_s[s], g.ddK, kv_tile_idx, k_arrived[s]); + warp::tma::load_async(dv_s[s], g.ddV, kv_tile_idx, k_arrived[s]); + wait(compute_done[kv_idx % KV_STAGES], (kv_idx / KV_STAGES) % 2); + } + } + } else { + // consumer wgs + warpgroup::increase_registers<224>(); + + wait(q_semaphore, 0); + + col_vec> lv, dvv, dDv, r1, r2, r3, Bv, tv; + warp::zero(dDv); warp::zero(r1); warp::zero(r2); warp::zero(r3); + rt_fl<16, HEAD_DIM> accQ, accO; + warp::zero(accQ); warp::zero(accO); + const float scale2 = g.scale * LOG2E; + + warpgroup::load(lv, l_s[wgid]); + warpgroup::load(dvv, d_s[wgid]); + + const int q16 = qblk * (TK_CWG * (TK_TQ/16)) + wgid * (TK_TQ/16) + (wid % 4); + for (int kv_idx = 0; kv_idx < total_k_tiles; kv_idx++) { + const int kt = kv_idx % k_tiles; // position in the key sequence + const int ring_idx = kv_idx % KV_STAGES; + + wait(k_arrived[ring_idx], (kv_idx / KV_STAGES) % 2); + + rt_fl<16, TK_TK> S, ddS, dP, dPa, intermediate; + + warpgroup::mm_ABt(S, q_s[wgid], k_s[ring_idx]); // S = Q K^T + warpgroup::mm_ABt(ddS, dq_s[wgid], k_s[ring_idx]); // ddS = ddQ K^T ... + warpgroup::mma_ABt(ddS, q_s[wgid], dk_s[ring_idx]); // ... + Q ddK^T + warpgroup::mm_ABt(dP, do_s[wgid], v_s[ring_idx]); // dP = dO V^T + warpgroup::mm_ABt(dPa, do_s[wgid], dv_s[ring_idx]); // dPa = dO ddV^T + warpgroup::mma_commit_group(); + warpgroup::mma_async_wait(); + + warp::mul(S, S, scale2); + warp::mul(ddS, ddS, g.scale); + tk_cmask(S, q16, kt); + warp::sub_row(S, S, lv); + warp::exp2(S, S); // S -> P vioa + + if (kv_idx < k_tiles) { + warp::mul(intermediate, ddS, S); + warp::row_sum(dDv, intermediate, dDv); // dD += sum ddS*P + warp::mul(intermediate, intermediate, dP); + warp::row_sum(r3, intermediate, r3); // r3 += sum dP*ddS*P + warp::mul(intermediate, dPa, S); + warp::row_sum(r1, intermediate, r1); // r1 += sum (dO.ddV)*P + warp::mul(intermediate, dP, S); + warp::row_sum(r2, intermediate, r2); // r2 += sum dP*P + } else { + // second sweep + warp::mul_row(intermediate, dP, dDv); + warp::sub(dPa, dPa, intermediate); + warp::mul_row(intermediate, ddS, dvv); + warp::sub(dPa, dPa, intermediate); + warp::mul(intermediate, dP, ddS); + warp::add(dPa, dPa, intermediate); + + warp::sub_row(dP, dP, dvv); // dS (into dP) + warp::mul(dP, dP, S); + warp::mul(dP, dP, g.scale); + warp::sub_row(dPa, dPa, Bv); // dS2 (into dPa) + warp::mul(dPa, dPa, S); + warp::mul(dPa, dPa, g.scale); + warp::sub_row(ddS, ddS, dDv); // ddP (into ddS) + warp::mul(ddS, ddS, S); + + rt_bf<16, TK_TK> mb0, mb1, mb2, mb3; + warp::copy(mb0, dP); + warp::copy(mb1, dPa); + warp::copy(mb2, ddS); + warp::copy(mb3, S); + warpgroup::mma_AB(accQ, mb0, dk_s[ring_idx]); // dQ2 += dS @ ddK + warpgroup::mma_AB(accQ, mb1, k_s[ring_idx]); // dQ2 += dS2 @ K + warpgroup::mma_AB(accO, mb2, v_s[ring_idx]); // ddO += ddP @ V + warpgroup::mma_AB(accO, mb3, dv_s[ring_idx]); // ddO += P @ ddV + warpgroup::mma_commit_group(); + warpgroup::mma_async_wait(); + } + + if (warpgroup::laneid() == 0) { + arrive(compute_done[ring_idx], 1); + } + + if (kv_idx == k_tiles - 1) { + warp::mul(tv, dDv, r2); + warp::sub(Bv, r1, tv); + warp::mul(tv, dvv, dDv); + warp::sub(Bv, Bv, tv); + warp::add(Bv, Bv, r3); + } + } + + // write out + warpgroup::store(q_s[wgid], accQ); + warpgroup::store(do_s[wgid], accO); + warpgroup::store(l_s[wgid], dDv); + warpgroup::store(d_s[wgid], Bv); + group<4>::sync(wgid + 4); + if (wid % 4 == 0) { + coord idx = {batch, head, qblk * CONSUMER_WARPGROUPS + wgid, 0}; + warp::tma::store_async(g.dQ2, q_s[wgid], idx); + warp::tma::store_async(g.ddO, do_s[wgid], idx); + } + for (int i = warpgroup::laneid(); i < TK_TQ; i += 128) { + g.dDo[vbase + wgid * TK_TQ + i] = l_s[wgid][i]; + g.Bo[vbase + wgid * TK_TQ + i] = d_s[wgid][i]; + } + warp::tma::store_async_wait(); + } +} + +} // namespace stage1 +} // namespace flash_hog_tk \ No newline at end of file diff --git a/flash_hog/csrc/tk_bwdbwd/stage2.cuh b/flash_hog/csrc/tk_bwdbwd/stage2.cuh new file mode 100644 index 0000000..0fa7b98 --- /dev/null +++ b/flash_hog/csrc/tk_bwdbwd/stage2.cuh @@ -0,0 +1,238 @@ +// ThunderKittens kernel: stage 2 of the causal-attention double-backward (SM90 only). + +#pragma once + +#include "kittens.cuh" + +namespace flash_hog_tk { +namespace stage2 { + +constexpr int HEAD_DIM = 64; + +template struct bwd_over_bwd_stage2_tile_dims {}; +template<> struct bwd_over_bwd_stage2_tile_dims<64> { + constexpr static int tile_width = (64); // head_dim + constexpr static int kv_height = (4*16); // keys per consumer wg + constexpr static int qo_height = (2*16); // 32 streamed queries per ring stage + constexpr static int stages = (4); +}; + +template struct bwd_over_bwd_stage2_globals { + using dims = bwd_over_bwd_stage2_tile_dims; + using k_tile = kittens::st_bf; + using q_tile = kittens::st_bf; // stream over q tiles + using per_query_vec = kittens::sv_fl; + using k_gl = kittens::gl; + using q_gl = kittens::gl; + using v_gl = kittens::gl; // L/D/dD/B +}; + +constexpr int CONSUMER_WARPGROUPS = (2); +constexpr int PRODUCER_WARPGROUPS = (1); +constexpr int NUM_WARPGROUPS = (CONSUMER_WARPGROUPS + PRODUCER_WARPGROUPS); +constexpr int NUM_WORKERS = (NUM_WARPGROUPS * kittens::WARPGROUP_WARPS); + +constexpr int TK_TKEY = bwd_over_bwd_stage2_tile_dims::kv_height; +constexpr int TK_TQ = bwd_over_bwd_stage2_tile_dims::qo_height; +constexpr int Q_STAGES = bwd_over_bwd_stage2_tile_dims::stages; +constexpr int TK_CWG = CONSUMER_WARPGROUPS; +constexpr int TK_WORKERS = NUM_WORKERS; +constexpr int SMEM_BYTES = 160 * 1024; + +using tk_ktile = bwd_over_bwd_stage2_globals::k_tile; +using tk_qtile = bwd_over_bwd_stage2_globals::q_tile; +using tk_kgl = bwd_over_bwd_stage2_globals::k_gl; +using tk_qgl = bwd_over_bwd_stage2_globals::q_gl; +using tk_vec = bwd_over_bwd_stage2_globals::per_query_vec; +using tk_vgl = bwd_over_bwd_stage2_globals::v_gl; + +struct tk_globals { // field order is load-bearing (aggregate init) + tk_qgl Q, dO, ddQ; // query-side (streamed, 32-row tiles) + tk_kgl K, V, ddK, ddV; // key-side (resident, 64-row tiles) + tk_kgl dK2, dV2; // outputs (64-row tiles) + tk_vgl L, D, dD, Bv; // per-query scalars, gl dims (B,H,1,T) + int N; + float scale; +}; + + +__device__ static inline void tk_cmask_t(kittens::rt_fl<16, TK_TQ>& S, int k16, int qt) { + #pragma unroll + for (int j = 0; j < TK_TQ / 16; j++) { + int q16 = qt * (TK_TQ / 16) + j; + auto& sub = reinterpret_cast&>(S.tiles[0][j]); + if (q16 < k16) kittens::warp::neg_infty(sub); + else if (q16 == k16) + kittens::warp::make_causal_t(sub, sub, kittens::base_types::constants::neg_infty()); + } +} + +__global__ __launch_bounds__(TK_WORKERS * 32, 1) +void tk_stage2(const __grid_constant__ tk_globals g) { + using namespace kittens; + extern __shared__ int __shm[]; + tma_swizzle_allocator al((int*)&__shm[0]); + const int wid = warpid(), wgid = wid / 4; // wg 0,1 = consumers; wg 2 = producer + const int kblk = blockIdx.x, head = blockIdx.y, batch = blockIdx.z; + constexpr float LOG2E = 1.44269504089f; + + using G = bwd_over_bwd_stage2_globals; + using q_tile = G::q_tile; + using k_tile = G::k_tile; + using per_query_vec = G::per_query_vec; + + k_tile (&k_s)[CONSUMER_WARPGROUPS] = al.allocate(); + k_tile (&v_s)[CONSUMER_WARPGROUPS] = al.allocate(); + k_tile (&dk_s)[CONSUMER_WARPGROUPS] = al.allocate(); // ddK + k_tile (&dv_s)[CONSUMER_WARPGROUPS] = al.allocate(); // ddV + q_tile (&q_s)[Q_STAGES] = al.allocate(); + q_tile (&dq_s)[Q_STAGES] = al.allocate(); // ddQ + q_tile (&do_s)[Q_STAGES] = al.allocate(); + per_query_vec (&l_s)[Q_STAGES] = al.allocate(); + per_query_vec (&d_s)[Q_STAGES] = al.allocate(); + per_query_vec (&dd_s)[Q_STAGES] = al.allocate(); + per_query_vec (&b_s)[Q_STAGES] = al.allocate(); + + const int q_start = (TK_CWG * TK_TKEY / TK_TQ) * kblk; // first q tile (causal) + const int TOTAL = g.N / TK_TQ - q_start; // q tiles to stream + + __shared__ kittens::semaphore k_semaphore, q_arrived[Q_STAGES], compute_done[Q_STAGES]; + + if (threadIdx.x == 0) { + init_semaphore(k_semaphore, 0, 1); + for (int s = 0; s < Q_STAGES; s++) { + init_semaphore(q_arrived[s], 0, 1); + init_semaphore(compute_done[s], TK_CWG, 0); + } + + tma::expect_bytes(k_semaphore, 4 * CONSUMER_WARPGROUPS * sizeof(k_tile)); + for (int w = 0; w < CONSUMER_WARPGROUPS; w++) { + coord kidx = {batch, head, kblk * CONSUMER_WARPGROUPS + w, 0}; + tma::load_async(k_s[w], g.K, kidx, k_semaphore); + tma::load_async(v_s[w], g.V, kidx, k_semaphore); + tma::load_async(dk_s[w], g.ddK, kidx, k_semaphore); + tma::load_async(dv_s[w], g.ddV, kidx, k_semaphore); + } + + for (int j = 0; j < Q_STAGES - 1 && j < TOTAL; j++) { + tma::expect_bytes(q_arrived[j], 3 * sizeof(q_tile) + 4 * sizeof(per_query_vec)); + int qt = q_start + j; + coord qidx = {batch, head, qt, 0}; + coord vidx = {batch, head, 0, qt}; + tma::load_async(q_s[j], g.Q, qidx, q_arrived[j]); + tma::load_async(do_s[j], g.dO, qidx, q_arrived[j]); + tma::load_async(dq_s[j], g.ddQ, qidx, q_arrived[j]); + tma::load_async(l_s[j], g.L, vidx, q_arrived[j]); + tma::load_async(d_s[j], g.D, vidx, q_arrived[j]); + tma::load_async(dd_s[j], g.dD, vidx, q_arrived[j]); + tma::load_async(b_s[j], g.Bv, vidx, q_arrived[j]); + } + } + + __syncthreads(); + + if (wgid == NUM_WARPGROUPS - 1) { + // producer wg + warpgroup::decrease_registers<32>(); + if (wid == CONSUMER_WARPGROUPS * 4) { // one warp issues the TMA loads + for (int j = Q_STAGES - 2; j <= TOTAL - 2; j++) { + int it = j + 1, s = it % Q_STAGES; + int qt = it + q_start; + coord qidx = {batch, head, qt, 0}; + coord vidx = {batch, head, 0, qt}; + warp::tma::expect_bytes(q_arrived[s], 3 * sizeof(q_tile) + 4 * sizeof(per_query_vec)); + warp::tma::load_async(q_s[s], g.Q, qidx, q_arrived[s]); + warp::tma::load_async(do_s[s], g.dO, qidx, q_arrived[s]); + warp::tma::load_async(dq_s[s], g.ddQ, qidx, q_arrived[s]); + warp::tma::load_async(l_s[s], g.L, vidx, q_arrived[s]); + warp::tma::load_async(d_s[s], g.D, vidx, q_arrived[s]); + warp::tma::load_async(dd_s[s], g.dD, vidx, q_arrived[s]); + warp::tma::load_async(b_s[s], g.Bv, vidx, q_arrived[s]); + wait(compute_done[(it - 1) % Q_STAGES], ((it - 1) / Q_STAGES) % 2); + } + } + } else { + // consumer wgs + warpgroup::increase_registers<224>(); + const int k16 = kblk * (TK_CWG * 4) + wgid * 4 + (wid % 4); // this warp's key rows + const float scale2 = g.scale * LOG2E; + + rt_fl<16, HEAD_DIM> accK, accV; + warp::zero(accK); + warp::zero(accV); + wait(k_semaphore, 0); + + for (int q_idx = 0; q_idx < TOTAL; q_idx++) { + const int ring_idx = q_idx % Q_STAGES; + + wait(q_arrived[ring_idx], (q_idx / Q_STAGES) % 2); + + rt_fl<16, TK_TQ> S, ddS, dP, dPa, intermediate; + + warpgroup::mm_ABt(S, k_s[wgid], q_s[ring_idx]); // S^T = K Q^T + warpgroup::mm_ABt(ddS, k_s[wgid], dq_s[ring_idx]); // ddS^T = K ddQ^T ... + warpgroup::mma_ABt(ddS, dk_s[wgid], q_s[ring_idx]); // ... + ddK Q^T + warpgroup::mm_ABt(dP, v_s[wgid], do_s[ring_idx]); // dP^T = V dO^T + warpgroup::mm_ABt(dPa, dv_s[wgid], do_s[ring_idx]); // dPa^T = ddV dO^T + warpgroup::mma_commit_group(); + warpgroup::mma_async_wait(); + + // per-query vectors broadcast along COLUMNS + row_vec> lrv, drv, ddrv, brv; + warp::load(lrv, l_s[ring_idx]); + warp::load(drv, d_s[ring_idx]); + warp::load(ddrv, dd_s[ring_idx]); + warp::load(brv, b_s[ring_idx]); + warp::mul(lrv, lrv, LOG2E); // exp2 path + + warp::mul(S, S, scale2); + warp::mul(ddS, ddS, g.scale); + tk_cmask_t(S, k16, q_idx + q_start); + warp::sub_col(S, S, lrv); + warp::exp2(S, S); // S := P + + // dP2 (into dPa) = dPa - dP*dD - ddS*D + dP*ddS + warp::mul_col(intermediate, dP, ddrv); + warp::sub(dPa, dPa, intermediate); + warp::mul_col(intermediate, ddS, drv); + warp::sub(dPa, dPa, intermediate); + warp::mul(intermediate, dP, ddS); + warp::add(dPa, dPa, intermediate); + + warp::sub_col(dP, dP, drv); // dS (into dP) + warp::mul(dP, dP, S); + warp::mul(dP, dP, g.scale); + warp::sub_col(dPa, dPa, brv); // dS2 (into dPa) + warp::mul(dPa, dPa, S); + warp::mul(dPa, dPa, g.scale); + warp::sub_col(ddS, ddS, ddrv); // ddP (into ddS) + warp::mul(ddS, ddS, S); + + rt_bf<16, TK_TQ> mb0, mb1, mb2; // own temp per async mma + warp::copy(mb0, dP); + warp::copy(mb1, dPa); + warp::copy(mb2, ddS); + warpgroup::mma_AB(accK, mb0, dq_s[ring_idx]); // dK2 += dS @ ddQ + warpgroup::mma_AB(accK, mb1, q_s[ring_idx]); // dK2 += dS2 @ Q + warpgroup::mma_AB(accV, mb2, do_s[ring_idx]); // dV2 += ddP @ dO + warpgroup::mma_commit_group(); + warpgroup::mma_async_wait(); + + if (warpgroup::laneid() == 0) arrive(compute_done[ring_idx], 1); + } + + // write out + warpgroup::store(k_s[wgid], accK); + warpgroup::store(v_s[wgid], accV); + group<4>::sync(wgid + 4); + if (wid % 4 == 0) { + coord idx = {batch, head, kblk * TK_CWG + wgid, 0}; + warp::tma::store_async(g.dK2, k_s[wgid], idx); + warp::tma::store_async(g.dV2, v_s[wgid], idx); + } + warp::tma::store_async_wait(); + } +} + +} // namespace stage2 +} // namespace flash_hog_tk \ No newline at end of file diff --git a/flash_hog/jax/_tk_build.py b/flash_hog/jax/_tk_build.py new file mode 100644 index 0000000..119d5cf --- /dev/null +++ b/flash_hog/jax/_tk_build.py @@ -0,0 +1,159 @@ +"""JIT build of the ThunderKittens double-backward plugin. + +Compiles csrc/tk_bwdbwd/ffi.cu with pip-provided CUDA tools (`flash-hog[tk]`); +a host C++ compiler is the only non-pip requirement. ThunderKittens (header-only, +pinned commit) is fetched once as a tarball, or set THUNDERKITTENS_PATH. +Outputs are cached under ~/.cache/flash_hog/ keyed on sources + TK commit + arch. +The plugin links with `-cudart none`; driver/cudart symbols resolve at load time +(_tk_gpu preloads them RTLD_GLOBAL). +""" + +from __future__ import annotations + +import hashlib +import os +import shutil +import subprocess +import sys +import tarfile +import tempfile +import urllib.request +from pathlib import Path + +TK_COMMIT = "34b15f7e7012de25ae162c8d9dc85296dd342676" +_TK_TARBALL = f"https://github.com/HazyResearch/ThunderKittens/archive/{TK_COMMIT}.tar.gz" + +_SRC_DIR = Path(__file__).resolve().parent.parent / "csrc" / "tk_bwdbwd" +_SOURCES = ("stage1.cuh", "stage2.cuh", "ffi.cu") +_ARCH = "sm_90a" + +# CUDA pip wheels: consolidated cu13-era layout (nvidia/cu13/...) or per-component +# cu12-era layout (nvidia/cuda_nvcc/..., nvidia/cuda_runtime/...) +_WHEEL_ROOTS = { + "nvcc_bin": ("cu13/bin", "cu12/bin", "cuda_nvcc/bin"), + "include": ("cu13/include", "cu12/include", "cuda_runtime/include"), + "lib": ("cu13/lib", "cu12/lib", "cuda_runtime/lib"), +} + + +def _cache_root() -> Path: + return Path(os.environ.get("XDG_CACHE_HOME", str(Path.home() / ".cache"))) / "flash_hog" + + +def nvidia_wheel_paths(kind: str) -> list[Path]: + out = [] + for base in sys.path: + for sub in _WHEEL_ROOTS[kind]: + p = Path(base) / "nvidia" / sub + if p.exists() and p not in out: + out.append(p) + return out + + +def _find_nvcc() -> str | None: + for d in nvidia_wheel_paths("nvcc_bin"): + if (d / "nvcc").exists(): + return str(d / "nvcc") + cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH") + if cuda_home and (Path(cuda_home) / "bin" / "nvcc").exists(): + return str(Path(cuda_home) / "bin" / "nvcc") + return shutil.which("nvcc") + + +def _download(url: str, dest: Path) -> None: + import ssl + + try: + with urllib.request.urlopen(url) as r, open(dest, "wb") as f: + shutil.copyfileobj(r, f) + return + except urllib.error.URLError as e: + if not isinstance(getattr(e, "reason", None), ssl.SSLCertVerificationError): + raise + try: + import certifi + ctx = ssl.create_default_context(cafile=certifi.where()) + except ImportError: + raise RuntimeError( + f"downloading ThunderKittens failed ({url}): SSL verification failed and " + "certifi is unavailable. Install ca-certificates or set THUNDERKITTENS_PATH." + ) from None + with urllib.request.urlopen(url, context=ctx) as r, open(dest, "wb") as f: + shutil.copyfileobj(r, f) + + +def _thunderkittens() -> Path: + env = os.environ.get("THUNDERKITTENS_PATH") + if env: + return Path(env) + dst = _cache_root() / f"ThunderKittens-{TK_COMMIT}" + if (dst / "include").exists(): + return dst + dst.parent.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(dir=dst.parent) as td: + tar_path = Path(td) / "tk.tar.gz" + _download(_TK_TARBALL, tar_path) + with tarfile.open(tar_path) as tf: + tf.extractall(td) + extracted = Path(td) / f"ThunderKittens-{TK_COMMIT}" + if not (extracted / "include").exists(): + raise RuntimeError(f"unexpected ThunderKittens tarball layout under {extracted}") + os.replace(extracted, dst) # atomic against concurrent builders + return dst + + +def _source_hash() -> str: + h = hashlib.sha256() + for name in _SOURCES: + h.update((_SRC_DIR / name).read_bytes()) + h.update(TK_COMMIT.encode()) + h.update(_ARCH.encode()) + return h.hexdigest()[:16] + + +def cached_so() -> Path | None: + so = _cache_root() / "tk_bwdbwd" / _source_hash() / "libtk_bwdbwd.so" + return so if so.exists() else None + + +def build(force: bool = False, verbose: bool = True) -> Path: + """Build (or reuse) the plugin; returns the path to libtk_bwdbwd.so.""" + out_dir = _cache_root() / "tk_bwdbwd" / _source_hash() + so = out_dir / "libtk_bwdbwd.so" + if so.exists() and not force: + return so + + nvcc = _find_nvcc() + if nvcc is None: + raise RuntimeError("nvcc not found — `pip install 'flash-hog[tk]'` (or set CUDA_HOME).") + if not (shutil.which("g++") or shutil.which("clang++") or shutil.which("c++")): + raise RuntimeError("no host C++ compiler found (g++/clang++); nvcc needs one.") + + import jax # XLA FFI headers + + tk = _thunderkittens() + out_dir.mkdir(parents=True, exist_ok=True) + tmp = out_dir / "libtk_bwdbwd.so.tmp" + + cmd = [ + nvcc, "-shared", "-Xcompiler", "-fPIC", "-std=c++20", "-O3", "--use_fast_math", + "--expt-relaxed-constexpr", "--expt-extended-lambda", + "-forward-unknown-to-host-compiler", "-Xcompiler=-fno-strict-aliasing", + "-Xcompiler=-Wno-psabi", "-DNDEBUG", "-DKITTENS_SM90", + "-gencode", f"arch=compute_90a,code={_ARCH}", + f"-I{tk / 'include'}", f"-I{tk / 'prototype'}", f"-I{jax.ffi.include_dir()}", + "-cudart", "none", + ] + cmd += [f"-I{inc}" for inc in nvidia_wheel_paths("include")] + cmd += [str(_SRC_DIR / "ffi.cu"), "-o", str(tmp)] + + if verbose: + print(f"[flash-hog] building TK plugin (one-time, ~10-60 s): {so}", flush=True) + proc = subprocess.run(cmd, capture_output=True, text=True) + if proc.returncode != 0: + raise RuntimeError( + f"TK plugin build failed (rc={proc.returncode}):\n" + f"{' '.join(cmd)}\n{proc.stdout[-2000:]}\n{proc.stderr[-6000:]}" + ) + os.replace(tmp, so) + return so diff --git a/flash_hog/jax/_tk_gpu.py b/flash_hog/jax/_tk_gpu.py new file mode 100644 index 0000000..1d01eac --- /dev/null +++ b/flash_hog/jax/_tk_gpu.py @@ -0,0 +1,163 @@ +"""ThunderKittens (Hopper/SM90) double-backward for causal attention. + +Default path is Pallas; opt in to TK: + + pip install 'flash-hog[tk]' # CUDA build tools from PyPI + + from flash_hog.jax import _tk_gpu as tk + tk.enable() # JIT-builds the plugin on first use (cached); TK kernels live + tk.disable() # back to Pallas + +enable() swaps flash-hog's double-backward rule for the TK kernels where +supported() (causal, head_dim 64, seq % 128 == 0, no GQA, Hopper) and falls back +to the Pallas rule per-call otherwise. FLASH_HOG_TK_LIB overrides the plugin path. +""" + +from __future__ import annotations + +import ctypes +import functools +import os + +import jax +import jax.numpy as jnp +import numpy as np + +_LIB_ENV = "FLASH_HOG_TK_LIB" +_FFI_TARGET = "tk_bwdbwd" + + +def _preload_cuda_libs() -> None: + # plugin is built with `-cudart none`: resolve driver/cudart symbols from the + # process by loading them RTLD_GLOBAL first + from flash_hog.jax import _tk_build + + try: + ctypes.CDLL("libcuda.so.1", mode=ctypes.RTLD_GLOBAL) + except OSError: + pass + for libdir in _tk_build.nvidia_wheel_paths("lib"): + for so in sorted(libdir.glob("libcudart.so.*")): + try: + ctypes.CDLL(str(so), mode=ctypes.RTLD_GLOBAL) + return + except OSError: + continue + + +@functools.cache +def _lib() -> ctypes.CDLL | None: + from flash_hog.jax import _tk_build + + path = os.environ.get(_LIB_ENV) or _tk_build.cached_so() + if path is None or not os.path.exists(path): + return None + _preload_cuda_libs() + lib = ctypes.CDLL(str(path)) + jax.ffi.register_ffi_target(_FFI_TARGET, jax.ffi.pycapsule(lib.TkBwdBwd), platform="CUDA") + return lib + + +@functools.cache +def _on_hopper() -> bool: + try: + devices = jax.devices("gpu") + except RuntimeError: + return False + return all(getattr(d, "compute_capability", "") == "9.0" for d in devices) + + +def supported(*, is_causal: bool, seq_len: int, head_dim: int, + num_q_heads: int, num_kv_heads: int) -> bool: + """True iff the TK kernels can serve this shape on this machine.""" + return ( + is_causal + and head_dim == 64 + and seq_len % 128 == 0 + and num_q_heads == num_kv_heads # no GQA + and _on_hopper() + and _lib() is not None + ) + + +def flash_bwdbwd(*, Q, K, V, O, dO, ddQ, ddK, ddV, L, scale: float): + """Causal-attention double-backward. Arguments are BTNH; returns dQ2, dK2, dV2, ddO.""" + B, T, N, Hd = Q.shape + + def to_bhtd(x): + return jnp.transpose(x, (0, 2, 1, 3)) + + Qb, Kb, Vb, dOb, ddQb, ddKb, ddVb = ( + to_bhtd(x).astype(jnp.bfloat16) for x in (Q, K, V, dO, ddQ, ddK, ddV) + ) + D = jnp.sum(to_bhtd(dO).astype(jnp.float32) * to_bhtd(O).astype(jnp.float32), axis=-1) + Lf = L.reshape(B, N, T).astype(jnp.float32) + + outs = jax.ffi.ffi_call( + _FFI_TARGET, + [jax.ShapeDtypeStruct((B, N, T, Hd), jnp.bfloat16)] * 4 # dQ2, ddO, dK2, dV2 + + [jax.ShapeDtypeStruct((B, N, T), jnp.float32)] * 2, # dD, B (scratch) + )(Qb, Kb, Vb, dOb, ddQb, ddKb, ddVb, Lf, D, scale=np.float32(scale)) + + dQ2, ddO, dK2, dV2 = (to_bhtd(x).astype(Q.dtype) for x in outs[:4]) + return dQ2, dK2, dV2, ddO + + +_original_rule = None # non-None iff the TK rule is installed + + +def enable() -> None: + """Opt in: route the double-backward through the TK kernels (per-call fallback).""" + global _original_rule + if _lib() is None: + from flash_hog.jax import _tk_build + + _tk_build.build() + _lib.cache_clear() + if _lib() is None: + raise RuntimeError("TK plugin built but failed to load") + if _original_rule is not None: + return + + from jax._src.cudnn.fused_attention_stablehlo import MaskType + + from flash_hog.jax import _attention_impl as impl + + pallas_rule = impl.dot_product_attention_bwd_rule_bwd_rule + + def tk_rule(mask_type, scale, res, g): + query, key, value, out, activation, dO = res + if not supported( + is_causal=(mask_type == MaskType.CAUSAL), + seq_len=query.shape[1], + head_dim=query.shape[3], + num_q_heads=query.shape[2], + num_kv_heads=key.shape[2], + ): + return pallas_rule(mask_type, scale, res, g) + ddQ, ddK, ddV = g + dQ2, dK2, dV2, ddO = flash_bwdbwd( + Q=query, K=key, V=value, O=out, dO=dO, + ddQ=ddQ, ddK=ddK, ddV=ddV, L=activation, scale=scale, + ) + return (dQ2, dK2, dV2, None, None), ddO + + _original_rule = pallas_rule + impl.dot_product_attention_bwd_rule_bwd_rule = tk_rule + jax.clear_caches() # traced double-backwards captured the old rule + + +def disable() -> None: + """Restore the stock Pallas double-backward rule.""" + global _original_rule + if _original_rule is None: + return + from flash_hog.jax import _attention_impl as impl + + impl.dot_product_attention_bwd_rule_bwd_rule = _original_rule + _original_rule = None + jax.clear_caches() + + +def is_enabled() -> bool: + return _original_rule is not None diff --git a/pyproject.toml b/pyproject.toml index 63cde39..90b5e89 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,6 +11,14 @@ dependencies = [ "jax>=0.9.0; sys_platform == 'darwin'", ] +[project.optional-dependencies] +# CUDA build tools for JIT-compiling the ThunderKittens double-backward plugin +# (flash_hog/csrc/tk_bwdbwd) at first _tk_gpu.enable(); Hopper + linux only. +tk = [ + "nvidia-cuda-nvcc>=13; sys_platform == 'linux'", + "nvidia-cuda-runtime>=13; sys_platform == 'linux'", +] + [build-system] requires = ["uv_build>=0.10.6,<0.11.0"] build-backend = "uv_build" @@ -18,6 +26,8 @@ build-backend = "uv_build" [tool.uv.build-backend] module-name = "flash_hog" module-root = "" +# don't package a locally built TK plugin (csrc/tk_bwdbwd/libtk_bwdbwd.so) into the wheel +wheel-exclude = ["**/*.so"] [tool.uv]