Flash Attention (Fused Tiled Attention)

Computes attention in a single kernel by tiling over the sequence dimension, using shared memory for intermediate results (scores, softmax, weighted sum).

Equivalence to Standard Attention

Standard attention (3 separate kernels):

scores[h,s] = scale * Σ_d Q[h,d] * K[kvHead,s,d]     -- score kernel
attn[h,s] = softmax(scores[h,:])                       -- softmax kernel
output[h,d] = Σ_s attn[h,s] * V[kvHead,s,d]           -- apply kernel

Flash attention (1 fused kernel): Same computation, but scores and attn stay in shared memory. No global memory write for intermediate scores/attn.

Proof of Equivalence

The CPU spec functions scaledDotForward, softmaxForward, and attentionForward are composed. Flash attention computes the same composition but without materializing intermediates:

flashAttention(Q, K, V, scale) = standard_attention(Q, K, V, scale)

This is verified numerically by verifyFlashEquivalence.

Memory Savings

Standard: O(numHeads × seqLen) global memory for scores + attn Flash: O(workgroupSize) shared memory only For seqLen=2048, numHeads=20: 160KB → ~4KB (40x reduction)

CPU Spec (for equivalence proof)

def Hesper.WGSL.FlashAttention.standardAttention (q : Array Float) (kCache vCache : Array (Array Float)) (scale : Float) : Array Float

Standard attention: score → softmax → apply (3 steps)

Equations

• Hesper.WGSL.FlashAttention.standardAttention q kCache vCache scale = Hesper.Training.VerifiedBackward.attentionForward q kCache vCache scale

def Hesper.WGSL.FlashAttention.flashAttentionSpec (q : Array Float) (kCache vCache : Array (Array Float)) (scale : Float) : Array Float

Flash attention CPU spec: same result, computed differently. This is intentionally written to show the tiled computation pattern.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.FlashAttention.verifyFlashEquivalence (tol : Float := 1e-4) : Bool

Verify flash attention produces same output as standard attention

Equations

• One or more equations did not get rendered due to their size.

GPU Kernel: Flash Attention Forward (single-token KV cache)

def Hesper.WGSL.FlashAttention.flashAttentionDynamicKernel (numHeads numKVHeads maxSeqLen headDim cacheLen : Nat) (scale : Float) (workgroupSize : Nat := 256) : Monad.ShaderM Unit

Flash attention forward kernel for single-token query with KV cache. One workgroup per head. Each workgroup:

1. Loads Q for this head from global memory
2. Iterates over cached K/V positions, computing online softmax
3. Writes final output to global memory

No intermediate scores/attn buffers needed.

@param numHeads Number of query heads @param numKVHeads Number of KV heads (GQA) @param cacheLen Number of positions in KV cache @param headDim Dimension per head @param scale 1/sqrt(headDim)

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.FlashAttention.flashAttentionDynamicParamsKernel (numHeads numKVHeads maxSeqLen headDim : Nat) (scale : Float) (workgroupSize : Nat := 256) : Monad.ShaderM Unit

Like flashAttentionDynamicKernel but reads cacheLen from params buffer. PTX is fixed regardless of cacheLen → fully cacheable. Supports headDim > workgroupSize via strided loops (shared_out).

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.FlashAttention.flashAttentionVecParamsKernel (numHeads numKVHeads maxSeqLen headDim : Nat) (scale : Float) : Monad.ShaderM Unit

doc 60 Session 1: warp-shuffle vec kernel.

Like flashAttentionDynamicParamsKernel but reduces the q·K[s] dot product via subgroupAdd (warp shuffle) instead of a shared-memory tree reduce. workgroupSize is fixed at 128 (4 warps × 32 lanes).

why this is faster

The legacy 256-thread tree reduce inside the cacheLen-step loop has log2(256)=8 barriers per K position. For Gemma 4 (head_dim=256, cacheLen up to ~200) that is 8 × cacheLen ≈ 1600 barriers per head per token. Replacing the tree with subgroupAdd cuts each K-step reduce to:

warp-shuffle sum : 5 shfl.bfly, 0 barriers lane-0 → smem write : 1 barrier warp-0 cross-warp : 1 subgroupAdd over numWarps=4 lanes (cheap) barrier : 1 ────────────────────── per-K total : 2 barriers (vs the old 8)

