diff --git a/benchmarks/bench_flash_attention.py b/benchmarks/bench_flash_attention.py new file mode 100644 index 0000000..a28b6ab --- /dev/null +++ b/benchmarks/bench_flash_attention.py @@ -0,0 +1,162 @@ +"""Benchmark + accuracy: fkl.flash_attention vs PyTorch SDPA (FlashAttention-2 +backend) and, if installed, the official flash-attn package. + +Run with the torch venv: + .venv-torch/bin/python benchmarks/bench_flash_attention.py + +Measures (batch*heads=32, head_dim=64, fp32 for fkl / fp16 for references): + - accuracy of every implementation vs an fp64 oracle + - latency + effective TFLOPS + - the FUSION advantage: attention + post-op as ONE fkl kernel vs + SDPA + separate post-op kernel (what handmade FA forces you to do) + - compressed KV cache: memory + decode-step latency +""" +import math +import time + +import numpy as np +import torch + +import fkl + +DEV = "cuda" +BH, D = 32, 64 + + +def oracle(q, k, v, causal, scale): + q64 = q.double().cpu().numpy() + k64 = k.double().cpu().numpy() + v64 = v.double().cpu().numpy() + out = np.zeros_like(q64) + for b in range(q64.shape[0]): + s = q64[b] @ k64[b].T * scale + if causal: + s[np.triu_indices_from(s, 1)] = -np.inf + s -= s.max(1, keepdims=True) + p = np.exp(s) + p /= p.sum(1, keepdims=True) + out[b] = p @ v64[b] + return out + + +def timeit(fn, iters=50): + fn() # warmup + compile + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(iters): + fn() + torch.cuda.synchronize() + return (time.perf_counter() - t0) / iters * 1e3 # ms + + +def tflops(ms, seq_q, seq_k, causal): + flops = 4.0 * BH * seq_q * seq_k * D * (0.5 if causal else 1.0) + return flops / (ms * 1e-3) / 1e12 + + +def main(): + torch.manual_seed(7) + print(f"config: batch*heads={BH}, head_dim={D}, RTX PRO 6000 (sm_120)\n") + + # ---------- accuracy ---------- + print("== accuracy vs fp64 oracle (seq 256, causal) ==") + q = torch.empty(BH, 256, D, device=DEV).uniform_(-1, 1) + k = torch.empty(BH, 256, D, device=DEV).uniform_(-1, 1) + v = torch.empty(BH, 256, D, device=DEV).uniform_(-1, 1) + scale = 1 / math.sqrt(D) + ref = oracle(q, k, v, True, scale) + + out_fkl = torch.from_dlpack(fkl.flash_attention(q, k, v, causal=True)) + err_fkl = np.abs(out_fkl.cpu().numpy().reshape(ref.shape) - ref).max() + print(f"fkl fp32 dense : {err_fkl:.3e}") + + with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.FLASH_ATTENTION): + out_sdpa16 = torch.nn.functional.scaled_dot_product_attention( + q.half(), k.half(), v.half(), is_causal=True) + err_sdpa16 = np.abs(out_sdpa16.float().cpu().numpy() - ref).max() + print(f"sdpa fp16 flash : {err_sdpa16:.3e} (PyTorch FA-2 kernel)") + + kc, vc = fkl.compress_kv(k), fkl.compress_kv(v) + out_q = torch.from_dlpack(fkl.flash_attention(q, kc, vc, causal=True)) + err_q = np.abs(out_q.cpu().numpy().reshape(ref.shape) - ref).max() + print(f"fkl int8-KV : {err_q:.3e} (4x smaller KV cache)") + + try: + from flash_attn import flash_attn_func + # flash-attn wants (batch, seq, heads, dim) fp16 + q4 = q.half().view(4, 8, 256, D).permute(0, 2, 1, 3).contiguous() + k4 = k.half().view(4, 8, 256, D).permute(0, 2, 1, 3).contiguous() + v4 = v.half().view(4, 8, 256, D).permute(0, 2, 1, 3).contiguous() + out_official = flash_attn_func(q4, k4, v4, causal=True) + oo = out_official.permute(0, 2, 1, 3).reshape(BH, 256, D) + err_off = np.abs(oo.float().cpu().numpy() - ref).max() + print(f"flash-attn fp16 : {err_off:.3e} (official Dao-AILab)") + have_official = True + except ImportError: + print("flash-attn : not installed (skipping)") + have_official = False + + # ---------- performance ---------- + print("\n== latency / TFLOPS (seq 1024, causal) ==") + q = torch.empty(BH, 1024, D, device=DEV).uniform_(-1, 1) + k = torch.empty(BH, 1024, D, device=DEV).uniform_(-1, 1) + v = torch.empty(BH, 1024, D, device=DEV).uniform_(-1, 1) + qh, kh, vh = q.half(), k.half(), v.half() + + fa = fkl.FlashAttention(D) # reuse compiled kernel + out_buf = fkl.DeviceBuffer(D, 1024, "float32", planes=BH) + ms_fkl = timeit(lambda: fa(q, k, v, out=out_buf, causal=True)) + print(f"fkl fp32 SIMT : {ms_fkl:7.3f} ms {tflops(ms_fkl,1024,1024,True):6.2f} TFLOPS") + + def run_sdpa(): + with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.FLASH_ATTENTION): + return torch.nn.functional.scaled_dot_product_attention(qh, kh, vh, is_causal=True) + ms_sdpa = timeit(run_sdpa) + print(f"sdpa fp16 flash : {ms_sdpa:7.3f} ms {tflops(ms_sdpa,1024,1024,True):6.2f} TFLOPS") + + if have_official: + q4 = qh.view(4, 8, 1024, D).permute(0, 2, 1, 3).contiguous() + k4 = kh.view(4, 8, 1024, D).permute(0, 2, 1, 3).contiguous() + v4 = vh.view(4, 8, 1024, D).permute(0, 2, 1, 3).contiguous() + ms_off = timeit(lambda: flash_attn_func(q4, k4, v4, causal=True)) + print(f"flash-attn fp16 : {ms_off:7.3f} ms {tflops(ms_off,1024,1024,True):6.2f} TFLOPS") + + # ---------- the fusion advantage ---------- + print("\n== fusion: attention + Mul(2)+Add(0.5) post-op (seq 1024) ==") + ms_fused = timeit(lambda: fa(q, k, v, out=out_buf, causal=True)) + fa_ep = fkl.FlashAttention(D, epilogue=[fkl.Mul(2.0), fkl.Add(0.5)]) + ms_fused_ep = timeit(lambda: fa_ep(q, k, v, out=out_buf, causal=True)) + print(f"fkl FA alone : {ms_fused:7.3f} ms") + print(f"fkl FA+epilogue FUSED : {ms_fused_ep:7.3f} ms " + f"(overhead {100*(ms_fused_ep-ms_fused)/ms_fused:+.1f}%)") + + def sdpa_then_postop(): + o = run_sdpa() + return o * 2.0 + 0.5 # extra kernel(s) + DRAM trip + ms_sdpa_post = timeit(sdpa_then_postop) + print(f"sdpa + separate post-op: {ms_sdpa_post:7.3f} ms " + f"(post-op cost {100*(ms_sdpa_post-ms_sdpa)/ms_sdpa:+.1f}% vs sdpa alone)") + + # ---------- compressed KV decode ---------- + print("\n== decode step (seq_q=1) vs 8192-token KV cache ==") + kl = torch.empty(BH, 8192, D, device=DEV).uniform_(-1, 1) + vl = torch.empty(BH, 8192, D, device=DEV).uniform_(-1, 1) + q1 = torch.empty(BH, 1, D, device=DEV).uniform_(-1, 1) + dense_mb = kl.numel() * 4 * 2 / 2**20 + kc, vc = fkl.compress_kv(kl), fkl.compress_kv(vl) + comp_mb = (kc.nbytes + vc.nbytes) / 2**20 + print(f"KV cache: dense fp32 {dense_mb:.0f} MB -> int8 {comp_mb:.0f} MB " + f"({dense_mb/comp_mb:.2f}x)") + + fa_dec = fkl.FlashAttention(D) + out1 = fkl.DeviceBuffer(D, 1, "float32", planes=BH) + ms_dense = timeit(lambda: fa_dec(q1, kl, vl, out=out1), iters=200) + fa_decq = fkl.FlashAttention(D, kv_layout="int8") + ms_int8 = timeit(lambda: fa_decq(q1, kc, vc, out=out1), iters=200) + print(f"decode dense fp32 : {ms_dense*1e3:7.1f} us") + print(f"decode int8 KV : {ms_int8*1e3:7.1f} us " + f"({ms_dense/ms_int8:.2f}x faster — bandwidth-bound)") + + +if __name__ == "__main__": + main() diff --git a/fkl/__init__.py b/fkl/__init__.py index 50ebded..13bb6c7 100644 --- a/fkl/__init__.py +++ b/fkl/__init__.py @@ -23,6 +23,7 @@ from .jit import compose, compose_divergent, FusedKernel, DivergentKernel from .circular import CircularTensor from .attention import flash_attention, FlashAttention, compress_kv, CompressedKV +from .flex import ScoreMod, ALiBi, SoftCap, SlidingWindow, BlockMask from .backend import CompilerBackend, set_backend, get_backend, clear_cache from .tensor import DeviceBuffer from .types import DType, dtype @@ -39,6 +40,7 @@ "compose", "compose_divergent", "FusedKernel", "DivergentKernel", "CircularTensor", "flash_attention", "FlashAttention", "compress_kv", "CompressedKV", + "ScoreMod", "ALiBi", "SoftCap", "SlidingWindow", "BlockMask", "CompilerBackend", "set_backend", "get_backend", "clear_cache", "DeviceBuffer", "DType", "dtype", ] diff --git a/fkl/attention.py b/fkl/attention.py index 82b05f4..76a5884 100644 --- a/fkl/attention.py +++ b/fkl/attention.py @@ -30,12 +30,13 @@ class CompressedKV: - """int8-per-token compressed KV tensor + per-token scales (on GPU).""" + """Compressed KV tensor + per-token scales (on GPU). fmt: 'int8' | 'fp8'.""" def __init__(self, data: DeviceBuffer, scales: DeviceBuffer, - batch_heads: int, seq: int, head_dim: int): + batch_heads: int, seq: int, head_dim: int, fmt: str = "int8"): self.data, self.scales = data, scales self.batch_heads, self.seq, self.head_dim = batch_heads, seq, head_dim + self.fmt = fmt @property def nbytes(self) -> int: @@ -46,16 +47,17 @@ def nbytes(self) -> int: def _quant_so(): - """Tiny standalone kernel for GPU-side per-token int8 quantization.""" + """Tiny standalone kernel for GPU-side per-token int8/fp8 quantization.""" global _quant_lib if _quant_lib is not None: return _quant_lib src = r""" #include +#include #include -extern "C" { -__global__ void quant_kernel(const float* dense, signed char* q8, float* scales, - int tokens, int headDim) { +template +__global__ void quant_kernel_t(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]; @@ -68,30 +70,46 @@ def _quant_so(): 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; + const float sc = smax[0] > 0.f ? smax[0] / (FP8 ? 448.f : 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); + const float x = dense[(long)t * headDim + d] / sc; + if (FP8) { + const __nv_fp8_e4m3 f8(x); + q8[(long)t * headDim + d] = (signed char)f8.__x; + } else { + q8[(long)t * headDim + d] = (signed char)nearbyintf(x); + } } } +extern "C" { void quantize(const float* dense, signed char* q8, float* scales, int tokens, int headDim, void* stream) { - quant_kernel<<>>(dense, q8, scales, tokens, headDim); + quant_kernel_t<<>>(dense, q8, scales, tokens, headDim); + cudaStreamSynchronize((cudaStream_t)stream); +} +void quantize_fp8(const float* dense, signed char* q8, float* scales, + int tokens, int headDim, void* stream) { + quant_kernel_t<<>>(dense, q8, scales, tokens, headDim); cudaStreamSynchronize((cudaStream_t)stream); } } """ - so = get_backend().compile(src, f"kvquant;cg={CODEGEN_VERSION}") + so = get_backend().compile(src, f"kvquant;cg={CODEGEN_VERSION};v=2") 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] + for fn in (lib.quantize, lib.quantize_fp8): + fn.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: +def compress_kv(kv, fmt: str = "int8") -> 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.""" + int8- or fp8(e4m3)-per-token layout ON the GPU. ~4x smaller than fp32. + fp8 keeps more precision near zero (attention tails) at the same size.""" + if fmt not in ("int8", "fp8"): + raise ValueError("fmt must be 'int8' or 'fp8'") v = as_device_view(kv) if v.dtype.base != "float32": raise TypeError("compress_kv expects float32 input (cast first)") @@ -101,40 +119,67 @@ def compress_kv(kv) -> CompressedKV: 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) + fn = _quant_so().quantize_fp8 if fmt == "fp8" else _quant_so().quantize + fn(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, fmt) class FlashAttention: - """Compiled FA-2 forward for one (head_dim, dtype, layout, epilogue) shape.""" + """Compiled FA-2 forward for one (head_dim, dtype, layout, prologue, + epilogue) shape. + + PROLOGUE FUSION (mirrors the C++ IOp-first API): prologue_q/k/v are + lists of fkl compute ops fused onto the Q/K/V *Read IOps* — every + element enters the kernel through read.then(op1).then(op2)..., + in-register, at load time. With kv_layout="int8" the K/V prologue + chains AFTER the dequantizing Int8TokenDequantRead. Op VALUES live in + params[] — changing them never recompiles.""" def __init__(self, head_dim: int, kv_layout: str = "dense", - epilogue=None): + epilogue=None, prologue_q=None, prologue_k=None, + prologue_v=None, mma: bool = False, score_mod=None, + block_sparse: bool = False): 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'") + if kv_layout not in ("dense", "int8", "fp8"): + raise ValueError("kv_layout must be 'dense', 'int8' or 'fp8'") + if (score_mod is not None or block_sparse) and not mma: + raise ValueError("score_mod/block_mask require mma=True " + "(tensor-core path)") self.head_dim = head_dim self.kv_layout = kv_layout + self.mma = bool(mma) # tensor-core (bf16 mma.sync) path + self.score_mod = score_mod # flex-attention score mod + self.block_sparse = bool(block_sparse) self.epilogue = list(epilogue or []) - for op in self.epilogue: + self.prologue_q = list(prologue_q or []) + self.prologue_k = list(prologue_k or []) + self.prologue_v = list(prologue_v or []) + for op in (self.epilogue + self.prologue_q + self.prologue_k + + self.prologue_v): if op.role in (READ, WRITE): - raise ValueError("epilogue ops must be compute-only") + raise ValueError("prologue/epilogue ops must be compute-only") self._lib = None # ---- public ------------------------------------------------------------ def __call__(self, q, k, v, out=None, causal=False, scale=None, - stream=None): + stream=None, block_mask=None): + if self.block_sparse != (block_mask is not None): + raise ValueError("kernel compiled with block_sparse=" + f"{self.block_sparse} but block_mask is " + f"{'set' if block_mask else '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 self.kv_layout in ("int8", "fp8"): if not isinstance(k, CompressedKV) or not isinstance(v, CompressedKV): - raise TypeError("int8 layout expects CompressedKV (use fkl.compress_kv)") + raise TypeError("compressed layout expects CompressedKV (use fkl.compress_kv)") + if k.fmt != self.kv_layout or v.fmt != self.kv_layout: + raise TypeError(f"CompressedKV fmt mismatch: kernel={self.kv_layout}, " + f"k={k.fmt}, v={v.fmt}") seq_k = k.seq kptr, vptr = k.data.ptr, v.data.ptr ks, vs = k.scales.ptr, v.scales.ptr @@ -149,93 +194,158 @@ def __call__(self, q, k, v, out=None, causal=False, scale=None, 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) + for ops in (self.prologue_q, self.prologue_k, self.prologue_v, + self.epilogue): + dt = _dtype("float32") + for op in ops: + 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) + + # flex score-mod runtime values (slope/cap/window — never recompiles) + mvals = list(self.score_mod.values) if self.score_mod else [] + mbuf = (ctypes.c_float * max(1, len(mvals)))(*mvals) + + # block sparsity: mask ptr + geometry (0s = dense) + if block_mask is not None: + mk_ptr, nqb, nkb = block_mask.ptr, block_mask.n_q_blocks, block_mask.n_kv_blocks + mbq, mbkv = block_mask.block_q, block_mask.block_kv + else: + mk_ptr, nqb, nkb, mbq, mbkv = 0, 0, 0, 128, 128 + 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.cast(mbuf, ctypes.c_void_p), + ctypes.c_void_p(mk_ptr), nqb, nkb, mbq, mbkv, 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() + def _chain(self, ops, pbase): + """ops -> (list of C++ exprs, tokens, new pbase).""" st = ChainState(_dtype("float32"), self.head_dim, 1, 1) - exprs, pbase, toks = [], 0, [] + exprs, toks = [], [] dt = _dtype("float32") - for op in self.epilogue: + for op in ops: 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})" + return exprs, toks, pbase + + 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 + + # prologue chains fuse onto the Read IOps; epilogue onto the output. + pbase = 0 + q_ex, q_tk, pbase = self._chain(self.prologue_q, pbase) + k_ex, k_tk, pbase = self._chain(self.prologue_k, pbase) + v_ex, v_tk, pbase = self._chain(self.prologue_v, pbase) + e_ex, e_tk, pbase = self._chain(self.epilogue, pbase) + + def fuse(base, exprs): + return base + "".join(f".then({e})" for e in exprs) + + if self.kv_layout == "int8": + k_read = "makeInt8KVRead((const signed char*)k, kScale, bh, seqK, HD)" + v_read = "makeInt8KVRead((const signed char*)v, vScale, bh, seqK, HD)" + elif self.kv_layout == "fp8": + k_read = "makeFp8KVRead(k, kScale, bh, seqK, HD)" + v_read = "makeFp8KVRead(v, vScale, bh, seqK, HD)" + else: + k_read = "makeAttentionRead((const float*)k, bh, seqK, HD)" + v_read = "makeAttentionRead((const float*)v, bh, seqK, HD)" + q_build = fuse("makeAttentionRead(q, bh, seqQ, HD)", q_ex) + k_build = fuse(k_read, k_ex) + v_build = fuse(v_read, v_ex) + + if e_ex: + ep_build = fuse(e_ex[0], e_ex[1:]) 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" + exec_fn = ("executeFlashAttentionMma" if self.mma + else "executeFlashAttention") + + # flex score mod + block sparsity (mma path only; validated in ctor) + mod_ctor = (self.score_mod.cpp_ctor if self.score_mod + else "NoScoreMod{}") + if self.mma: + sparse_decl = (" const BlockSparsity sparse{ blockMask, nQB, nKB, " + "maskBQ, maskBKV };\n") + extra_args = ", epilogue, scoreMod, sparse" + else: + sparse_decl = "" + extra_args = ", epilogue" + + call_tail = extra_args if self.mma else ", epilogue" src = f"""// AUTO-GENERATED by fkl-python (FlashAttention). Host+device TU. #include #include #include #include +#include using namespace fk; +constexpr int HD = {self.head_dim}; + 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 float* params, const float* modParams, + const unsigned char* blockMask, int nQB, int nKB, + int maskBQ, int maskBKV, void* ext_stream) {{ + // PROLOGUES: Q/K/V are Read IOps (possibly .then-fused); every element + // enters the DPP through them, in-register, at load time. + const auto qIOp = {q_build}; + const auto kIOp = {k_build}; + const auto vIOp = {v_build}; const auto epilogue = {ep_build}; - using Ep = {ep_type}; - if (ext_stream != nullptr) {{ + const auto scoreMod = {mod_ctor}; + (void)scoreMod; (void)modParams; +{sparse_decl} 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); + {exec_fn}(qIOp, kIOp, vIOp, o, bh, seqQ, seqK, + causal != 0, stream, scale{call_tail}); }} 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); + {exec_fn}(qIOp, kIOp, vIOp, o, bh, seqQ, seqK, + causal != 0, stream, scale{call_tail}); stream.sync(); }} }} }} // extern "C" """ + mod_tok = self.score_mod.token() if self.score_mod else "none" sig = (f"flashattn;arch={_ARCH};cg={CODEGEN_VERSION};d={self.head_dim};" - f"kv={self.kv_layout};ep=" + "|".join(toks)) + f"kv={self.kv_layout};mma={int(self.mma)};mod={mod_tok};" + f"bs={int(self.block_sparse)};pq=" + "|".join(q_tk) + + ";pk=" + "|".join(k_tk) + ";pv=" + "|".join(v_tk) + + ";ep=" + "|".join(e_tk)) 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] + ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, + ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, + ctypes.c_void_p] self._lib = lib @@ -243,18 +353,48 @@ def _ensure_compiled(self, in_dtype: str): def flash_attention(q, k, v, out=None, causal=False, scale=None, stream=None, - epilogue=None, kv_layout=None): + epilogue=None, kv_layout=None, + prologue_q=None, prologue_k=None, prologue_v=None, + mma=False, score_mod=None, block_mask=None): """One-shot API. Auto-detects head_dim and kv layout; compiled kernels - are cached per (head_dim, layout, epilogue shape).""" + are cached per (head_dim, layout, mma, prologue shapes, epilogue shape). + + prologue_q/k/v: compute op chains fused onto the Q/K/V Read IOps — + every element is preprocessed in-register at load time (Oscar's + prologue: the input is a ReadOperation, not a pointer). + mma=True: tensor-core path (bf16 mma.sync m16n8k16, fp32 accum) — + ~1e-3 abs error on unit inputs (vs ~5e-7 SIMT fp32) but order-of- + magnitude faster. Same prologue/epilogue fusion.""" if kv_layout is None: - kv_layout = "int8" if isinstance(k, CompressedKV) else "dense" + kv_layout = k.fmt if isinstance(k, CompressedKV) else "dense" + if (score_mod is not None or block_mask is not None) and not mma: + mma = True # flex/sparse live on the tensor-core path vq = as_device_view(q) - ep_key = tuple(type(op).__name__ for op in (epilogue or [])) - key = (vq.width, kv_layout, ep_key) + + def _key(ops): + return tuple(type(op).__name__ for op in (ops or [])) + + key = (vq.width, kv_layout, bool(mma), _key(prologue_q), _key(prologue_k), + _key(prologue_v), _key(epilogue), + type(score_mod).__name__ if score_mod else "none", + block_mask is not None) fa = _fa_cache.get(key) if fa is None: - fa = FlashAttention(vq.width, kv_layout, epilogue) + fa = FlashAttention(vq.width, kv_layout, epilogue, + prologue_q, prologue_k, prologue_v, mma=mma, + score_mod=score_mod, + block_sparse=block_mask is not None) _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) + else: # same shapes, fresh values (never recompiles) + if epilogue: + fa.epilogue = list(epilogue) + if prologue_q: + fa.prologue_q = list(prologue_q) + if prologue_k: + fa.prologue_k = list(prologue_k) + if prologue_v: + fa.prologue_v = list(prologue_v) + if score_mod is not None: + fa.score_mod = score_mod # same TYPE (key) — fresh values + return fa(q, k, v, out=out, causal=causal, scale=scale, stream=stream, + block_mask=block_mask) diff --git a/fkl/codegen.py b/fkl/codegen.py index cf6b732..fc5ca59 100644 --- a/fkl/codegen.py +++ b/fkl/codegen.py @@ -21,7 +21,7 @@ from .types import DType # bump when generate_cu's emitted C++ changes for the SAME signature inputs -CODEGEN_VERSION = 7 +CODEGEN_VERSION = 8 def plan(ops: List[Op], in_dtype: DType, in_shape: Tuple[int, int, int], @@ -72,116 +72,47 @@ def generate_cu(ops: List[Op], in_dtype: DType, states, out_st = plan(ops, in_dtype, in_shape, n_inputs) in_st = states[0] - # emit build() expressions for everything except read/write (those are - # constructed from the IO pointers below) + # emit build() expressions for everything except read/write. The read + # and write IOps construct their OWN host-side IO objects via + # emit_read()/emit_write() (the IO is part of the IOp, mirroring how the + # C++ side carries the buffer inside the read/write IOp type) instead of + # a closed if/elif ladder here. iop_exprs = [] batch_fuse_head = False # batch ReadBack ops need explicit fuse() w/ read - fuse_with_read = None # op that consumes the read expr (BorderReader) + needs_read = None # op requiring the read spliced in (CONSTANT border) pbase = 0 + pbase_by_op = [] for op, st in zip(ops, states): if hasattr(op, "bind"): op.bind(st.dtype) - nvals = len(op.values) + pbase_by_op.append(pbase) if op.role not in (READ, WRITE): - if getattr(op, "_fuse_with_read", False): - if fuse_with_read is not None: - raise ValueError("only one read-fusing op per chain") - fuse_with_read = (op, st, pbase) + if getattr(op, "_needs_read", False): + if needs_read is not None: + raise ValueError("only one read-splicing op per chain") + needs_read = (op, st, pbase) iop_exprs.append(None) # placeholder, filled below else: iop_exprs.append(op.cpp(st, pbase)) if getattr(op, "_batch", False): batch_fuse_head = True - pbase += nvals + pbase += len(op.values) read_op, write_op = ops[0], ops[-1] - in_t, out_t = in_st.dtype.ctype, out_st.dtype.ctype MEM = "MemType::Device" if target == "gpu" else "MemType::Host" DPP = ("TransformDPP" if thread_fusion else "TransformDPP<>") - # ---- IO construction (host) ---- - if read_op.name == "ReadSet": - # constant generator: no input pointer at all - in_decl = "// ReadSet: no DRAM input" - read_expr = ops[0].cpp(in_st, 0) - elif read_op.name == "TensorPack": - # planar (CHW) input read back as packed vector pixels - base_in = in_st.dtype.base_ctype - packed_t = ops[0].out_dtype(in_st.dtype).ctype - in_decl = (f"Tensor<{base_in}> input(({base_in}*)d_in, (uint)dims->in_w, " - f"(uint)dims->in_h, (uint)dims->in_planes, " - f"(uint){ops[0]._ch}, {MEM});") - read_expr = f"TensorPack<{packed_t}>::build(input)" - elif n_inputs > 1: - # HF over a batch of images: d_in is void** (array of device ptrs). - # std::array, B> -> PerThreadRead inside FKL. - elems = ", ".join( - f"Ptr2D<{in_t}>((({in_t}**)d_in)[{i}], (uint)dims->in_w, " - f"(uint)dims->in_h, (uint)(dims->in_w * sizeof({in_t})), {MEM})" - for i in range(n_inputs)) - in_decl = (f"const std::array, {n_inputs}> " - f"input{{ {elems} }};") - read_expr = f"PerThreadRead::build(input)" - elif in_st.planes > 1: - in_decl = (f"Tensor<{in_t}> input(({in_t}*)d_in, (uint)dims->in_w, " - f"(uint)dims->in_h, (uint)dims->in_planes, 1, {MEM});") - read_expr = f"TensorRead<{in_t}>::build(input)" - else: - in_decl = (f"Ptr2D<{in_t}> input(({in_t}*)d_in, (uint)dims->in_w, " - f"(uint)dims->in_h, (uint)(dims->in_w * sizeof({in_t})), " - f"{MEM});") - read_expr = f"PerThreadRead::build(input)" - - if write_op.name == "TensorSplit": - base_t = out_st.dtype.base_ctype - # planes = batch (thread.z); color_planes = channels (split offsets) - out_decl = (f"Tensor<{base_t}> output(({base_t}*)d_out, (uint)dims->out_w, " - f"(uint)dims->out_h, (uint)dims->out_planes, " - f"(uint){out_st.dtype.channels}, {MEM});") - write_expr = f"TensorSplit<{out_t}>::build(output)" - elif write_op.name == "TensorTSplit": - base_t = out_st.dtype.base_ctype - ch = out_st.dtype.channels - # TensorT's (data, w, h, planes, cp) ctor leaves pitches at 0 (they - # are only filled by h_malloc_init on allocation). For an external - # pointer, build the RawPtr with explicit pitches and pass it - # straight to build(params). - pitch = f"(uint)(dims->out_w * sizeof({base_t}))" - out_decl = ( - f"PtrDims outDims((uint)dims->out_w, (uint)dims->out_h, " - f"(uint)dims->out_planes, (uint){ch});\n" - f" outDims.pitch = {pitch};\n" - f" outDims.plane_pitch = outDims.pitch * outDims.height;\n" - f" outDims.color_planes_pitch = outDims.plane_pitch * outDims.planes;\n" - f" RawPtr outRaw{{ ({base_t}*)d_out, outDims }};") - write_expr = f"TensorTSplit<{out_t}>::build(outRaw)" - elif write_op.name == "SplitWrite": - base_t = out_st.dtype.base_ctype - ch = out_st.dtype.channels - plane = f"(dims->out_w * dims->out_h)" - ptrs = ", ".join( - f"Ptr2D<{base_t}>((({base_t}*)d_out) + {c} * {plane}, " - f"(uint)dims->out_w, (uint)dims->out_h, " - f"(uint)(dims->out_w * sizeof({base_t})), {MEM})" - for c in range(ch)) - out_decl = f"const std::vector> output{{ {ptrs} }};" - write_expr = f"SplitWrite::build(output)" - elif out_st.planes > 1: - out_decl = (f"Tensor<{out_t}> output(({out_t}*)d_out, (uint)dims->out_w, " - f"(uint)dims->out_h, (uint)dims->out_planes, 1, {MEM});") - write_expr = f"TensorWrite<{out_t}>::build(output)" - else: - out_decl = (f"Ptr2D<{out_t}> output(({out_t}*)d_out, (uint)dims->out_w, " - f"(uint)dims->out_h, (uint)(dims->out_w * sizeof({out_t})), " - f"{MEM});") - write_expr = f"PerThreadWrite::build(output)" + # ---- IO construction (host): delegated to the read/write IOps ---- + in_decl, read_expr = read_op.emit_read(in_st, MEM, n_inputs) + out_decl, write_expr = write_op.emit_write(out_st, MEM, pbase_by_op[-1]) - if fuse_with_read is not None: - # Per Oscar: the BorderReader takes the Image/Ptr2D read IOp as its - # backIOp: BorderReader::build(readIOp[, value]). The combined - # expression REPLACES the read at the head of the chain. - op_f, st_f, pb_f = fuse_with_read + if needs_read is not None: + # CONSTANT BorderReader: FKL's incomplete-const builder doesn't + # compile, so it needs the COMPLETE form build(readIOp, value). The + # combined expression REPLACES the read at the head of the chain. + # (All other border modes fuse via C++ fuse_back — no splice.) + op_f, st_f, pb_f = needs_read read_expr = op_f.cpp_with_read(st_f, pb_f, read_expr) iop_exprs = [e for e in iop_exprs if e is not None] diff --git a/fkl/flex.py b/fkl/flex.py new file mode 100644 index 0000000..1a4cb8e --- /dev/null +++ b/fkl/flex.py @@ -0,0 +1,143 @@ +"""FlexAttention-style score mods and block-sparse masks for fkl. + +SCORE MODS (the PyTorch FlexAttention idea, fused at COMPILE TIME): +a score mod transforms each attention score after scaling and +causal/bounds masking, before the online softmax: + + s' = mod(s, q_idx, kv_idx) + +Built-ins map 1:1 to the C++ functors in flash_attention_mma.h: + + fkl.flash_attention(q, k, v, mma=True, score_mod=fkl.ALiBi(slope=0.0625)) + fkl.flash_attention(q, k, v, mma=True, score_mod=fkl.SoftCap(20.0)) + fkl.flash_attention(q, k, v, mma=True, score_mod=fkl.SlidingWindow(256)) + +Mod VALUES (slope, cap, window) are runtime kernel arguments — changing +them never recompiles. Only the mod TYPE is part of the compile key. + +BLOCK SPARSITY: BlockMask wraps a (bh, nQB, nKB) uint8 numpy array (or +GPU buffer); inactive tiles are skipped entirely (loads AND math): + + mask = fkl.BlockMask(dense_mask_uint8, block_q=128, block_kv=128) + fkl.flash_attention(q, k, v, mma=True, block_mask=mask) +""" + +from __future__ import annotations + +import ctypes + +# NOTE: numpy is imported lazily inside the methods that need it. The core +# package contract is "import fkl works with zero python deps" (see +# pyproject: dependencies = []); numpy is only required if you actually +# construct a BlockMask from a host array. +from .tensor import DeviceBuffer, as_device_view + + +# --------------------------------------------------------------------------- +# score mods +# --------------------------------------------------------------------------- +class ScoreMod: + """Base. Subclasses define: name (C++ functor), values (runtime params).""" + name = "NoScoreMod" + cpp_ctor = "NoScoreMod{}" + + @property + def values(self): + return [] + + def token(self) -> str: + return self.name + + +class ALiBi(ScoreMod): + """s - slope * (q_idx - kv_idx). One slope per kernel call; for + per-head slopes launch per-head or use the C++ API directly.""" + name = "ALiBiScoreMod" + + def __init__(self, slope: float): + self.slope = float(slope) + + @property + def cpp_ctor(self): + return "ALiBiScoreMod{ modParams[0] }" + + @property + def values(self): + return [self.slope] + + +class SoftCap(ScoreMod): + """Gemma-2 style logit soft capping: cap * tanh(s / cap).""" + name = "SoftCapScoreMod" + + def __init__(self, cap: float): + self.cap = float(cap) + + @property + def cpp_ctor(self): + return "SoftCapScoreMod{ modParams[0] }" + + @property + def values(self): + return [self.cap] + + +class SlidingWindow(ScoreMod): + """mask_mod: keep only the last `window` keys for each query.""" + name = "SlidingWindowMask" + + def __init__(self, window: int): + self.window = int(window) + + @property + def cpp_ctor(self): + return "SlidingWindowMask{ (int)modParams[0] }" + + @property + def values(self): + return [float(self.window)] + + +# --------------------------------------------------------------------------- +# block sparsity +# --------------------------------------------------------------------------- +class BlockMask: + """Block-sparse attention mask at (block_q x block_kv) granularity. + + mask: numpy uint8/bool array (bh, nQB, nKB) — 1 = attend, 0 = skip — + or anything with __cuda_array_interface__ already on the GPU. + Inactive tiles skip global loads AND tensor-core math; runtime scales + with density. block_q/block_kv must be multiples of the kernel tiles + (128 is always safe for the defaults).""" + + def __init__(self, mask, block_q: int = 128, block_kv: int = 128): + import numpy as np + if isinstance(mask, np.ndarray): + if mask.ndim != 3: + raise ValueError("mask must be (bh, nQBlocks, nKVBlocks)") + m = np.ascontiguousarray(mask.astype(np.uint8)) + self.bh, self.n_q_blocks, self.n_kv_blocks = m.shape + self._buf = DeviceBuffer(self.n_kv_blocks, self.n_q_blocks, + "uint8", planes=self.bh) + self._buf.copy_from_host(m.tobytes()) + self.ptr = as_device_view(self._buf).ptr + else: + v = as_device_view(mask) + self._buf = mask # keep alive + self.bh, self.n_q_blocks, self.n_kv_blocks = v.planes, v.height, v.width + self.ptr = v.ptr + self.block_q = int(block_q) + self.block_kv = int(block_kv) + + @staticmethod + def causal(bh: int, seq: int, block: int = 128) -> "BlockMask": + """Convenience: block-causal mask (tile fully above diagonal -> skip).""" + import numpy as np + n = (seq + block - 1) // block + m = np.zeros((bh, n, n), dtype=np.uint8) + for qb in range(n): + m[:, qb, :qb + 1] = 1 + return BlockMask(m, block, block) + + def density(self) -> float: + return -1.0 # unknown without download; informational only diff --git a/fkl/jit.py b/fkl/jit.py index 59e9583..38779a3 100644 --- a/fkl/jit.py +++ b/fkl/jit.py @@ -269,12 +269,15 @@ def __init__(self, plane_map, chains): self._variants = {} def _selector_cpp(self): - # FKL convention (see SequenceSelectorType in circular_tensor.h): - # at(z) is uint, 1-based sequence index. Generate a chain of ternaries. - terms = [] - for z, s in enumerate(self.plane_map): - terms.append((z, s)) - expr = f"{self.plane_map[-1]}u" + # DivergentBatchTransformDPP is 0-BASED: exec() calls + # divergent_operate<0>(z, seqs...) and runs the sequence whose + # 0-based position matches at(z) (see data_parallel_patterns.h; the + # upstream regression test uses at(index)=index==0?0u:1u, i.e. 0 picks + # the FIRST sequence). plane_map is given 1-based at the API for + # readability, so emit (s-1) here. (NOTE: this is a DIFFERENT selector + # contract than circular_tensor.h's SequenceSelectorType.) + terms = [(z, s - 1) for z, s in enumerate(self.plane_map)] + expr = f"{terms[-1][1]}u" for z, s in reversed(terms[:-1]): expr = f"(index == {z}u ? {s}u : {expr})" return ( @@ -384,7 +387,9 @@ def _get_variant(self, dt, shape, B): "|".join(op.token(st) for op, st in zip(ch, _plan(ch, dt, shape, B)[0])) for ch in self.chains) - sig = (f"divergent;arch={_ARCH};in={dt}x{B};map={tuple(self.plane_map)};" + # sv (selector version): bump when _selector_cpp's emitted C++ changes + # for the same plane_map, so stale .so files are not reused. + sig = (f"divergent;arch={_ARCH};sv=2;in={dt}x{B};map={tuple(self.plane_map)};" f"chains={chain_toks}") so = get_backend().compile(src, sig) lib = ctypes.CDLL(str(so)) diff --git a/fkl/operations.py b/fkl/operations.py index bc73468..f3f4b72 100644 --- a/fkl/operations.py +++ b/fkl/operations.py @@ -366,21 +366,58 @@ def token(self, state): # ===================== memory ops ========================================= class _ReadOp(Op): + """A chain-head memory Read IOp. + + Owns its own host-side IO construction via emit_read(): it returns the + C++ declaration of the input buffer object plus the build() expression + that consumes it. Keeping this on the descriptor (instead of a closed + if/elif ladder in codegen) means a Read IOp is self-contained — exactly + how the C++ side carries the read in its IOp type. emit_read is the read + counterpart to cpp() for COMPUTE ops. + """ role = READ + def emit_read(self, state, mem, n_inputs): + """-> (in_decl, read_expr). in_decl declares the host buffer object + from d_in/dims; read_expr is the FKL build() that consumes it.""" + raise NotImplementedError(f"{self.name} does not implement emit_read") + class _WriteOp(Op): + """A chain-tail memory Write IOp. emit_write() is the write counterpart + to _ReadOp.emit_read().""" role = WRITE + def emit_write(self, state, mem, pbase): + """-> (out_decl, write_expr). out_decl declares the host output + buffer object from d_out/dims; write_expr is the build() consuming it.""" + raise NotImplementedError(f"{self.name} does not implement emit_write") + class TensorRead(_ReadOp): - """Read input. 2D Ptr2D -> PerThreadRead<_2D>; 3D Tensor -> TensorRead.""" + """Read input. 2D Ptr2D -> PerThreadRead<_2D>; 3D Tensor -> TensorRead; + batch of B same-size 2D images -> std::array -> PerThreadRead + (BatchRead under the hood, the HF read).""" name = "TensorRead" - def cpp(self, state, pbase): + def emit_read(self, state, mem, n_inputs): + t = state.dtype.ctype + if n_inputs > 1: + # HF over a batch of images: d_in is void** (array of device ptrs). + elems = ", ".join( + f"Ptr2D<{t}>((({t}**)d_in)[{i}], (uint)dims->in_w, " + f"(uint)dims->in_h, (uint)(dims->in_w * sizeof({t})), {mem})" + for i in range(n_inputs)) + decl = (f"const std::array, {n_inputs}> " + f"input{{ {elems} }};") + return decl, f"PerThreadRead::build(input)" if state.planes > 1: - return f"TensorRead<{state.dtype}>::build(in_tensor)" - return f"PerThreadRead::build(input.ptr())" + decl = (f"Tensor<{t}> input(({t}*)d_in, (uint)dims->in_w, " + f"(uint)dims->in_h, (uint)dims->in_planes, 1, {mem});") + return decl, f"TensorRead<{t}>::build(input)" + decl = (f"Ptr2D<{t}> input(({t}*)d_in, (uint)dims->in_w, " + f"(uint)dims->in_h, (uint)(dims->in_w * sizeof({t})), {mem});") + return decl, f"PerThreadRead::build(input)" def token(self, state): return f"Read<{state.dtype},p{state.planes}>" @@ -389,10 +426,15 @@ def token(self, state): class TensorWrite(_WriteOp): name = "TensorWrite" - def cpp(self, state, pbase): + def emit_write(self, state, mem, pbase): + t = state.dtype.ctype if state.planes > 1: - return f"TensorWrite<{state.dtype}>::build(out_tensor)" - return f"PerThreadWrite::build(output.ptr())" + decl = (f"Tensor<{t}> output(({t}*)d_out, (uint)dims->out_w, " + f"(uint)dims->out_h, (uint)dims->out_planes, 1, {mem});") + return decl, f"TensorWrite<{t}>::build(output)" + decl = (f"Ptr2D<{t}> output(({t}*)d_out, (uint)dims->out_w, " + f"(uint)dims->out_h, (uint)(dims->out_w * sizeof({t})), {mem});") + return decl, f"PerThreadWrite::build(output)" def token(self, state): return f"Write<{state.dtype},p{state.planes}>" @@ -400,14 +442,19 @@ def token(self, state): class TensorSplit(_WriteOp): """Write vector pixels as separate planes (planar layout for DNNs). - uchar3 HxW -> 3 planes of HxW uchar.""" + uchar3 HxW -> 3 planes of HxW uchar. planes = batch (thread.z); + color_planes = channels (split offsets).""" name = "TensorSplit" def out_shape(self, shape): return shape # planes encoded by channels at alloc time - def cpp(self, state, pbase): - return f"TensorSplit<{state.dtype}>::build(out_tensor)" + def emit_write(self, state, mem, pbase): + base_t = state.dtype.base_ctype + decl = (f"Tensor<{base_t}> output(({base_t}*)d_out, (uint)dims->out_w, " + f"(uint)dims->out_h, (uint)dims->out_planes, " + f"(uint){state.dtype.channels}, {mem});") + return decl, f"TensorSplit<{state.dtype}>::build(output)" def token(self, state): return f"TensorSplit<{state.dtype}>" @@ -552,8 +599,17 @@ class SplitWrite(_WriteOp): contiguously (like TensorSplit but via SplitWrite params).""" name = "SplitWrite" - def cpp(self, state, pbase): - return None # constructed by codegen from output pointer + def emit_write(self, state, mem, pbase): + base_t = state.dtype.base_ctype + ch = state.dtype.channels + plane = "(dims->out_w * dims->out_h)" + ptrs = ", ".join( + f"Ptr2D<{base_t}>((({base_t}*)d_out) + {c} * {plane}, " + f"(uint)dims->out_w, (uint)dims->out_h, " + f"(uint)(dims->out_w * sizeof({base_t})), {mem})" + for c in range(ch)) + decl = f"const std::vector> output{{ {ptrs} }};" + return decl, f"SplitWrite::build(output)" def token(self, state): return f"SplitWrite<{state.dtype}>" @@ -620,6 +676,11 @@ def cpp(self, state, pbase): f"(uint)params[{pbase + n + 1}], 1u)") return f"ReadSet<{self._dt}>::build({val}, {at})" + def emit_read(self, state, mem, n_inputs): + # constant generator: no input pointer at all. Read params live at + # the head of params[] (pbase 0), like any chain-head op. + return "// ReadSet: no DRAM input", self.cpp(state, 0) + def token(self, state): return f"ReadSet<{self._dt}>" @@ -637,8 +698,13 @@ def __init__(self, channels: int): def out_dtype(self, dt): return dt.with_channels(self._ch) - def cpp(self, state, pbase): - return None # built by codegen from the input pointer + def emit_read(self, state, mem, n_inputs): + base_in = state.dtype.base_ctype + packed_t = self.out_dtype(state.dtype).ctype + decl = (f"Tensor<{base_in}> input(({base_in}*)d_in, (uint)dims->in_w, " + f"(uint)dims->in_h, (uint)dims->in_planes, " + f"(uint){self._ch}, {mem});") + return decl, f"TensorPack<{packed_t}>::build(input)" def token(self, state): return f"TensorPack<{state.dtype}c{self._ch}>" @@ -649,8 +715,22 @@ class TensorTSplit(_WriteOp): layout [C][planes][H][W] instead of TensorSplit's [planes][C][H][W].""" name = "TensorTSplit" - def cpp(self, state, pbase): - return None # built by codegen from the output pointer + def emit_write(self, state, mem, pbase): + base_t = state.dtype.base_ctype + ch = state.dtype.channels + # TensorT's (data, w, h, planes, cp) ctor leaves pitches at 0 (they + # are only filled by h_malloc_init on allocation). For an external + # pointer, build the RawPtr with explicit pitches and pass it + # straight to build(params). + pitch = f"(uint)(dims->out_w * sizeof({base_t}))" + decl = ( + f"PtrDims outDims((uint)dims->out_w, (uint)dims->out_h, " + f"(uint)dims->out_planes, (uint){ch});\n" + f" outDims.pitch = {pitch};\n" + f" outDims.plane_pitch = outDims.pitch * outDims.height;\n" + f" outDims.color_planes_pitch = outDims.plane_pitch * outDims.planes;\n" + f" RawPtr outRaw{{ ({base_t}*)d_out, outDims }};") + return decl, f"TensorTSplit<{state.dtype}>::build(outRaw)" def token(self, state): return f"TensorTSplit<{state.dtype}>" @@ -660,12 +740,28 @@ def token(self, state): class BorderReader(Op): - """Out-of-bounds policy fused onto the read (BVF). Per Oscar: the input - is an Image/Ptr2D read, and BorderReader takes that complete Read IOp as - its backIOp directly: BorderReader::build(readIOp[, defaultValue]). - Codegen consumes the read expression (marker: _fuse_with_read).""" + """Out-of-bounds policy, fused backwards onto the read (BVF). + + Value-less modes (replicate/reflect/wrap/reflect101) are emitted as a + plain IncompleteReadBack IOp: + BorderReader::build() + FKL's BackFuser::fuse_back (called inside executeOperations) detects the + IncompleteReadBackType and fuses it with the preceding Read IOp via + fk::fuse -> BorderReader::build(readIOp, selfIOp). Python no longer + splices the read in for these — that duplicated the C++ fusion. (Oscar + review, point 2.) + + CONSTANT is the exception: FKL's incomplete builder + BorderReader::build(value) does not compile (border_reader.h + passes NullType where a backIOp is required), so the library only + supports the COMPLETE form build(readIOp, value). For that one mode we + still hand the read expression in (marker: _needs_read). All other modes + take the clean fuse_back path. + """ name = "BorderReader" - _fuse_with_read = True + + _BT = {"constant": "CONSTANT", "replicate": "REPLICATE", + "reflect": "REFLECT", "wrap": "WRAP", "reflect101": "REFLECT_101"} def __init__(self, border: str = "replicate", value=0.0): b = border.lower() @@ -673,6 +769,9 @@ def __init__(self, border: str = "replicate", value=0.0): raise ValueError(f"border must be one of {_BORDER_TYPES}") self._b = b self._value = value + # only CONSTANT needs the read spliced in (FKL incomplete-const is + # unsupported); other modes fuse via C++ fuse_back. + self._needs_read = (b == "constant") @property def values(self): @@ -684,18 +783,18 @@ def bind(self, dt): if self._b == "constant": self._vals = pack_value(dt, self._value) - def cpp_with_read(self, state, pbase, read_expr: str) -> str: - bt = {"constant": "CONSTANT", "replicate": "REPLICATE", - "reflect": "REFLECT", "wrap": "WRAP", - "reflect101": "REFLECT_101"}[self._b] - if self._b == "constant": - self.bind(state.dtype) - return (f"BorderReader::build(" - f"{read_expr}, {make_expr(state.dtype, pbase)})") - return f"BorderReader::build({read_expr})" - def cpp(self, state, pbase): - raise RuntimeError("BorderReader must be emitted via cpp_with_read") + # value-less borders: incomplete IOp, C++ fuse_back supplies the read. + if self._b != "constant": + return f"BorderReader::build()" + raise RuntimeError("CONSTANT BorderReader must be emitted via " + "cpp_with_read (FKL incomplete-const unsupported)") + + def cpp_with_read(self, state, pbase, read_expr: str) -> str: + # complete form: BorderReader::build(readIOp, value) + self.bind(state.dtype) + return (f"BorderReader::build(" + f"{read_expr}, {make_expr(state.dtype, pbase)})") def token(self, state): return f"BorderReader<{self._b},{state.dtype}>" diff --git a/fkl/tensor.py b/fkl/tensor.py index 19a1767..e028029 100644 --- a/fkl/tensor.py +++ b/fkl/tensor.py @@ -8,7 +8,7 @@ from dataclasses import dataclass import ctypes -from .types import DType, from_cai, _BASES +from .types import DType, from_cai, from_shape, _BASES @dataclass @@ -23,6 +23,22 @@ class DeviceView: def as_device_view(x) -> "DeviceView": + # FAST PATH: torch.Tensor. Avoids __cuda_array_interface__, which + # builds a fresh dict per access (~2.5us each — dominated the per-call + # overhead of fkl.flash_attention; measured 28us -> ~7us with this). + tt = _torch_tensor_type() + if tt is not None and type(x) is tt: + if not x.is_cuda: + raise TypeError("torch tensor must be on CUDA") + if not x.is_contiguous(): + raise ValueError("non-contiguous arrays not supported; call .contiguous() first") + base = _TORCH_DTYPE_TO_BASE.get(x.dtype) + if base is None: + raise TypeError(f"unsupported torch dtype {x.dtype}") + dt, w, h, p = from_shape(tuple(x.shape), base) + dev = x.device.index + return DeviceView(x.data_ptr(), w, h, p, dt, x, 0 if dev is None else dev) + cai = getattr(x, "__cuda_array_interface__", None) if cai is None: raise TypeError( @@ -35,6 +51,29 @@ def as_device_view(x) -> "DeviceView": return DeviceView(int(cai["data"][0]), w, h, p, dt, x, _device_of(x)) +_torch_tensor_cls = False # False = not probed yet; None = torch absent +_TORCH_DTYPE_TO_BASE = {} + + +def _torch_tensor_type(): + global _torch_tensor_cls + if _torch_tensor_cls is False: + import sys + torch = sys.modules.get("torch") # never import torch ourselves + if torch is None: + _torch_tensor_cls = None + else: + _torch_tensor_cls = torch.Tensor + _TORCH_DTYPE_TO_BASE.update({ + torch.float32: "float32", torch.float16: "float16", + torch.bfloat16: "bfloat16", torch.uint8: "uint8", + torch.int8: "int8", torch.int16: "int16", + torch.int32: "int32", torch.int64: "int64", + torch.float64: "float64", + }) + return _torch_tensor_cls + + def _device_of(x) -> int: """Best-effort device index: DeviceBuffer.device, torch .device.index, cupy .device.id; default 0 (CAI itself carries no device id).""" diff --git a/fkl/types.py b/fkl/types.py index d15de79..b6955cf 100644 --- a/fkl/types.py +++ b/fkl/types.py @@ -74,14 +74,9 @@ def dtype(spec) -> DType: return DType(spec, 1) -def from_cai(shape, typestr) -> tuple["DType", int, int, int]: - """(shape, typestr) -> (DType, width, height, planes). - - Trailing dim of size 2..4 is folded into a vector pixel type. - """ - base = _TYPESTR_TO_BASE.get(typestr) - if base is None: - raise TypeError(f"unsupported typestr {typestr!r}") +def from_shape(shape, base) -> tuple["DType", int, int, int]: + """(shape, base name) -> (DType, width, height, planes). Same folding + rules as from_cai but without typestr parsing (fast path).""" ch = 1 dims = list(shape) if len(dims) >= 2 and 2 <= dims[-1] <= 4: @@ -99,6 +94,17 @@ def from_cai(shape, typestr) -> tuple["DType", int, int, int]: return DType(base, ch), int(w), int(h), int(p) +def from_cai(shape, typestr) -> tuple["DType", int, int, int]: + """(shape, typestr) -> (DType, width, height, planes). + + Trailing dim of size 2..4 is folded into a vector pixel type. + """ + base = _TYPESTR_TO_BASE.get(typestr) + if base is None: + raise TypeError(f"unsupported typestr {typestr!r}") + return from_shape(shape, base) + + # constants for value packing in codegen def pack_value(dt: DType, value) -> list[float]: """Normalize a python scalar/tuple to dt.channels floats for params[].""" diff --git a/skills/fkl-python-extending/SKILL.md b/skills/fkl-python-extending/SKILL.md index a39121d..9173835 100644 --- a/skills/fkl-python-extending/SKILL.md +++ b/skills/fkl-python-extending/SKILL.md @@ -8,8 +8,20 @@ description: Add a new FKL operation to the fkl-python front-end (operations.py ## Architecture recap - `fkl/operations.py` - one Python descriptor class per C++ Operation. + - COMPUTE ops implement `cpp(state, pbase)` (the build() expression). + - Read ops (subclass `_ReadOp`) implement `emit_read(state, mem, n_inputs)` + -> `(in_decl, read_expr)`: they construct their OWN host-side input + buffer object (Ptr2D / Tensor / std::array / ...) AND the + build() that consumes it. A Read IOp is self-contained. + - Write ops (subclass `_WriteOp`) implement `emit_write(state, mem, pbase)` + -> `(out_decl, write_expr)`, the write counterpart. - `fkl/codegen.py` - threads ChainState (dtype+shape) through the chain and emits ONE host+device .cu with extern "C" fkl_entry. + generate_cu() just calls read_op.emit_read() / + write_op.emit_write() — there is NO if/elif ladder on + op .name. Add a new IO layout by adding a Read/Write + descriptor with emit_read/emit_write, not by editing + codegen. - `fkl/backend.py` - clang/nvcc single-step compile + disk cache. - `fkl/jit.py` - compose() + FusedKernel hot path (ctypes). @@ -86,6 +98,16 @@ def t_myop(): whose InstanceType is ReadType -> BackFuser does NOT auto-fuse them with the previous read. Codegen wraps them as `fuse(read_expr, batch_expr)`. This is keyed off the op having a truthy `_batch` attribute. +2b. ReadBack ops fuse for free via C++. executeOperations() internally calls + BackFuser::fuse_back(iOps...) (executors.h), and a value-less + IncompleteReadBack op (e.g. BorderReader::build()) placed as a + plain IOp in the flat list gets fused with the preceding Read automatically + — do NOT splice the read in from Python. EXCEPTION: BorderReader + is broken upstream — the incomplete builder build(value) does not compile + (border_reader.h:72 passes NullType where a backIOp is required), so it + needs the COMPLETE form build(readIOp, value). That one mode uses the + `_needs_read` marker so codegen splices the read expression in; every other + border mode takes the clean fuse_back path. 3. FKL's LTS-C++17 branch is the supported base (it carries the fix for issue #244, ColorConversion FusedOperation aliases). On stale `main` the BGR2GRAY-family aliases fail to compile (see #244/#248). @@ -102,7 +124,14 @@ def t_myop(): over a const tuple -> rvalue-ref binding error (upstream bug). The divergent codegen launches launchDivergentBatchTransformDPP_Kernel directly with buildOperationSequence(...) lvalues and a generated - PySequenceSelector (1-based at(z), FKL convention). + PySequenceSelector. The selector is 0-BASED: exec() calls + divergent_operate<0>(z, seqs...) and runs the sequence at the 0-based + position equal to at(z) (data_parallel_patterns.h; upstream regression + test MySelector::at returns index==0?0u:1u). compose_divergent takes a + 1-based plane_map at the API for readability, so _selector_cpp emits + (s-1). Do NOT confuse this with circular_tensor.h's SequenceSelectorType + (a different contract). Bump the `sv=` token in the divergent signature + when the emitted selector C++ changes, or stale .so files get reused. 8. IOpSequences passed to the divergent kernel must be LVALUES (const auto seqN = buildOperationSequence(...)), not temporaries inlined in the launch expression. diff --git a/tests/test_flash_attention.py b/tests/test_flash_attention.py index 72c26c1..6baff86 100644 --- a/tests/test_flash_attention.py +++ b/tests/test_flash_attention.py @@ -137,7 +137,218 @@ def t_single_query_decode(): check_true(f"FA decode step s_q=1 vs s_k=256 int8 (err={err:.2e})", err < 2e-2) +def t_prologue_q(): + """Q prologue Mul(2): equals oracle on 2*Q (fused in the Read IOp).""" + q, k, v = _mk(2, 24, 48, 32, 7) + got = _from_gpu( + fkl.flash_attention(_to_gpu(q), _to_gpu(k), _to_gpu(v), + prologue_q=[fkl.Mul(2.0)]), q.shape) + ref = _oracle(q * 2.0, k, v, False, 1 / math.sqrt(32)) + err = np.abs(got - ref).max() + check_true(f"FA Q-prologue Mul(2) == oracle(2Q) (err={err:.2e})", err < 5e-6) + + +def t_prologue_v_affine(): + """V prologue Mul(3).then(Add(1)) => 3*out + 1 (since sum p_j = 1).""" + q, k, v = _mk(2, 24, 48, 32, 8) + got = _from_gpu( + fkl.flash_attention(_to_gpu(q), _to_gpu(k), _to_gpu(v), + prologue_v=[fkl.Mul(3.0), fkl.Add(1.0)]), q.shape) + ref = _oracle(q, k, v, False, 1 / math.sqrt(32)) + err = np.abs(got - (3.0 * ref + 1.0)).max() + check_true(f"FA V-prologue Mul(3).Add(1) == 3*out+1 (err={err:.2e})", err < 5e-6) + + +def t_prologue_int8_kv(): + """Prologue chains AFTER int8 dequant: V int8 + Mul(2) => 2*out_int8.""" + q, k, v = _mk(2, 16, 64, 32, 9) + kc, vc = fkl.compress_kv(_to_gpu(k)), fkl.compress_kv(_to_gpu(v)) + base = _from_gpu(fkl.flash_attention(_to_gpu(q), kc, vc), q.shape) + got = _from_gpu(fkl.flash_attention(_to_gpu(q), kc, vc, + prologue_v=[fkl.Mul(2.0)]), q.shape) + err = np.abs(got - 2.0 * base).max() + check_true(f"FA int8-KV + V-prologue Mul(2) == 2*base (err={err:.2e})", + err < 1e-5) + + +def t_prologue_values_no_recompile(): + """Changing prologue values reuses the cached kernel.""" + import time + q, k, v = _mk(1, 16, 16, 32, 10) + _ = fkl.flash_attention(_to_gpu(q), _to_gpu(k), _to_gpu(v), + prologue_q=[fkl.Mul(2.0)]) # compiles + t0 = time.perf_counter() + a = _from_gpu(fkl.flash_attention(_to_gpu(q), _to_gpu(k), _to_gpu(v), + prologue_q=[fkl.Mul(4.0)]), q.shape) + dt = time.perf_counter() - t0 + ref = _oracle(q * 4.0, k, v, False, 1 / math.sqrt(32)) + ok = np.abs(a - ref).max() < 5e-6 and dt < 1.0 + check_true(f"FA prologue values change without recompile ({dt*1e3:.0f}ms)", ok) + + +def t_mma_dense_causal(): + """Tensor-core path (bf16 mma): same oracle, bf16-class tolerance.""" + q, k, v = _mk(2, 128, 128, 64, 11) + out = fkl.flash_attention(_to_gpu(q), _to_gpu(k), _to_gpu(v), + causal=True, mma=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 mma d64 causal (err={err:.2e})", err < 2e-2) + + +def t_mma_prologue_epilogue(): + """mma path keeps the FKL superpowers: V-prologue + epilogue fused.""" + q, k, v = _mk(2, 64, 128, 64, 12) + got = _from_gpu( + fkl.flash_attention(_to_gpu(q), _to_gpu(k), _to_gpu(v), mma=True, + prologue_v=[fkl.Mul(3.0), fkl.Add(1.0)], + epilogue=[fkl.Mul(2.0)]), q.shape) + ref = _oracle(q, k, v, False, 1 / math.sqrt(64)) + err = np.abs(got - 2.0 * (3.0 * ref + 1.0)).max() + check_true(f"FA mma V-prologue+epilogue == 2*(3*out+1) (err={err:.2e})", + err < 5e-2) + + +def t_fp8_kv(): + """fp8 e4m3 KV cache: exact vs dequantized oracle is impractical from + python (needs fp8 decode), so verify quant-bounded e2e accuracy and the + memory ratio; mma path exercises the QUANT_KV cp.async schedule.""" + q, k, v = _mk(2, 48, 96, 64, 20) + kc = fkl.compress_kv(_to_gpu(k), fmt="fp8") + vc = fkl.compress_kv(_to_gpu(v), fmt="fp8") + ratio = (k.size * 4) / kc.nbytes + check_true(f"FP8 KV compression ratio {ratio:.2f}x (>3.5x)", ratio > 3.5) + + got = _from_gpu(fkl.flash_attention(_to_gpu(q), kc, vc, causal=True), q.shape) + ref = _oracle(q, k, v, True, 1 / math.sqrt(64)) + err = np.abs(got - ref).max() + # e4m3 has 3 mantissa bits (~6% rel err); on UNIFORM data int8's 127 + # uniform levels are tighter. fp8 wins on outlier-heavy real caches. + check_true(f"FA fp8-KV e2e quant-bounded (err={err:.2e})", err < 6e-2) + + got2 = _from_gpu(fkl.flash_attention(_to_gpu(q), kc, vc, causal=True, + mma=True), q.shape) + err2 = np.abs(got2 - ref).max() + check_true(f"FA fp8-KV mma path (err={err2:.2e})", err2 < 6e-2) + + +def _oracle_mod(q, k, v, causal, scale, mod): + """fp64 reference with score transform mod(s, q_idx, kv_idx).""" + q, k, v = q.astype(np.float64), k.astype(np.float64), v.astype(np.float64) + bh, sq, d = q.shape + sk = k.shape[1] + qi = np.arange(sq)[:, None] + ki = np.arange(sk)[None, :] + out = np.zeros((bh, sq, d)) + for b in range(bh): + s = q[b] @ k[b].T * scale + s = mod(s, qi, ki) + if causal: + s = np.where(ki > qi, -np.inf, s) + 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 t_flex_alibi(): + q, k, v = _mk(2, 128, 128, 64, 30) + got = _from_gpu(fkl.flash_attention(_to_gpu(q), _to_gpu(k), _to_gpu(v), + causal=True, + score_mod=fkl.ALiBi(0.0625)), q.shape) + ref = _oracle_mod(q, k, v, True, 1 / math.sqrt(64), + lambda s, qi, ki: s - 0.0625 * (qi - ki)) + err = np.abs(got - ref).max() + check_true(f"flex ALiBi causal (err={err:.2e})", err < 5e-3) + + +def t_flex_softcap(): + q, k, v = _mk(2, 128, 128, 64, 31) + got = _from_gpu(fkl.flash_attention(_to_gpu(q), _to_gpu(k), _to_gpu(v), + causal=True, + score_mod=fkl.SoftCap(20.0)), q.shape) + ref = _oracle_mod(q, k, v, True, 1 / math.sqrt(64), + lambda s, qi, ki: 20.0 * np.tanh(s / 20.0)) + err = np.abs(got - ref).max() + check_true(f"flex SoftCap(20) causal (err={err:.2e})", err < 5e-3) + + +def t_flex_sliding_window(): + q, k, v = _mk(2, 128, 128, 64, 32) + got = _from_gpu(fkl.flash_attention(_to_gpu(q), _to_gpu(k), _to_gpu(v), + causal=True, + score_mod=fkl.SlidingWindow(32)), + q.shape) + ref = _oracle_mod(q, k, v, True, 1 / math.sqrt(64), + lambda s, qi, ki: np.where(qi - ki >= 32, -np.inf, s)) + err = np.abs(got - ref).max() + check_true(f"flex SlidingWindow(32) causal (err={err:.2e})", err < 5e-3) + + +def t_flex_values_no_recompile(): + """Same mod TYPE with new value -> same compiled kernel, fresh result.""" + q, k, v = _mk(1, 128, 128, 64, 33) + g1 = _from_gpu(fkl.flash_attention(_to_gpu(q), _to_gpu(k), _to_gpu(v), + score_mod=fkl.SoftCap(20.0)), q.shape) + g2 = _from_gpu(fkl.flash_attention(_to_gpu(q), _to_gpu(k), _to_gpu(v), + score_mod=fkl.SoftCap(5.0)), q.shape) + r2 = _oracle_mod(q, k, v, False, 1 / math.sqrt(64), + lambda s, qi, ki: 5.0 * np.tanh(s / 5.0)) + diff = np.abs(g1 - g2).max() + err = np.abs(g2 - r2).max() + check_true(f"flex value swap no-recompile (diff={diff:.2e}, err={err:.2e})", + diff > 1e-4 and err < 5e-3) + + +def t_block_sparse(): + bh, s, d, blk = 2, 512, 64, 128 + q, k, v = _mk(bh, s, s, d, 34) + n = s // blk + rng = np.random.default_rng(35) + mask = (rng.uniform(size=(bh, n, n)) < 0.5).astype(np.uint8) + mask[:, :, 0] = 1 # guarantee a live block per row + bm = fkl.BlockMask(mask, blk, blk) + got = _from_gpu(fkl.flash_attention(_to_gpu(q), _to_gpu(k), _to_gpu(v), + block_mask=bm), q.shape) + + def mod(srow, qi, ki): + keep = mask_b[qi // blk, ki // blk].astype(bool) + return np.where(keep, srow, -np.inf) + + ref = np.zeros_like(q, dtype=np.float64) + for b in range(bh): + mask_b = mask[b] + ref[b] = _oracle_mod(q[b:b+1], k[b:b+1], v[b:b+1], False, + 1 / math.sqrt(d), mod)[0] + err = np.abs(got - ref).max() + check_true(f"block-sparse 50% d64 s512 (err={err:.2e})", err < 5e-3) + + +def t_block_sparse_causal_helper(): + bh, s, d, blk = 2, 256, 64, 128 + q, k, v = _mk(bh, s, s, d, 36) + bm = fkl.BlockMask.causal(bh, s, blk) + # block-causal == full attention inside diagonal blocks, none above + got = _from_gpu(fkl.flash_attention(_to_gpu(q), _to_gpu(k), _to_gpu(v), + block_mask=bm), q.shape) + + def mod(srow, qi, ki): + return np.where((ki // blk) > (qi // blk), -np.inf, srow) + + ref = _oracle_mod(q, k, v, False, 1 / math.sqrt(d), mod) + err = np.abs(got - ref).max() + check_true(f"BlockMask.causal helper (err={err:.2e})", err < 5e-3) + + 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") + t_single_query_decode, t_prologue_q, t_prologue_v_affine, + t_prologue_int8_kv, t_prologue_values_no_recompile, + t_mma_dense_causal, t_mma_prologue_epilogue, t_fp8_kv, + t_flex_alibi, t_flex_softcap, t_flex_sliding_window, + t_flex_values_no_recompile, t_block_sparse, + t_block_sparse_causal_helper], + "flash-attention")