include/brotensor/ops/flash_attention.h
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #pragma once | ||
| 2 | |||
| 3 | // brotensor ops/flash_attention.h — Flash-attention family: tiled/windowed/varlen/qkvo/decode + KV cache. | ||
| 4 | |||
| 5 | #include "../tensor.h" | ||
| 6 | #include <cstdint> | ||
| 7 | |||
| 8 | namespace brotensor { | ||
| 9 | |||
| 10 | |||
| 11 | // Flash-attention-style fused attention (FP16, inference-only). Q, K, V are | ||
| 12 | // already projected. Tiled online softmax over Lk (no Lk-long materialisation); | ||
| 13 | // FP32 accumulation. | ||
| 14 | // Q: (Lq,D); K, V: (Lk,D) — FP16. | ||
| 15 | // d_mask: optional length-Lk FP32 mask (1 valid / 0 invalid); may be null. | ||
| 16 | // num_heads divides D. | ||
| 17 | // causal: if true, key k attends to query q only when k <= q (requires | ||
| 18 | // Lq == Lk); combines multiplicatively with d_mask. | ||
| 19 | // O: (Lq,D) FP16, resized as needed. | ||
| 20 | void flash_attention_forward(const Tensor& Q, | ||
| 21 | const Tensor& K, | ||
| 22 | const Tensor& V, | ||
| 23 | const float* d_mask, | ||
| 24 | int num_heads, | ||
| 25 | bool causal, | ||
| 26 | Tensor& O); | ||
| 27 | |||
| 28 | |||
| 29 | // GQA generalisation of flash_attention_forward: grouped-query self-attention | ||
| 30 | // over pre-projected Q/K/V, causal OR bidirectional. Q carries num_q_heads, | ||
| 31 | // K/V carry num_kv_heads (num_kv_heads must divide num_q_heads; equal == plain | ||
| 32 | // MHA). causal == false is full bidirectional attention over all Lk keys — the | ||
| 33 | // encoder prefill a bidirectional-ified decoder needs (LLM2Vec). causal == true | ||
| 34 | // reproduces flash_attention_forward's causal mask with GQA (requires Lq == Lk). | ||
| 35 | // Tiled online softmax, FP16/BF16/FP32, FP32 accumulation; shares the windowed | ||
| 36 | // kernel's GQA + key-mask machinery. | ||
| 37 | // Q: (Lq, num_q_heads*head_dim); K, V: (Lk, num_kv_heads*head_dim). | ||
| 38 | // d_mask: optional length-Lk FP32 key mask (1 valid / 0 invalid); may be null. | ||
| 39 | // num_kv_heads divides num_q_heads; head_dim = Q.cols/num_q_heads must equal | ||
| 40 | // K.cols/num_kv_heads. Lk >= Lq; causal additionally requires Lq == Lk. | ||
| 41 | // O: (Lq, num_q_heads*head_dim), same dtype as Q, resized as needed. | ||
| 42 | void flash_attention_gqa_forward(const Tensor& Q, | ||
| 43 | const Tensor& K, | ||
| 44 | const Tensor& V, | ||
| 45 | const float* d_mask, | ||
| 46 | int num_q_heads, | ||
| 47 | int num_kv_heads, | ||
| 48 | bool causal, | ||
| 49 | Tensor& O); | ||
| 50 | |||
| 51 | |||
| 52 | // Sliding-window self-attention (FP32/FP16/BF16, inference-only) — the local | ||
| 53 | // attention of streaming neural codecs (e.g. Qwen3-TTS / Mimi), autoregressive decode, | ||
| 54 | // and bidirectional encoder models (e.g. ModernBERT). Q, K, V already projected, | ||
| 55 | // (L, num_heads*head_dim). | ||
| 56 | // The Lq queries occupy the last Lq positions of a length-Lk sequence | ||
| 57 | // (q_offset = Lk - Lq): query row r is at absolute position aq = r + q_offset. | ||
| 58 | // - causal == true: attends keys [max(0, aq-window+1), aq]. window <= 0 means | ||
| 59 | // unbounded causal — identical to flash_attention_forward with causal=true. | ||
| 60 | // - causal == false: attends keys [max(0, aq - window/2), min(Lk - 1, aq + window/2)] | ||
| 61 | // when window > 0, or all keys [0, Lk) when window <= 0 (full bidirectional). | ||
| 62 | // - Lq == Lk: self-attention (prefill / codec sliding window / encoder). | ||
| 63 | // - Lq < Lk: incremental decode of an Lq-token block over a K/V cache; | ||
| 64 | // Lq == 1 with window <= 0 attends every cached key (full cache attention), | ||
| 65 | // replacing a varlen call with no cu_seqlens upload. Requires Lk >= Lq. | ||
| 66 | // d_mask: optional length-Lk FP32 key mask (1 valid / 0 invalid), combined | ||
| 67 | // multiplicatively with the window; may be null. | ||
| 68 | // num_heads divides D. O: (Lq, num_heads*head_dim), resized as needed. | ||
| 69 | void flash_attention_windowed_forward(const Tensor& Q, | ||
| 70 | const Tensor& K, | ||
| 71 | const Tensor& V, | ||
| 72 | const float* d_mask, | ||
| 73 | int num_heads, | ||
| 74 | int window, | ||
| 75 | Tensor& O, | ||
| 76 | bool causal = true); | ||
| 77 | |||
| 78 | |||
| 79 | // Packed variable-length multi-head attention, forward only (Qwen3-VL window | ||
| 80 | // attention). All sequences in a batch live contiguously in one packed tensor; | ||
| 81 | // per-sequence boundaries come from `cu_seqlens_*` INT32 prefix-sum buffers of | ||
| 82 | // length `batch_size + 1`. Sequence b covers Q rows | ||
| 83 | // [cu_seqlens_q[b], cu_seqlens_q[b+1]) and attends to K/V rows | ||
| 84 | // [cu_seqlens_k[b], cu_seqlens_k[b+1]); no cross-sequence attention. Mirrors | ||
| 85 | // flash_attn_varlen_func semantics. | ||
| 86 | // Q: (total_tokens_q, num_heads * head_dim) — FP16/BF16/FP32 (GPU) / FP32 (CPU). | ||
| 87 | // K, V: (total_tokens_k, num_heads * head_dim) — same packing, num_heads | ||
| 88 | // matches Q (no GQA — that's a future cleanup). | ||
| 89 | // cu_seqlens_q, cu_seqlens_k: DEVICE pointers (CUDA/Metal device, raw host | ||
| 90 | // pointers on CPU), length batch_size + 1, INT32 prefix sums. Same | ||
| 91 | // convention as `const float* d_mask` elsewhere — caller owns the | ||
| 92 | // buffer; the op does not allocate or copy. | ||
| 93 | // max_seqlen_q, max_seqlen_k: bound the longest sequence's length — used by | ||
| 94 | // the GPU kernel for block sizing; the CPU impl ignores them. | ||
| 95 | // causal: if true, key k in sequence b attends to query q only when | ||
| 96 | // (k - cu_seqlens_k[b]) <= (q - cu_seqlens_q[b]); only meaningful when | ||
| 97 | // the per-sequence Q and K lengths match. | ||
| 98 | // O: (total_tokens_q, num_heads * head_dim) — overwritten, resized + dtype-set. | ||
| 99 | void flash_attention_varlen_forward(const Tensor& Q, | ||
| 100 | const Tensor& K, | ||
| 101 | const Tensor& V, | ||
| 102 | const int32_t* cu_seqlens_q, | ||
| 103 | const int32_t* cu_seqlens_k, | ||
| 104 | int batch_size, | ||
| 105 | int max_seqlen_q, | ||
| 106 | int max_seqlen_k, | ||
| 107 | int num_heads, | ||
| 108 | int head_dim, | ||
| 109 | bool causal, | ||
| 110 | Tensor& O); | ||
| 111 | |||
| 112 | |||
| 113 | // Packed variable-length bidirectional self-attention straight off a fused QKV | ||
| 114 | // projection — the encoder-batching entry point (many independent sequences of | ||
| 115 | // any lengths packed back to back, no padding, one launch per layer). | ||
| 116 | // QKV: (L, 3*num_heads*head_dim) — per row [q heads | k heads | v heads], the | ||
| 117 | // layout a fused Wqkv linear produces. FP16/BF16 on the GPU (head_dim 64 | ||
| 118 | // takes the tensor-core kernel, other widths a generic one), FP32 on | ||
| 119 | // both CPU and GPU. | ||
| 120 | // seq_bounds: (L, 2) INT32 on the same device — row r's sequence occupies | ||
| 121 | // rows [seq_bounds[r,0], seq_bounds[r,1]) and must contain r. Rows only | ||
| 122 | // ever attend inside their own sequence. | ||
| 123 | // window > 0: row r attends keys j of its sequence with |r - j| <= window/2 | ||
| 124 | // (ModernBERT / flash_attention_windowed_forward's bidirectional window); | ||
| 125 | // window <= 0: every key of its sequence. | ||
| 126 | // O: (L, num_heads*head_dim), same dtype as QKV, resized as needed. | ||
| 127 | // Row → sequence membership is data, not launch shape: the grid depends only | ||
| 128 | // on L, so a CUDA graph captured for one L replays for any packing of it. | ||
| 129 | void flash_attention_packed_qkv_forward(const Tensor& QKV, | ||
| 130 | const Tensor& seq_bounds, | ||
| 131 | int num_heads, | ||
| 132 | int window, | ||
| 133 | Tensor& O); | ||
| 134 | |||
| 135 | |||
| 136 | // Backward of flash_attention_packed_qkv_forward — the same packing, bounds | ||
| 137 | // and window, so a packed encoder can be trained (or differentiated to its | ||
| 138 | // inputs) through the exact attention it runs. Recompute-based: consumes no | ||
| 139 | // forward caches; re-derives each row's softmax statistics, then accumulates | ||
| 140 | // dQ per query row and dK / dV per key row (no atomics, deterministic). | ||
| 141 | // QKV: the forward's input (L, 3*num_heads*head_dim), post-RoPE. | ||
| 142 | // dO: (L, num_heads*head_dim) upstream gradient of the forward's O, same | ||
| 143 | // dtype as QKV. | ||
| 144 | // seq_bounds, num_heads, window: exactly as passed to the forward. | ||
| 145 | // dQKV: (L, 3*num_heads*head_dim) OVERWRITTEN (resized + dtype-set to | ||
| 146 | // QKV's) — per row [dQ heads | dK heads | dV heads]. | ||
| 147 | // FP16/BF16/FP32 on the GPU (FP32 math), FP32 on the CPU. Cost is | ||
| 148 | // O(L * span * head_dim) with span the keys a row attends (window + 1, or its | ||
| 149 | // sequence length). | ||
| 150 | void flash_attention_packed_qkv_backward(const Tensor& QKV, | ||
| 151 | const Tensor& dO, | ||
| 152 | const Tensor& seq_bounds, | ||
| 153 | int num_heads, | ||
| 154 | int window, | ||
| 155 | Tensor& dQKV); | ||
| 156 | |||
| 157 | |||
| 158 | // Backward of flash_attention_varlen_forward — packed variable-length attention | ||
| 159 | // over pre-projected Q/K/V (no projection weights, no biases — projections are | ||
| 160 | // handled by the caller's linear layer; bias gradients belong to that layer). | ||
| 161 | // Recompute-based: consumes no forward caches, re-runs the per-sequence | ||
| 162 | // softmax then reverses it. Matches flash_attention_backward's per-sequence | ||
| 163 | // math; only the cu_seqlens scatter/gather and per-sequence causal differ. | ||
| 164 | // Q: (total_tokens_q, num_heads*head_dim) — FP16/BF16/FP32 (GPU) / FP32 (CPU). | ||
| 165 | // K, V: (total_tokens_k, num_heads*head_dim) — same dtype, same packing. | ||
| 166 | // O: (total_tokens_q, ...) forward output — currently unused (kept for API | ||
| 167 | // symmetry with flash_attention_backward). | ||
| 168 | // dO: (total_tokens_q, ...) upstream gradient, same dtype as Q/K/V. | ||
| 169 | // cu_seqlens_q, cu_seqlens_k: DEVICE pointers on GPU, host pointers on CPU, | ||
| 170 | // length batch_size + 1, INT32 prefix sums. Same convention as the | ||
| 171 | // forward. | ||
| 172 | // max_seqlen_q, max_seqlen_k: bound the longest per-sequence length — used | ||
| 173 | // by the GPU kernel for block/scratch sizing; the CPU impl ignores them. | ||
| 174 | // causal: per-sequence; only meaningful when per-sequence Lq == Lk. | ||
| 175 | // dQ: (total_tokens_q, ...); dK, dV: (total_tokens_k, ...) — OVERWRITTEN | ||
| 176 | // (resized + dtype-set), matching flash_attention_backward's contract. | ||
| 177 | // Out-of-range positions (causal-excluded, or in a fully empty K sequence) | ||
| 178 | // contribute nothing to any gradient. No cross-sequence attention. | ||
| 179 | void flash_attention_varlen_backward(const Tensor& Q, | ||
| 180 | const Tensor& K, | ||
| 181 | const Tensor& V, | ||
| 182 | const Tensor& O, | ||
| 183 | const Tensor& dO, | ||
| 184 | const int32_t* cu_seqlens_q, | ||
| 185 | const int32_t* cu_seqlens_k, | ||
| 186 | int batch_size, | ||
| 187 | int max_seqlen_q, | ||
| 188 | int max_seqlen_k, | ||
| 189 | int num_heads, | ||
| 190 | int head_dim, | ||
| 191 | bool causal, | ||
| 192 | Tensor& dQ, | ||
| 193 | Tensor& dK, | ||
| 194 | Tensor& dV); | ||
| 195 | |||
| 196 | |||
| 197 | // Flash-attention with QKV and output projections fused at the boundary. | ||
| 198 | // Projects X->Q, Ctx->K,V (or X->Q,K,V when Ctx is null), runs the tiled core, | ||
| 199 | // then projects with Wo. FP16 throughout. | ||
| 200 | // X: (Lq,D) FP16, query source. | ||
| 201 | // Ctx: (Lk,D_ctx) FP16 or null. Null => self-attention (Ctx<-X). D_ctx may | ||
| 202 | // differ from D (e.g. SD1.5 cross-attention). | ||
| 203 | // Wq, Wo: (D,D); Wk, Wv: (D,D_ctx) — FP16. | ||
| 204 | // bq, bk, bv, bo: optional (D,1) FP16 biases; null to skip. | ||
| 205 | // d_mask: optional length-Lk FP32 mask. num_heads divides D. | ||
| 206 | // causal: see flash_attention_forward (typically with Ctx == null). | ||
| 207 | // O: (Lq,D) FP16, resized as needed. | ||
| 208 | void flash_attention_qkvo_forward(const Tensor& X, | ||
| 209 | const Tensor* Ctx, | ||
| 210 | const Tensor& Wq, const Tensor* bq, | ||
| 211 | const Tensor& Wk, const Tensor* bk, | ||
| 212 | const Tensor& Wv, const Tensor* bv, | ||
| 213 | const Tensor& Wo, const Tensor* bo, | ||
| 214 | const float* d_mask, | ||
| 215 | int num_heads, | ||
| 216 | bool causal, | ||
| 217 | Tensor& O); | ||
| 218 | |||
| 219 | |||
| 220 | // Backward of flash_attention_qkvo_forward. Recompute-style: consumes no | ||
| 221 | // forward caches — re-runs the attention math from the inputs, then reverses | ||
| 222 | // it. FP16 storage, FP32 accumulation. All shape / dtype / Ctx-null / | ||
| 223 | // rectangular-Wk-Wv / causal / optional-bias rules match the forward; pass the | ||
| 224 | // same values. | ||
| 225 | // dO: (Lq,D) FP16 upstream. | ||
| 226 | // dX: (Lq,D) FP16 overwritten. For self-attention (Ctx null) dX absorbs | ||
| 227 | // the K/V-projection gradients too: dX = dQ.Wq + dK.Wk + dV.Wv. | ||
| 228 | // dCtx: (Lk,D_ctx) FP16 overwritten; must be null iff Ctx is null. For | ||
| 229 | // cross-attention dCtx = dK.Wk + dV.Wv. | ||
| 230 | // dWq, dWo: (D,D); dWk, dWv: (D,D_ctx) — FP16, accumulated (caller zeros). | ||
| 231 | // dbq, dbk, dbv, dbo: (D,1) FP16, accumulated iff the matching forward bias | ||
| 232 | // was non-null; pass null otherwise (the null/non-null symmetry must | ||
| 233 | // be exact — a mismatch is rejected). | ||
| 234 | // Causal- and mask-excluded positions contribute nothing to any gradient. | ||
| 235 | void flash_attention_qkvo_backward( | ||
| 236 | const Tensor& X, const Tensor* Ctx, | ||
| 237 | const Tensor& Wq, const Tensor* bq, | ||
| 238 | const Tensor& Wk, const Tensor* bk, | ||
| 239 | const Tensor& Wv, const Tensor* bv, | ||
| 240 | const Tensor& Wo, const Tensor* bo, | ||
| 241 | const float* d_mask, | ||
| 242 | int num_heads, | ||
| 243 | bool causal, | ||
| 244 | const Tensor& dO, | ||
| 245 | Tensor& dX, Tensor* dCtx, | ||
| 246 | Tensor& dWq, Tensor* dbq, | ||
| 247 | Tensor& dWk, Tensor* dbk, | ||
| 248 | Tensor& dWv, Tensor* dbv, | ||
| 249 | Tensor& dWo, Tensor* dbo); | ||
| 250 | |||
| 251 | |||
| 252 | // Backward of flash_attention_forward — bare attention core, no projection | ||
| 253 | // weights (what LoRA-style adapters need; projections are wrapped externally). | ||
| 254 | // Recompute-based; FP16 storage, FP32 accumulation. | ||
| 255 | // Q: (Lq,D); K, V: (Lk,D) — pre-projected forward inputs, FP16. | ||
| 256 | // O: (Lq,D) forward output — currently unused (kept for API symmetry). | ||
| 257 | // dO: (Lq,D) FP16 upstream. | ||
| 258 | // d_mask: optional length-Lk FP32 mask (null for unmasked); positions with | ||
| 259 | // mask[k] <= 0.5 are dropped. | ||
| 260 | // num_heads divides D. causal: match the forward (requires Lq == Lk). | ||
| 261 | // dQ: (Lq,D); dK, dV: (Lk,D) — FP16, overwritten (resized + dtype-set). | ||
| 262 | // Causal- and mask-excluded positions contribute nothing to dQ/dK/dV. | ||
| 263 | void flash_attention_backward(const Tensor& Q, | ||
| 264 | const Tensor& K, | ||
| 265 | const Tensor& V, | ||
| 266 | const Tensor& O, | ||
| 267 | const Tensor& dO, | ||
| 268 | const float* d_mask, | ||
| 269 | int num_heads, | ||
| 270 | bool causal, | ||
| 271 | Tensor& dQ, | ||
| 272 | Tensor& dK, | ||
| 273 | Tensor& dV); | ||
| 274 | |||
| 275 | |||
| 276 | // Project a context tensor through Wk/Wv into the exact (Lk,D) FP16 K/V buffers | ||
| 277 | // flash_attention_forward consumes. Used to precompute cross-attention K/V once | ||
| 278 | // per generate() (text context is fixed across denoising steps). Numerically | ||
| 279 | // identical to the K/V projection stage of flash_attention_qkvo_forward. | ||
| 280 | // ctx: (Lk,D_ctx) FP16. Wk, Wv: (D,D_ctx) FP16. bk, bv: (D,1) FP16 or null. | ||
| 281 | // K_out, V_out: (Lk,D) FP16, resized as needed. | ||
| 282 | void flash_attention_project_kv(const Tensor& ctx, | ||
| 283 | const Tensor& Wk, const Tensor* bk, | ||
| 284 | const Tensor& Wv, const Tensor* bv, | ||
| 285 | Tensor& K_out, | ||
| 286 | Tensor& V_out); | ||
| 287 | |||
| 288 | |||
| 289 | // Like flash_attention_qkvo_forward but K and V are already projected by the | ||
| 290 | // caller (typically via flash_attention_project_kv). Projects X->Q with Wq/bq, | ||
| 291 | // runs the tiled core against the supplied K/V, applies Wo/bo. FP16 throughout. | ||
| 292 | // X: (Lq,D); K, V: (Lk,D) — FP16. Wq, Wo: (D,D) FP16; bq, bo: (D,1) FP16 or null. | ||
| 293 | // d_mask: optional length-Lk FP32 mask. num_heads divides D. | ||
| 294 | // causal: see flash_attention_forward. O: (Lq,D) FP16, resized as needed. | ||
| 295 | void flash_attention_q_with_kv_cached_forward(const Tensor& X, | ||
| 296 | const Tensor& K, | ||
| 297 | const Tensor& V, | ||
| 298 | const Tensor& Wq, const Tensor* bq, | ||
| 299 | const Tensor& Wo, const Tensor* bo, | ||
| 300 | const float* d_mask, | ||
| 301 | int num_heads, | ||
| 302 | bool causal, | ||
| 303 | Tensor& O); | ||
| 304 | |||
| 305 | |||
| 306 | // KV-cache append (FP16): copy K_new, V_new into rows [cur_len, cur_len+L_new) | ||
| 307 | // of K_cache, V_cache. | ||
| 308 | // K_new, V_new: (L_new,D) FP16. K_cache, V_cache: (L_max,D) FP16 — must be | ||
| 309 | // pre-sized (not resized); cur_len + L_new <= L_max. | ||
| 310 | void kv_cache_append(const Tensor& K_new, const Tensor& V_new, | ||
| 311 | int cur_len, Tensor& K_cache, Tensor& V_cache); | ||
| 312 | |||
| 313 | |||
| 314 | // Causal flash-attention against a partially-filled KV cache (FP16, fwd-only). | ||
| 315 | // Runs the tiled core against rows [0, valid_len) of the caches. Query position | ||
| 316 | // p_q = (valid_len - L_q) + i attends to cache positions [0, p_q]. | ||
| 317 | // Q: (L_q, num_q_heads*head_dim) FP16 — L_q == 1 for token-by-token decode, | ||
| 318 | // L_q > 1 supported. | ||
| 319 | // K_cache, V_cache: (L_max, num_kv_heads*head_dim) FP16 — only rows | ||
| 320 | // [0, valid_len) are read. | ||
| 321 | // valid_len >= L_q. num_kv_heads must divide num_q_heads (GQA); each KV head | ||
| 322 | // serves num_q_heads/num_kv_heads consecutive query heads. num_kv_heads == | ||
| 323 | // num_q_heads is plain MHA. head_dim = Q.cols/num_q_heads must equal | ||
| 324 | // K_cache.cols/num_kv_heads. | ||
| 325 | // O: (L_q, num_q_heads*head_dim) FP16, resized as needed. | ||
| 326 | // GQA is native on all three backends: the decode kernel maps query head h to KV | ||
| 327 | // head h/(num_q_heads/num_kv_heads), reading the n_kv-wide cache directly — no | ||
| 328 | // KV-head widening needed. num_kv_heads == num_q_heads is plain MHA. | ||
| 329 | // | ||
| 330 | // Gemma-2 extensions (both default to today's exact behaviour): | ||
| 331 | // attn_softcap > 0: tanh logit soft-capping. After the 1/sqrt(head_dim) scale | ||
| 332 | // and BEFORE the causal/window mask and the online-softmax max/exp, each | ||
| 333 | // raw score s is replaced by attn_softcap * tanh(s / attn_softcap). | ||
| 334 | // attn_softcap == 0 (default) disables it — bit-identical to before. | ||
| 335 | // window > 0: sliding-window causal masking. Query at absolute position p | ||
| 336 | // (p = valid_len - L_q + q) attends key j only when j <= p AND | ||
| 337 | // j > p - window, i.e. keys in [max(0, p-window+1), p]. window <= 0 | ||
| 338 | // (default) is unbounded causal — bit-identical to before. Composes | ||
| 339 | // multiplicatively with the causal mask. | ||
| 340 | void flash_attention_decode(const Tensor& Q, | ||
| 341 | const Tensor& K_cache, const Tensor& V_cache, | ||
| 342 | int valid_len, int num_q_heads, int num_kv_heads, | ||
| 343 | Tensor& O, | ||
| 344 | float attn_softcap = 0.0f, int window = 0); | ||
| 345 | |||
| 346 | |||
| 347 | // Single-token decode attention over a FIXED-CAPACITY masked KV cache — the | ||
| 348 | // CUDA-graph-capturable twin of flash_attention_decode. The kernel always | ||
| 349 | // sees the full (L_max, ·) cache buffers, so its launch shape never changes | ||
| 350 | // as generation advances; validity lives in `d_mask`, a device-resident | ||
| 351 | // length-L_max FP32 key mask (1 valid / 0 invalid) the caller updates | ||
| 352 | // between graph replays. Masked keys are dropped before the dot product and | ||
| 353 | // softmax (their weights underflow to exact zeros), so with | ||
| 354 | // mask = [1]*valid_len + [0]*(L_max-valid_len) the result is bit-identical | ||
| 355 | // to flash_attention_decode(Q, K, V, valid_len, ...) at L_q == 1. Fully | ||
| 356 | // masked key tiles are skipped without touching K/V. | ||
| 357 | // Q: (1, num_q_heads*head_dim) — single query row only. | ||
| 358 | // K_cache, V_cache: (L_max, num_kv_heads*head_dim); same GQA rules as | ||
| 359 | // flash_attention_decode. | ||
| 360 | // d_mask: device pointer, length L_max, FP32; must not be null and must | ||
| 361 | // mark at least one key valid. | ||
| 362 | // O: (1, num_q_heads*head_dim), resized as needed. | ||
| 363 | // Gemma-2 extensions match flash_attention_decode (see above): | ||
| 364 | // attn_softcap > 0: tanh logit soft-capping applied to each raw score before | ||
| 365 | // the mask and softmax. attn_softcap == 0 (default) disables it. | ||
| 366 | // window > 0: sliding-window masking. The query's absolute position p is the | ||
| 367 | // highest valid key index (the last 1 in d_mask); a key j is kept only | ||
| 368 | // when it is valid AND j > p - window, i.e. j in [max(0,p-window+1), p]. | ||
| 369 | // window <= 0 (default) keeps the full valid set — bit-identical to | ||
| 370 | // before. Combines multiplicatively with the validity mask. | ||
| 371 | void flash_attention_decode_masked(const Tensor& Q, | ||
| 372 | const Tensor& K_cache, | ||
| 373 | const Tensor& V_cache, | ||
| 374 | const float* d_mask, | ||
| 375 | int num_q_heads, int num_kv_heads, | ||
| 376 | Tensor& O, | ||
| 377 | float attn_softcap = 0.0f, int window = 0); | ||
| 378 | |||
| 379 | |||
| 380 | // Single-token decode attention over the first *d_len keys of a | ||
| 381 | // fixed-capacity KV cache, the length read on the device at run time. Like | ||
| 382 | // the masked op, the launch shape never changes as generation advances, so a | ||
| 383 | // CUDA graph captured once replays for every step (the caller writes the new | ||
| 384 | // length between replays); unlike it, the work divides the VALID keys, so a | ||
| 385 | // cache sized for a long generation costs nothing for its unused rows. | ||
| 386 | // Results agree with flash_attention_decode(.., *d_len, ..) to rounding (the | ||
| 387 | // split points differ, so not bit for bit). | ||
| 388 | // Q may carry R > 1 rows: consecutive positions whose K/V are already the | ||
| 389 | // last R of the *d_len valid rows, row r attending causally to the first | ||
| 390 | // *d_len - R + 1 + r keys (flash_attention_decode's L_q > 1 rule) — a block | ||
| 391 | // of draft tokens verified in one pass, each cache row read once for all R. | ||
| 392 | // On CUDA that needs the grouped kernel's shapes (head_dim 64/128/256, | ||
| 393 | // group 1/2/4/8); *d_len >= R. | ||
| 394 | // Q: (R, num_q_heads*head_dim); K_cache, V_cache: (L_max, num_kv_heads*head_dim) | ||
| 395 | // FP16/BF16 on CUDA (head_dim a multiple of 32, at most 256), FP32 on CPU. | ||
| 396 | // d_len: device INT32 (a raw host pointer on CPU), 0 <= *d_len <= L_max; | ||
| 397 | // a length of 0 writes zeros. | ||
| 398 | // O: (1, num_q_heads*head_dim), resized as needed. | ||
| 399 | void flash_attention_decode_len(const Tensor& Q, const Tensor& K_cache, const Tensor& V_cache, | ||
| 400 | const int32_t* d_len, int num_q_heads, int num_kv_heads, Tensor& O); | ||
| 401 | |||
| 402 | |||
| 403 | // Back-compat overload: num_kv_heads defaults to num_heads (plain MHA). | ||
| 404 | 2 | inline void flash_attention_decode(const Tensor& Q, | |
| 405 | const Tensor& K_cache, const Tensor& V_cache, | ||
| 406 | int valid_len, int num_heads, Tensor& O, | ||
| 407 | float attn_softcap = 0.0f, int window = 0) { | ||
| 408 | 4 | flash_attention_decode(Q, K_cache, V_cache, valid_len, | |
| 409 | 2 | num_heads, /*num_kv_heads=*/num_heads, O, | |
| 410 | 2 | attn_softcap, window); | |
| 411 | 2 | } | |
| 412 | |||
| 413 | |||
| 414 | // ─── W8A16 flash-attention variants ──────────────────────────────────────── | ||
| 415 | // | ||
| 416 | // Same composition as flash_attention_project_kv / | ||
| 417 | // flash_attention_q_with_kv_cached_forward / flash_attention_qkvo_forward, but | ||
| 418 | // every projection takes an INT8 weight + per-output-row FP32 scale instead of | ||
| 419 | // an FP16 weight. The attention core stays FP16 (activations are never | ||
| 420 | // quantised). Each W*_int8 is (D, in_dim) with a matching scale (D,1); biases | ||
| 421 | // stay FP16 (D,1), optional. Masks, causal flag, and num_heads match the FP16 | ||
| 422 | // versions. | ||
| 423 | void flash_attention_project_kv_int8w_fp16(const Tensor& ctx, | ||
| 424 | const Tensor& Wk_int8, | ||
| 425 | const Tensor& sk, | ||
| 426 | const Tensor* bk, | ||
| 427 | const Tensor& Wv_int8, | ||
| 428 | const Tensor& sv, | ||
| 429 | const Tensor* bv, | ||
| 430 | Tensor& K_out, | ||
| 431 | Tensor& V_out); | ||
| 432 | |||
| 433 | |||
| 434 | void flash_attention_q_with_kv_cached_int8w_fp16(const Tensor& X, | ||
| 435 | const Tensor& K, | ||
| 436 | const Tensor& V, | ||
| 437 | const Tensor& Wq_int8, | ||
| 438 | const Tensor& sq, | ||
| 439 | const Tensor* bq, | ||
| 440 | const Tensor& Wo_int8, | ||
| 441 | const Tensor& so, | ||
| 442 | const Tensor* bo, | ||
| 443 | const float* d_mask, | ||
| 444 | int num_heads, | ||
| 445 | bool causal, | ||
| 446 | Tensor& O); | ||
| 447 | |||
| 448 | |||
| 449 | void flash_attention_qkvo_int8w_fp16(const Tensor& X, | ||
| 450 | const Tensor* Ctx, | ||
| 451 | const Tensor& Wq_int8, const Tensor& sq, const Tensor* bq, | ||
| 452 | const Tensor& Wk_int8, const Tensor& sk, const Tensor* bk, | ||
| 453 | const Tensor& Wv_int8, const Tensor& sv, const Tensor* bv, | ||
| 454 | const Tensor& Wo_int8, const Tensor& so, const Tensor* bo, | ||
| 455 | const float* d_mask, | ||
| 456 | int num_heads, | ||
| 457 | bool causal, | ||
| 458 | Tensor& O); | ||
| 459 | |||
| 460 | } // namespace brotensor | ||
| 461 |