so we expect ~4× fewer barriers in the inner loop, plus subgroupAdd is a register-level shuffle that does not touch shared memory.

algorithm (unchanged from the legacy dynamic kernel)

Same online softmax over a serial K loop; workgroups are still per head (gridX = numHeads, gridY = 1). This is the conservative Option B (doc 60): we keep cacheLen-serial, just switch the reduce primitive. A future revision can re-shape to llama.cpp's K-parallel layout (multiple warps each processing a distinct K position) for the remaining gap.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.FlashAttention.flashAttentionBatchKernel (numHeads numKVHeads maxSeqLen headDim seqLen : Nat) (scale : Float) (workgroupSize : Nat := 256) : Monad.ShaderM Unit

Batched flash-attention for prefill: processes seqLen query tokens in one dispatch, attending each over its own causal prefix of the K/V cache.

Layout (column-major Q/output, KV cache as in single-token kernel):

• q[col * (numHeads * headDim) + h * headDim + d] -- col = query token index
• output[col * (numHeads * headDim) + h * headDim + d]
• k_cache, v_cache: [numKVHeads, maxSeqLen, headDim] (same as single-token)

Grid: (numHeads, seqLen, 1). Each WG owns one (head, query token). cacheLen for query token col = startPos + col + 1 (causal).

params[0] = startPos — KV cache length BEFORE this batch was written. For prefill from scratch, startPos = 0 → cacheLen = col+1 per token.

PTX is fixed for given (numHeads, numKVHeads, maxSeqLen, headDim, seqLen).

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.FlashAttention.flashAttentionBatchKernelF16 (numHeads numKVHeads maxSeqLen headDim seqLen : Nat) (scale : Float) (workgroupSize : Nat := 256) : Monad.ShaderM Unit

f16 K/V cache version of flashAttentionBatchKernel. Reads K and V from the packed half2 cache (u32 per word holding 2 f16 values for consecutive dims) instead of f32. Same I/O contract for q/output/ params — only the cache buffers differ.

Cache layout (matches V11 + RopeKF16): cache[kvHead * maxSeqLen * (headDim/2) + pos * (headDim/2) + dPair] = pack2x16float(K[kvHead, pos, 2dPair], K[kvHead, pos, 2dPair+1])

Each thread loops 2 dims per inner iter (one u32 read covers a pair). Otherwise the algorithm is identical to the f32 version.

Equations

• One or more equations did not get rendered due to their size.

Subgroup Flash Attention (M=1 decode, no barriers, no shared mem)

def Hesper.WGSL.FlashAttention.flashAttentionSubgroupKernel (numHeads numKVHeads maxSeqLen headDim cacheLen : Nat) (scale : Float) : Monad.ShaderM Unit

Subgroup-based M=1 flash attention kernel. Replaces flashAttentionDynamicKernel's 256-thread-tree-reduce + 10-barriers-per-position scheme with a 32-thread (1 hardware subgroup) design:

• 1 workgroup per attention head, 32 threads
• headDim is partitioned across lanes with stride 32 — lane tid owns dims {tid, tid+32, tid+64, ...}. Gemma 4 has headDim=128 so that's exactly 4 dims per lane, fully unrolled.
• Q is read once per head and held in per-lane registers (q0..q3), no shared memory.
• Output accumulator is also per-lane registers (o0..o3).
• Per cached position:
  • lane reads its 4 K values, computes 4 FMAs against q*, then ONE subgroupAdd → score broadcast to every lane
  • scalar online softmax update in every lane (identical results because the subgroupAdd broadcasts the same dot product to all 32 lanes)
  • lane reads its 4 V values and applies the rescale/contrib factors to o0..o3
• No shared memory, no barriers, no tree reduction.

Constraint: headDim % 32 == 0 && headDim / 32 ≤ 8. Gemma 4's headDim=128 → 4 regs per lane, well within register budget. Also requires subgroup support (every desktop Vulkan/NVIDIA driver).

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.FlashAttention.flashAttentionSubgroupParamsKernel (numHeads numKVHeads maxSeqLen headDim : Nat) (scale : Float) : Monad.ShaderM Unit

