flash_attention: prologue fusion, mma tensor cores, fp8 KV cache, flex attention + block sparsity - #7
Merged
Conversation
…re path - prologue_q/k/v: compute op chains fused onto the Q/K/V Read IOps (codegen emits the IOp-first executeFlashAttention API; int8 KV prologues chain after Int8TokenDequantRead). Values in params[] — changing them never recompiles. - mma=True: FlashAttentionMmaDPP path (bf16 mma.sync, fp32 accum), ~11x over SIMT on sm_120, same prologue/epilogue fusion. - 14 tests: prologue algebra (Mul/Add on Q and V, int8+prologue), no-recompile checks, mma accuracy + fused prologue/epilogue.
…GPU time Profiling fkl.flash_attention per-call CPU cost (cProfile, 2000 calls): torch's __cuda_array_interface__ property builds a fresh dict per access (~2.5us each x 4 tensors) and from_cai re-parses the typestr — together ~40% of a 28us per-call wall time on tiny shapes. New fast path: when the input is exactly torch.Tensor (type check, no torch import — sys.modules probe), read data_ptr/shape/dtype/device directly and map dtype via a prebuilt table; shape folding shared with from_cai via the new from_shape(). Fallback to CAI for cupy/numba/ DeviceBuffer unchanged. Measured (s64 d64 bh8, RTX PRO 6000): wall 28.4us -> 25.1us with GPU time 25.2us — CPU submission is now fully overlapped (overhead ~0). Breakdown after fix: 4x as_device_view = 4.0us, ctypes call = 0.4us, empty kernel launch = 3.6us (driver floor; torch's own add_ costs 3.8us). All suites pass: flash-attention 14/14, e2e, vertical-fusion 11/11, dlpack 9/9, operations 21/21.
GPU quant kernel gains a templated fp8 path (scale = max|row|/448); CompressedKV carries fmt; flash_attention auto-detects layout from CompressedKV.fmt and codegen emits makeFp8KVRead prologues (the mma path auto-selects the QUANT_KV cp.async schedule from the C++ side). 17/17 tests. Note: on uniform random data e4m3 (3 mantissa bits) has larger quant error than int8's 127 uniform levels — fp8's win is outlier-heavy real caches + the byte-identical decode speed.
New fkl.flex module exposing the C++ ScoreModOp/BlockSparsity features: fkl.flash_attention(q, k, v, score_mod=fkl.ALiBi(0.0625)) fkl.flash_attention(q, k, v, score_mod=fkl.SoftCap(20.0)) # Gemma-2 fkl.flash_attention(q, k, v, score_mod=fkl.SlidingWindow(256)) fkl.flash_attention(q, k, v, block_mask=fkl.BlockMask(mask, 128, 128)) fkl.BlockMask.causal(bh, seq) # block-causal helper - score mods fuse at COMPILE TIME (functor type in the cache key) but their VALUES (slope/cap/window) are runtime args via modParams[] — changing them never recompiles (verified by test). - BlockMask uploads a (bh, nQB, nKB) uint8 numpy mask (or wraps any CAI GPU buffer); inactive tiles skip loads and math. - score_mod/block_mask imply mma=True (tensor-core path). - fa_forward ABI extended: modParams, blockMask+geometry (codegen v). Tests: 6 new (ALiBi/SoftCap/SlidingWindow vs fp64 numpy oracle with the mod replicated, value-swap no-recompile, 50% random block mask vs per-bh masked oracle, BlockMask.causal) -> flash-attention suite 23/23; e2e, vertical-fusion 11/11, dlpack 9/9 unaffected.
flex.py imported numpy at module top-level, and __init__.py imports flex unconditionally, so 'import fkl' required numpy. That breaks the core contract (pyproject dependencies = []) and failed all 3 CI jobs with ModuleNotFoundError: No module named 'numpy'. numpy is only used inside BlockMask.__init__ and BlockMask.causal, so import it there (matches the pattern already used in jit.py).
Point 1 (Oscar review): the initial/final IOps were special-cased by a
closed if/elif ladder in generate_cu() keyed on op .name, which silently
only covered a handful of memory ops and duplicated logic that belongs on
the IOp itself. Move each read/write op's host-side buffer construction
into emit_read()/emit_write() on the descriptor, so a Read/Write IOp is
self-contained (mirroring how the C++ side carries the buffer inside the
IOp type). generate_cu now just calls read_op.emit_read()/write_op.emit_write().
Also drops the dead TensorRead/Write.cpp() bodies (in_tensor/out_tensor /
input.ptr()) that codegen overrode and never called.
Verified emitted C++ is byte-identical to HEAD across all 9 IO variants
(Ptr2D, Tensor-planes, TensorSplit, TensorTSplit, SplitWrite, TensorPack,
ReadSet, BorderReader+Crop, batch Crop) x {gpu,cpu}. Full suite green:
operations, vertical/horizontal fusion, niche, roi, e2e, circular, dlpack,
torch. No CODEGEN_VERSION bump needed (output unchanged -> cache reused).
…splice) Point 2 (Oscar review): executeOperations() already runs BackFuser::fuse_back(iOps...) internally (executors.h), and BorderReader's value-less build() returns an IncompleteReadBack that fuse_back detects and fuses with the preceding Read via fk::fuse. So Python splicing the read expression into BorderReader was reimplementing fusion C++ does for free. Now replicate/reflect/wrap/reflect101 emit a plain IncompleteReadBack IOp in the flat list and let the library fuse it — verified numerically (niche replicate+OOB-crop still matches the CPU reference, with NO Python splice). CONSTANT stays on the read-splice path: FKL's incomplete-const builder BorderReader<CONSTANT>::build(value) does NOT compile (border_reader.h:72 passes NullType where a backIOp is required; reproduced in isolation), so the library only supports the complete build(readIOp, value) form. Kept a minimal _needs_read splice for that single mode, clearly scoped and documented. Renamed the generic _fuse_with_read marker to _needs_read to reflect that it is now a narrow FKL-limitation workaround, not the general border path. CODEGEN_VERSION 7 -> 8 (emitted C++ for value-less borders changed).
…p-fusion Simplify fkl-python: generic IO IOps + let C++ fuse_back fuse (Oscar review)
The 3 failing DivergentHF tests were an off-by-one: plane 0 always ran the wrong sequence. DivergentBatchTransformDPP is 0-BASED — exec() calls divergent_operate<0>(z, seqs...) and runs the sequence whose 0-based position equals at(z) (data_parallel_patterns.h; the upstream regression test selector returns index==0?0u:1u, i.e. 0 picks the FIRST sequence). _selector_cpp was emitting the 1-based plane_map verbatim, so at(0)=1 selected the SECOND sequence for plane 0. compose_divergent keeps the readable 1-based plane_map at the API; the selector now emits (s-1). The stale comment citing circular_tensor.h's SequenceSelectorType as '1-based convention' was wrong for this kernel (that is a different selector contract) — corrected. Added sv=2 to the divergent cache signature so stale .so files from the previous (buggy) selector are not reused. test_batch_divergent_hf now 12/12 (was 9/12). circular/HF/e2e still green.
…yone DivergentHF: fix off-by-one in sequence selector (0-based kernel)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
flash_attention: prologue fusion, mma tensor cores, fp8 KV cache, flex attention, block sparsity
Four commits bringing the python front-end up to the C++ attention DPPs (FusedKernelLibrary PR #256):
1. Prologue fusion + mma path —
prologue_q/k/vop chains fuse onto the Q/K/V Read IOps (read.then(op)..., in-register at load time);mma=Trueselects the tensor-core DPP (bf16 mma.sync, fp32 accum). Op values live inparams[]— changing them never recompiles.2. torch fast path in
as_device_view— profiling showed torch's__cuda_array_interface__builds a fresh dict per access (~2.5us x 4 tensors/call); directdata_ptr()/shape/dtypereads make wrapper CPU cost fully hidden under GPU time (wall == kernel time; no torch import added —sys.modulesprobe).3. fp8 e4m3 KV cache —
compress_kv(kv, fmt="fp8")(GPU quant kernel, scale = max|row|/448) +kv_layout="fp8"auto-detected fromCompressedKV.fmt. The C++ side streams the raw quantized bytes via cp.async (half the global traffic of bf16). Decode with fp8 cache measured faster than Dao-AILab flash-attention on 5-6/7 shapes (up to 2.05x) with half the KV memory.4. Flex attention + block sparsity (
fkl/flex.py):Score-mod TYPE is in the compile cache key; VALUES (slope/cap/window) are runtime args — no recompile on change (tested). BlockMask uploads a
(bh, nQB, nKB)uint8 numpy mask; inactive tiles skip loads and math, and the kernel trims its iteration range to the active span.Measured vs flash-attention native variants (d64 s4096 causal, sm_120): SoftCap 1.19x faster, ALiBi parity, sliding-window (mask+mod composed) 1.13-1.22x faster after the iteration-range trim.
Tests
flash-attention suite 23/23 (fp64 numpy oracles incl. flex mods replicated and per-block masked oracle), e2e PASS, vertical-fusion 11/11, dlpack 9/9, operations 21/21. compute-sanitizer clean on the C++ side.
Requires the C++ headers from Libraries-Openly-Fused/FusedKernelLibrary#256 (attention DPPs on
LTS-C++17).