From be16c169db597e225a7d94b019e837bee9109188 Mon Sep 17 00:00:00 2001 From: johnnynunez Date: Thu, 11 Jun 2026 19:45:01 +0200 Subject: [PATCH] fkl.flash_attention: FA-2 forward with fused epilogue + compressed int8 KV cache Python API over FusedKernelLibrary's new attention DPPs (FKL PR #256): out = fkl.flash_attention(q, k, v, causal=True) # FKL-only powers: kc, vc = fkl.compress_kv(k), fkl.compress_kv(v) # 4x smaller cache out = fkl.flash_attention(q, kc, vc, causal=True, epilogue=[fkl.Mul(2.0), fkl.Add(0.5)]) - epilogue=[ops]: any fkl compute ops run IN-REGISTER on the attention output inside the same kernel (codegen chains them with .then()); with handmade FA that is a second kernel + DRAM round-trip. Epilogue VALUES travel in params via the IOp's build() args -> changing them never recompiles (chain shape is part of the cache signature). - compress_kv(): GPU-side per-token int8 quantization (one fp32 scale per token; standalone JIT-compiled kernel). flash_attention() dequantizes in-register inside the fused kernel; the cache is never inflated. - FlashAttention class for explicit shape control; flash_attention() one-shot API with per-(head_dim, layout, epilogue-shape) kernel cache. tests/test_flash_attention.py (8 checks, fp64 numpy oracle): dense causal/ragged/cross, compression ratio >3.5x verified, int8-KV EXACT vs dequantized oracle (1e-7), end-to-end quant error bounded (4e-3), fused epilogue == host-applied (0 error), epilogue value changes without recompile, decode step (seq_q=1 vs 256-token compressed cache). Vendored headers refreshed from the feat/attention-dpps branch. --- fkl/__init__.py | 2 + fkl/attention.py | 260 ++++++++++++++++++++++++++++++++++ tests/test_flash_attention.py | 143 +++++++++++++++++++ 3 files changed, 405 insertions(+) create mode 100644 fkl/attention.py create mode 100644 tests/test_flash_attention.py diff --git a/fkl/__init__.py b/fkl/__init__.py index 4014458..50ebded 100644 --- a/fkl/__init__.py +++ b/fkl/__init__.py @@ -22,6 +22,7 @@ ) from .jit import compose, compose_divergent, FusedKernel, DivergentKernel from .circular import CircularTensor +from .attention import flash_attention, FlashAttention, compress_kv, CompressedKV from .backend import CompilerBackend, set_backend, get_backend, clear_cache from .tensor import DeviceBuffer from .types import DType, dtype @@ -37,6 +38,7 @@ "ColorConversion", "Crop", "Resize", "Warping", "BorderReader", "Deinterlace", "compose", "compose_divergent", "FusedKernel", "DivergentKernel", "CircularTensor", + "flash_attention", "FlashAttention", "compress_kv", "CompressedKV", "CompilerBackend", "set_backend", "get_backend", "clear_cache", "DeviceBuffer", "DType", "dtype", ] diff --git a/fkl/attention.py b/fkl/attention.py new file mode 100644 index 0000000..82b05f4 --- /dev/null +++ b/fkl/attention.py @@ -0,0 +1,260 @@ +"""fkl.attention — FlashAttention-2 forward on FKL's cooperative DPPs. + +What makes this different from calling a handmade FA kernel: + + 1. EPILOGUE FUSION: pass ops=[...] (any fkl compute ops) and they run + IN-REGISTER on the attention output inside the same kernel — no + second kernel, no DRAM round-trip: + fkl.flash_attention(q, k, v, epilogue=[fkl.Mul(2.0), fkl.Add(1.0)]) + 2. COMPRESSED KV CACHE: kv_layout="int8" stores K/V as int8 with one + fp32 scale per token (4x smaller than fp32, 2x than fp16) and + dequantizes in-register inside the fused kernel: + kc = fkl.compress_kv(k); vc = fkl.compress_kv(v) + out = fkl.flash_attention(q, kc, vc, causal=True) + +Layout: (batch*heads, seq, head_dim) C-contiguous; head_dim in {32,64,128,...}. +fp32 accumulation always. Targets SM 12x (SIMT mapping per fa-5090; no +TMEM). Compiled once per (dtype, head_dim, layout, epilogue-shape, arch); +runtime values (scales, sizes, causal flag, epilogue params) never +recompile. +""" +from __future__ import annotations +import ctypes +import math + +from .backend import get_backend +from .codegen import CODEGEN_VERSION +from .operations import READ, WRITE, ChainState +from .tensor import DeviceBuffer, as_device_view, stream_handle +from .types import dtype as _dtype + + +class CompressedKV: + """int8-per-token compressed KV tensor + per-token scales (on GPU).""" + + def __init__(self, data: DeviceBuffer, scales: DeviceBuffer, + batch_heads: int, seq: int, head_dim: int): + self.data, self.scales = data, scales + self.batch_heads, self.seq, self.head_dim = batch_heads, seq, head_dim + + @property + def nbytes(self) -> int: + return self.data._nbytes + self.scales._nbytes + + +_quant_lib = None + + +def _quant_so(): + """Tiny standalone kernel for GPU-side per-token int8 quantization.""" + global _quant_lib + if _quant_lib is not None: + return _quant_lib + src = r""" +#include +#include +extern "C" { +__global__ void quant_kernel(const float* dense, signed char* q8, float* scales, + int tokens, int headDim) { + const int t = blockIdx.x; + if (t >= tokens) return; + __shared__ float smax[256]; + float mx = 0.f; + for (int d = threadIdx.x; d < headDim; d += blockDim.x) + mx = fmaxf(mx, fabsf(dense[(long)t * headDim + d])); + smax[threadIdx.x] = mx; + __syncthreads(); + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + if (threadIdx.x < s) smax[threadIdx.x] = fmaxf(smax[threadIdx.x], smax[threadIdx.x + s]); + __syncthreads(); + } + const float sc = smax[0] > 0.f ? smax[0] / 127.f : 1.f; + if (threadIdx.x == 0) scales[t] = sc; + for (int d = threadIdx.x; d < headDim; d += blockDim.x) { + q8[(long)t * headDim + d] = (signed char)nearbyintf(dense[(long)t * headDim + d] / sc); + } +} +void quantize(const float* dense, signed char* q8, float* scales, + int tokens, int headDim, void* stream) { + quant_kernel<<>>(dense, q8, scales, tokens, headDim); + cudaStreamSynchronize((cudaStream_t)stream); +} +} +""" + so = get_backend().compile(src, f"kvquant;cg={CODEGEN_VERSION}") + lib = ctypes.CDLL(str(so)) + lib.quantize.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, + ctypes.c_int, ctypes.c_int, ctypes.c_void_p] + _quant_lib = lib + return lib + + +def compress_kv(kv) -> CompressedKV: + """Compress a (batch_heads, seq, head_dim) float32 K or V tensor to the + int8-per-token layout ON the GPU. ~4x memory reduction vs fp32.""" + v = as_device_view(kv) + if v.dtype.base != "float32": + raise TypeError("compress_kv expects float32 input (cast first)") + bh, seq, hd = v.planes, v.height, v.width + if bh == 1 and v.planes == 1: + raise ValueError("expected 3D (batch_heads, seq, head_dim)") + tokens = bh * seq + q8 = DeviceBuffer(hd, seq, "int8", planes=bh) + scales = DeviceBuffer(tokens, 1, "float32") + _quant_so().quantize(ctypes.c_void_p(v.ptr), ctypes.c_void_p(q8.ptr), + ctypes.c_void_p(scales.ptr), tokens, hd, + ctypes.c_void_p(0)) + return CompressedKV(q8, scales, bh, seq, hd) + + +class FlashAttention: + """Compiled FA-2 forward for one (head_dim, dtype, layout, epilogue) shape.""" + + def __init__(self, head_dim: int, kv_layout: str = "dense", + epilogue=None): + if head_dim % 32 != 0: + raise ValueError("head_dim must be a multiple of 32") + if kv_layout not in ("dense", "int8"): + raise ValueError("kv_layout must be 'dense' or 'int8'") + self.head_dim = head_dim + self.kv_layout = kv_layout + self.epilogue = list(epilogue or []) + for op in self.epilogue: + if op.role in (READ, WRITE): + raise ValueError("epilogue ops must be compute-only") + self._lib = None + + # ---- public ------------------------------------------------------------ + def __call__(self, q, k, v, out=None, causal=False, scale=None, + stream=None): + vq = as_device_view(q) + bh, seq_q, hd = vq.planes, vq.height, vq.width + if hd != self.head_dim: + raise ValueError(f"q head_dim {hd} != compiled {self.head_dim}") + + if self.kv_layout == "int8": + if not isinstance(k, CompressedKV) or not isinstance(v, CompressedKV): + raise TypeError("int8 layout expects CompressedKV (use fkl.compress_kv)") + seq_k = k.seq + kptr, vptr = k.data.ptr, v.data.ptr + ks, vs = k.scales.ptr, v.scales.ptr + else: + vk, vv = as_device_view(k), as_device_view(v) + seq_k = vk.height + kptr, vptr, ks, vs = vk.ptr, vv.ptr, 0, 0 + + self._ensure_compiled(vq.dtype.base) + if out is None: + out = DeviceBuffer(hd, seq_q, "float32", planes=bh) + vout = as_device_view(out) + + eps = [] + dt = _dtype("float32") + for op in self.epilogue: + if hasattr(op, "bind"): + op.bind(dt) + eps.extend(op.values) + dt = op.out_dtype(dt) + pbuf = (ctypes.c_float * max(1, len(eps)))(*eps) + + sc = float(scale) if scale else 1.0 / math.sqrt(self.head_dim) + self._lib.fa_forward(ctypes.c_void_p(vq.ptr), ctypes.c_void_p(kptr), + ctypes.c_void_p(vptr), ctypes.c_void_p(ks), + ctypes.c_void_p(vs), ctypes.c_void_p(vout.ptr), + bh, seq_q, seq_k, ctypes.c_float(sc), + 1 if causal else 0, + ctypes.cast(pbuf, ctypes.c_void_p), + ctypes.c_void_p(stream_handle(stream))) + return out + + # ---- compilation --------------------------------------------------------- + def _ensure_compiled(self, in_dtype: str): + if self._lib is not None: + return + if in_dtype != "float32": # base name + raise TypeError("flash_attention currently takes float32 q/k/v") + from .jit import _ARCH + + # epilogue chain -> C++ expression chained with .then() + st = ChainState(_dtype("float32"), self.head_dim, 1, 1) + exprs, pbase, toks = [], 0, [] + dt = _dtype("float32") + for op in self.epilogue: + if hasattr(op, "bind"): + op.bind(dt) + exprs.append(op.cpp(st, pbase)) + toks.append(op.token(st)) + pbase += len(op.values) + dt = op.out_dtype(dt) + if exprs: + ep_build = exprs[0] + "".join(f".then({e})" for e in exprs[1:]) + ep_type = f"decltype({ep_build})" + else: + ep_build = "AttentionIdentityEpilogue{}" + ep_type = "AttentionIdentityEpilogue" + + kvl = ("KVLayout::INT8_PER_TOKEN" if self.kv_layout == "int8" + else "KVLayout::DENSE") + kvt = "signed char" if self.kv_layout == "int8" else "float" + + src = f"""// AUTO-GENERATED by fkl-python (FlashAttention). Host+device TU. +#include +#include +#include +#include + +using namespace fk; + +extern "C" {{ +void fa_forward(const float* q, const void* k, const void* v, + const float* kScale, const float* vScale, float* o, + int bh, int seqQ, int seqK, float scale, int causal, + const float* params, void* ext_stream) {{ + const auto epilogue = {ep_build}; + using Ep = {ep_type}; + if (ext_stream != nullptr) {{ + Stream stream(reinterpret_cast(ext_stream)); + executeFlashAttention= 64 ? 32 : 32), 4, Ep>( + q, (const {kvt}*)k, (const {kvt}*)v, o, bh, seqQ, seqK, + causal != 0, stream, kScale, vScale, scale, epilogue); + }} else {{ + static Stream stream; + executeFlashAttention= 64 ? 32 : 32), 4, Ep>( + q, (const {kvt}*)k, (const {kvt}*)v, o, bh, seqQ, seqK, + causal != 0, stream, kScale, vScale, scale, epilogue); + stream.sync(); + }} +}} +}} // extern "C" +""" + sig = (f"flashattn;arch={_ARCH};cg={CODEGEN_VERSION};d={self.head_dim};" + f"kv={self.kv_layout};ep=" + "|".join(toks)) + so = get_backend().compile(src, sig) + lib = ctypes.CDLL(str(so)) + lib.fa_forward.argtypes = [ctypes.c_void_p] * 6 + [ + ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_float, + ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p] + self._lib = lib + + +_fa_cache = {} + + +def flash_attention(q, k, v, out=None, causal=False, scale=None, stream=None, + epilogue=None, kv_layout=None): + """One-shot API. Auto-detects head_dim and kv layout; compiled kernels + are cached per (head_dim, layout, epilogue shape).""" + if kv_layout is None: + kv_layout = "int8" if isinstance(k, CompressedKV) else "dense" + vq = as_device_view(q) + ep_key = tuple(type(op).__name__ for op in (epilogue or [])) + key = (vq.width, kv_layout, ep_key) + fa = _fa_cache.get(key) + if fa is None: + fa = FlashAttention(vq.width, kv_layout, epilogue) + _fa_cache[key] = fa + elif epilogue: + fa.epilogue = list(epilogue) # same shape, fresh values + return fa(q, k, v, out=out, causal=causal, scale=scale, stream=stream) diff --git a/tests/test_flash_attention.py b/tests/test_flash_attention.py new file mode 100644 index 0000000..72c26c1 --- /dev/null +++ b/tests/test_flash_attention.py @@ -0,0 +1,143 @@ +"""Tests for fkl.flash_attention: FA-2 forward vs fp64 numpy oracle, +dense + compressed (int8) KV cache + fused epilogue. Requires numpy.""" +import math +import struct +import sys + +try: + import numpy as np +except ImportError: + print("SKIP: numpy required") + sys.exit(0) + +import fkl +from harness import check_true, run + + +def _oracle(q, k, v, causal, scale): + """fp64 reference: softmax(scale * QK^T [causal]) V per (bh).""" + q, k, v = q.astype(np.float64), k.astype(np.float64), v.astype(np.float64) + bh, sq, d = q.shape + sk = k.shape[1] + out = np.zeros((bh, sq, d)) + for b in range(bh): + s = q[b] @ k[b].T * scale # (sq, sk) + if causal: + mask = np.triu(np.ones((sq, sk), dtype=bool), 1) + s[mask] = -np.inf + s = s - s.max(axis=1, keepdims=True) + p = np.exp(s) + p /= p.sum(axis=1, keepdims=True) + out[b] = p @ v[b] + return out + + +def _to_gpu(arr): + bh, seq, d = arr.shape + buf = fkl.DeviceBuffer(d, seq, "float32", planes=bh) + buf.copy_from_host(arr.astype(np.float32).tobytes()) + return buf + + +def _from_gpu(buf, shape): + return np.frombuffer(buf.copy_to_host(), dtype=np.float32).reshape(shape) + + +def _mk(bh, sq, sk, d, seed): + rng = np.random.default_rng(seed) + return (rng.uniform(-1, 1, (bh, sq, d)).astype(np.float32), + rng.uniform(-1, 1, (bh, sk, d)).astype(np.float32), + rng.uniform(-1, 1, (bh, sk, d)).astype(np.float32)) + + +def t_dense_causal(): + q, k, v = _mk(2, 64, 64, 64, 1) + out = fkl.flash_attention(_to_gpu(q), _to_gpu(k), _to_gpu(v), causal=True) + got = _from_gpu(out, q.shape) + ref = _oracle(q, k, v, True, 1 / math.sqrt(64)) + err = np.abs(got - ref).max() + check_true(f"FA dense d64 causal (err={err:.2e})", err < 5e-6) + + +def t_dense_cross_ragged(): + q, k, v = _mk(2, 33, 127, 32, 2) + out = fkl.flash_attention(_to_gpu(q), _to_gpu(k), _to_gpu(v)) + got = _from_gpu(out, q.shape) + ref = _oracle(q, k, v, False, 1 / math.sqrt(32)) + err = np.abs(got - ref).max() + check_true(f"FA dense d32 ragged 33/127 (err={err:.2e})", err < 5e-6) + + +def t_compressed_kv(): + """int8 KV cache: kernel must be EXACT vs oracle-on-dequantized values, + and END-TO-END accuracy vs the original values must be quantization- + bounded. Also verifies the 4x memory saving.""" + q, k, v = _mk(2, 48, 96, 64, 3) + kc = fkl.compress_kv(_to_gpu(k)) + vc = fkl.compress_kv(_to_gpu(v)) + + # memory: int8 data + 1 float/token vs fp32 dense + dense_bytes = k.size * 4 + comp_bytes = kc.nbytes + ratio = dense_bytes / comp_bytes + check_true(f"KV compression ratio {ratio:.2f}x (>3.5x)", ratio > 3.5) + + out = fkl.flash_attention(_to_gpu(q), kc, vc, causal=True) + got = _from_gpu(out, q.shape) + + # exactness vs dequantized oracle + k8 = np.frombuffer(kc.data.copy_to_host(), dtype=np.int8).reshape(k.shape) + ks = np.frombuffer(kc.scales.copy_to_host(), dtype=np.float32).reshape(2, 96, 1) + v8 = np.frombuffer(vc.data.copy_to_host(), dtype=np.int8).reshape(v.shape) + vs = np.frombuffer(vc.scales.copy_to_host(), dtype=np.float32).reshape(2, 96, 1) + kd, vd = k8 * ks, v8 * vs + ref_dq = _oracle(q, kd, vd, True, 1 / math.sqrt(64)) + err_exact = np.abs(got - ref_dq).max() + check_true(f"FA int8-KV exact vs dequantized oracle (err={err_exact:.2e})", + err_exact < 5e-6) + + # end-to-end: quantization-bounded vs original + ref_full = _oracle(q, k, v, True, 1 / math.sqrt(64)) + err_e2e = np.abs(got - ref_full).max() + check_true(f"FA int8-KV end-to-end quant error bounded (err={err_e2e:.2e})", + err_e2e < 2e-2) + + +def t_fused_epilogue(): + """epilogue=[Mul, Add] runs inside the kernel; equals host-applied.""" + q, k, v = _mk(1, 32, 32, 32, 4) + base = _from_gpu(fkl.flash_attention(_to_gpu(q), _to_gpu(k), _to_gpu(v)), + q.shape) + fused = _from_gpu( + fkl.flash_attention(_to_gpu(q), _to_gpu(k), _to_gpu(v), + epilogue=[fkl.Mul(2.0), fkl.Add(0.5)]), q.shape) + err = np.abs(fused - (base * 2.0 + 0.5)).max() + check_true(f"FA fused epilogue == host-applied (err={err:.2e})", err < 1e-6) + + +def t_epilogue_values_no_recompile(): + q, k, v = _mk(1, 16, 16, 32, 5) + a = _from_gpu(fkl.flash_attention(_to_gpu(q), _to_gpu(k), _to_gpu(v), + epilogue=[fkl.Mul(3.0)]), q.shape) + b = _from_gpu(fkl.flash_attention(_to_gpu(q), _to_gpu(k), _to_gpu(v), + epilogue=[fkl.Mul(5.0)]), q.shape) + base = _from_gpu(fkl.flash_attention(_to_gpu(q), _to_gpu(k), _to_gpu(v)), + q.shape) + ok = (np.abs(a - base * 3).max() < 1e-5 and np.abs(b - base * 5).max() < 1e-5) + check_true("FA epilogue values change without recompile", ok) + + +def t_single_query_decode(): + """seq_q=1 (autoregressive decode step) against a long compressed cache.""" + q, k, v = _mk(2, 1, 256, 64, 6) + kc, vc = fkl.compress_kv(_to_gpu(k)), fkl.compress_kv(_to_gpu(v)) + got = _from_gpu(fkl.flash_attention(_to_gpu(q), kc, vc), q.shape) + ref = _oracle(q, k, v, False, 1 / math.sqrt(64)) + err = np.abs(got - ref).max() + check_true(f"FA decode step s_q=1 vs s_k=256 int8 (err={err:.2e})", err < 2e-2) + + +if __name__ == "__main__": + run([t_dense_causal, t_dense_cross_ragged, t_compressed_kv, + t_fused_epilogue, t_epilogue_values_no_recompile, + t_single_query_decode], "flash-attention")