Params-buffer variant of flashAttentionSubgroupKernel: reads cacheLen from a 2-u32 params buffer (position 1 = cacheLen, matching the layout used by flashAttentionDynamicParamsKernel). This keeps the PTX shape-stable across decode positions — exactly what CUDA Graph capture needs to replay correctly past the initial cacheLen-at-capture boundary.

All other properties are identical to flashAttentionSubgroupKernel: 32-thread workgroup, per-lane headDim slicing, single subgroupAdd per position, online softmax in registers, no shared memory, no barriers.

Equations

• One or more equations did not get rendered due to their size.

SWA variant: subgroup flash attention with sliding window

def Hesper.WGSL.FlashAttention.flashAttentionSWASubgroupKernel (numHeads numKVHeads maxSeqLen headDim cacheLen windowSize currentPos : Nat) (scale : Float) : Monad.ShaderM Unit

SWA (Sliding Window Attention) variant of flashAttentionSubgroupKernel. Identical layout, but only positions in [currentPos - windowSize + 1, currentPos] contribute to the softmax. We implement that by masking: positions outside the window still loop but get score = -inf so their softmax weight is zero.

Simpler than llama.cpp's "only loop over window" optimisation (that would need a dynamic loop start and break the compile-time cacheLen specialisation), but the wasted positions are just a few scalar ops per masked-out index per lane, which is negligible vs. the K/V memory loads we skip.

Equations

• One or more equations did not get rendered due to their size.

Dynamic Flash Attention with params buffer (production)

def Hesper.WGSL.FlashAttention.flashAttentionParamsKernel (numHeads numKVHeads maxSeqLen headDim : Nat) (scale : Float) (workgroupSize : Nat := 256) : Monad.ShaderM Unit

Flash attention with dynamic cacheLen from params buffer. Same as in-place kernel but reads cacheLen from params[1] (u32). Uses diagnostic(off, derivative_uniformity) to allow barrier.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.FlashAttention.executeFlashAttentionWithParams {β : Type} [GPUBackend β] (ctx : β) (qBuf kCacheBuf vCacheBuf paramsBuf outputBuf : GPUBackend.Buf β) (numHeads numKVHeads maxSeqLen headDim : Nat) (scale : Float) (cacheRef : Option (IO.Ref (Option (GPUBackend.CachedDispatch β))) := none) : IO Unit

Execute flash attention with params buffer (dynamic cacheLen, 1 dispatch). Same WGSL source for all cacheLen → 100% pipeline cache hit rate. Takes an optional cacheRef so callers can share a CachedDispatch across decode steps; when none, falls back to the throwaway-ref anti-pattern (for first-call/prefill sites).

Equations

• One or more equations did not get rendered due to their size.

In-Place Flash Attention (single tile, no merge)

def Hesper.WGSL.FlashAttention.flashAttentionInPlaceKernel (numHeads numKVHeads maxSeqLen headDim cacheLen : Nat) (scale : Float) (workgroupSize : Nat := 256) : Monad.ShaderM Unit

Flash attention kernel where Q input and output share the same buffer. Q is loaded into shared memory first, then output overwrites the buffer. Single read-write buffer avoids WebGPU aliasing. 1 dispatch only.

Equations

• One or more equations did not get rendered due to their size.

Tiled Flash Attention (v2) — High Parallelism

def Hesper.WGSL.FlashAttention.flashAttentionTiledPhase1 (numHeads numKVHeads maxSeqLen headDim cacheLen tileSize : Nat) (scale : Float) (workgroupSize : Nat := 256) : Monad.ShaderM Unit

Tiled flash attention: Phase 1 — each tile computes partial online softmax. Dispatch: (numHeads, numTiles). Each workgroup processes tileSize positions. Outputs per tile: partial_output[headDim], partial_max[1], partial_sumexp[1]

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.FlashAttention.flashAttentionTiledPhase2 (numHeads headDim numTiles : Nat) : Monad.ShaderM Unit

Tiled flash attention: Phase 2 — merge partial results. Each thread handles one output dimension for one head. Dispatch: (numHeads * headDim)

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.FlashAttention.createFlashPartialBuffer {β : Type} [GPUBackend β] (ctx : β) (numHeads maxSeqLen headDim : Nat) (tileSize : Nat := 32) : IO (GPUBackend.Buf β)

