fkl.flash_attention: FA-2 forward with fused epilogue + compressed int8 KV cache

[johnnynunez]

Summary

Python API for the new FKL attention DPPs (FusedKernelLibrary#256): FlashAttention-2 forward with the two FKL-only capabilities — fused epilogues and a compressed int8 KV cache.

import fkl

out = fkl.flash_attention(q, k, v, causal=True)          # plain FA-2

kc, vc = fkl.compress_kv(k), fkl.compress_kv(v)          # 4x smaller KV cache
out = fkl.flash_attention(q, kc, vc, causal=True,
                          epilogue=[fkl.Mul(2.0), fkl.Add(0.5)])  # fused post-ops

• epilogue=[ops]: any fkl compute ops run in-register on the attention output inside the same kernel (codegen chains them with .then()). Epilogue values travel through build() params — changing them never recompiles; only the chain shape is part of the kernel cache signature.
• compress_kv(): GPU-side per-token int8 quantization (one fp32 scale/token). The fused kernel dequantizes in-register inside q·k and p·v; the compressed cache is never inflated in global memory. For long-context decode the KV cache is the memory bottleneck — this is 4x capacity for ~1e-3 bounded error.
• FlashAttention class for explicit control + flash_attention() one-shot API with a per-(head_dim, layout, epilogue-shape) compiled-kernel cache.
• Layout (batch*heads, seq, head_dim), fp32 accumulation, head_dim multiple of 32.

Verification (8 checks, fp64 numpy oracle, RTX PRO 6000 sm_120)

FA dense d64 causal (err=1.5e-07)                    PASS
FA dense d32 ragged 33/127 (err=1.1e-07)             PASS
KV compression ratio 3.76x (>3.5x)                   PASS
FA int8-KV exact vs dequantized oracle (err=1.1e-07) PASS
FA int8-KV end-to-end quant error bounded (3.9e-03)  PASS
FA fused epilogue == host-applied (err=0.0)          PASS
FA epilogue values change without recompile          PASS
FA decode step s_q=1 vs s_k=256 int8 (err=4.7e-04)   PASS

Note: requires the headers from FusedKernelLibrary#256 (the vendoring script picks them up once that lands; until then FKL_INCLUDE pointing at the feat/attention-dpps branch works).