Pre-allocate partial buffer for tiled flash attention. Call once during initialization, reuse across all tokens.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.FlashAttention.executeFlashAttentionTiled {β : Type} [GPUBackend β] (ctx : β) (qBuf kCacheBuf vCacheBuf outputBuf : GPUBackend.Buf β) (numHeads numKVHeads maxSeqLen headDim cacheLen : Nat) (scale : Float) (partialBuf : Option (GPUBackend.Buf β) := none) (phase1Ref phase2Ref : Option (IO.Ref (Option (GPUBackend.CachedDispatch β))) := none) (kcrLookup : Option (UInt64 → IO (IO.Ref (Option (GPUBackend.CachedDispatch β)))) := none) : IO Unit

Execute tiled flash attention (2 phases).

The phase1Ref / phase2Ref parameters are the per-call dispatch cache slots. Earlier code synthesised them via IO.mkRef none inside the function body when callers passed none — that turned every invocation into a cold cudaExecuteImpl miss because the ref was thrown away at end of call. Doc 57 §3b.4 traced the 4 ms graphs-OFF host-time gap to that pattern.

Now, when callers pass none, we provide a kcrLookup callback that resolves a (cacheKey → IO.Ref) for refs that survive across calls — typically routed through KernelCacheRefs from the model forward. As long as the same cacheLen recurs across forward invocations (and within a forward, all layers share one cacheLen) the dispatches hit the cache. Pass kcrLookup := none only if you want the legacy throwaway behaviour (e.g. a one-shot test).

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.FlashAttention.executeFlashAttentionVecParams {β : Type} [GPUBackend β] (ctx : β) (qBuf kCacheBuf vCacheBuf outputBuf paramsBuf : GPUBackend.Buf β) (numHeads numKVHeads maxSeqLen headDim : Nat) (scale : Float) (kcrLookup : Option (UInt64 → IO (IO.Ref (Option (GPUBackend.CachedDispatch β)))) := none) : IO Unit

doc 60 Session 1: launcher for flashAttentionVecParamsKernel.

Same call shape as executeFlashAttentionDynamic — single-token Q, cacheLen lives in state.paramsBuf[1] (so the PTX is fully cacheable and CUDA-Graph friendly). The only difference vs the legacy launcher is the workgroup size (128 instead of min 256 headDim) and the use of the new shader.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.FlashAttention.executeFlashAttentionDynamic {β : Type} [GPUBackend β] (ctx : β) (qBuf kCacheBuf vCacheBuf outputBuf : GPUBackend.Buf β) (numHeads numKVHeads maxSeqLen headDim cacheLen : Nat) (scale : Float) (kcrLookup : Option (UInt64 → IO (IO.Ref (Option (GPUBackend.CachedDispatch β)))) := none) : IO Unit

Equations

• One or more equations did not get rendered due to their size.

def Hesper.WGSL.FlashAttention.executeFlashAttention {β : Type} [GPUBackend β] (ctx : β) (qBuf kCacheBuf vCacheBuf outputBuf : GPUBackend.Buf β) (numHeads numKVHeads cacheLen headDim : Nat) (scale : Float) : IO Unit

Execute flash attention with static cacheLen (for testing). Uses dynamic kernel with a params buffer containing cacheLen.

Equations

• One or more equations did not get rendered due to their size.

Sliding Window Flash Attention

def Hesper.WGSL.FlashAttention.flashAttentionSWAKernel (numHeads numKVHeads maxSeqLen headDim cacheLen windowSize currentPos : Nat) (scale : Float) (workgroupSize : Nat := 256) : Monad.ShaderM Unit

Flash attention with sliding window masking. Only attends to positions within [max(0, pos - windowSize + 1), pos]. Uses compile-time cacheLen and windowSize for the loop bounds.

For SWA layers in Gemma 4 ISWA architecture. Positions outside the window get -inf score (masked out by online softmax).

@param numHeads Number of query heads @param numKVHeads Number of KV heads (GQA) @param maxSeqLen Maximum sequence length (cache buffer size) @param headDim Per-head dimension @param cacheLen Current number of cached positions @param windowSize Sliding window size (e.g., 512) @param currentPos Current query position @param scale Attention scale (1/sqrt(headDim))

Equations

• One or more equations did not get rendered due to their size.