include/brotensor/ops/flash_attention.h - GCC Code Coverage Report

Directory:	./
Coverage:	low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%

	Coverage	Exec / Excl / Total
Lines:	100.0%	5 / 0 / 5
Functions:	100.0%	1 / 0 / 1
Branches:	-%	0 / 0 / 0

    include/brotensor/ops/flash_attention.h
    
        Line
        Branch
        Exec
        Source
      
        #pragma once
      
        // brotensor ops/flash_attention.h — Flash-attention family: tiled/windowed/varlen/qkvo/decode + KV cache.
      
        #include "../tensor.h"
      
        #include <cstdint>
      
        namespace brotensor {
      
        // Flash-attention-style fused attention (FP16, inference-only). Q, K, V are
      
        // already projected. Tiled online softmax over Lk (no Lk-long materialisation);
      
        // FP32 accumulation.
      
        //   Q: (Lq,D); K, V: (Lk,D) — FP16.
      
        //   d_mask: optional length-Lk FP32 mask (1 valid / 0 invalid); may be null.
      
        //   num_heads divides D.
      
        //   causal: if true, key k attends to query q only when k <= q (requires
      
        //           Lq == Lk); combines multiplicatively with d_mask.
      
        //   O: (Lq,D) FP16, resized as needed.
      
        void flash_attention_forward(const Tensor& Q,
      
                                     const Tensor& K,
      
                                     const Tensor& V,
      
                                     const float* d_mask,
      
                                     int num_heads,
      
                                     bool causal,
      
                                     Tensor& O);
      
        // GQA generalisation of flash_attention_forward: grouped-query self-attention
      
        // over pre-projected Q/K/V, causal OR bidirectional. Q carries num_q_heads,
      
        // K/V carry num_kv_heads (num_kv_heads must divide num_q_heads; equal == plain
      
        // MHA). causal == false is full bidirectional attention over all Lk keys — the
      
        // encoder prefill a bidirectional-ified decoder needs (LLM2Vec). causal == true
      
        // reproduces flash_attention_forward's causal mask with GQA (requires Lq == Lk).
      
        // Tiled online softmax, FP16/BF16/FP32, FP32 accumulation; shares the windowed
      
        // kernel's GQA + key-mask machinery.
      
        //   Q: (Lq, num_q_heads*head_dim);  K, V: (Lk, num_kv_heads*head_dim).
      
        //   d_mask: optional length-Lk FP32 key mask (1 valid / 0 invalid); may be null.
      
        //   num_kv_heads divides num_q_heads; head_dim = Q.cols/num_q_heads must equal
      
        //     K.cols/num_kv_heads.  Lk >= Lq; causal additionally requires Lq == Lk.
      
        //   O: (Lq, num_q_heads*head_dim), same dtype as Q, resized as needed.
      
        void flash_attention_gqa_forward(const Tensor& Q,
      
                                         const Tensor& K,
      
                                         const Tensor& V,
      
                                         const float* d_mask,
      
                                         int num_q_heads,
      
                                         int num_kv_heads,
      
                                         bool causal,
      
                                         Tensor& O);
      
        // Sliding-window causal self-attention (FP32, inference-only) — the local
      
        // attention of streaming neural codecs (e.g. Qwen3-TTS / Mimi) and the
      
        // autoregressive decode step. Q, K, V already projected, (L, num_heads*head_dim).
      
        // Always causal. The Lq queries occupy the last Lq positions of a length-Lk
      
        // causal sequence (q_offset = Lk - Lq): query row r is at absolute position
      
        // r + q_offset and attends keys [max(0, pos-window+1), pos]. window <= 0 means
      
        // unbounded causal — identical to flash_attention_forward with causal=true.
      
        //   - Lq == Lk: self-attention (prefill / codec sliding window).
      
        //   - Lq  < Lk: incremental decode of an Lq-token block over a K/V cache;
      
        //     Lq == 1 with window <= 0 attends every cached key (full cache attention),
      
        //     replacing a varlen call with no cu_seqlens upload. Requires Lk >= Lq.
      
        //   d_mask: optional length-Lk FP32 key mask (1 valid / 0 invalid), combined
      
        //           multiplicatively with the window; may be null.
      
        //   num_heads divides D.  O: (Lq, num_heads*head_dim), resized as needed.
      
        void flash_attention_windowed_forward(const Tensor& Q,
      
                                              const Tensor& K,
      
                                              const Tensor& V,
      
                                              const float* d_mask,
      
                                              int num_heads,
      
                                              int window,
      
                                              Tensor& O);
      
        // Packed variable-length multi-head attention, forward only (Qwen3-VL window
      
        // attention). All sequences in a batch live contiguously in one packed tensor;
      
        // per-sequence boundaries come from `cu_seqlens_*` INT32 prefix-sum buffers of
      
        // length `batch_size + 1`. Sequence b covers Q rows
      
        // [cu_seqlens_q[b], cu_seqlens_q[b+1]) and attends to K/V rows
      
        // [cu_seqlens_k[b], cu_seqlens_k[b+1]); no cross-sequence attention. Mirrors
      
        // flash_attn_varlen_func semantics.
      
        //   Q: (total_tokens_q, num_heads * head_dim) — FP16/BF16/FP32 (GPU) / FP32 (CPU).
      
        //   K, V: (total_tokens_k, num_heads * head_dim) — same packing, num_heads
      
        //         matches Q (no GQA — that's a future cleanup).
      
        //   cu_seqlens_q, cu_seqlens_k: DEVICE pointers (CUDA/Metal device, raw host
      
        //         pointers on CPU), length batch_size + 1, INT32 prefix sums. Same
      
        //         convention as `const float* d_mask` elsewhere — caller owns the
      
        //         buffer; the op does not allocate or copy.
      
        //   max_seqlen_q, max_seqlen_k: bound the longest sequence's length — used by
      
        //         the GPU kernel for block sizing; the CPU impl ignores them.
      
        //   causal: if true, key k in sequence b attends to query q only when
      
        //         (k - cu_seqlens_k[b]) <= (q - cu_seqlens_q[b]); only meaningful when
      
        //         the per-sequence Q and K lengths match.
      
        //   O: (total_tokens_q, num_heads * head_dim) — overwritten, resized + dtype-set.
      
        void flash_attention_varlen_forward(const Tensor& Q,
      
                                            const Tensor& K,
      
                                            const Tensor& V,
      
                                            const int32_t* cu_seqlens_q,
      
                                            const int32_t* cu_seqlens_k,
      
                                            int batch_size,
      
                                            int max_seqlen_q,
      
                                            int max_seqlen_k,
      
                                            int num_heads,
      
                                            int head_dim,
      
                                            bool causal,
      
                                            Tensor& O);
      
        // Backward of flash_attention_varlen_forward — packed variable-length attention
      
        // over pre-projected Q/K/V (no projection weights, no biases — projections are
      
        // handled by the caller's linear layer; bias gradients belong to that layer).
      
        // Recompute-based: consumes no forward caches, re-runs the per-sequence
      
        // softmax then reverses it. Matches flash_attention_backward's per-sequence
      
        // math; only the cu_seqlens scatter/gather and per-sequence causal differ.
      
        //   Q: (total_tokens_q, num_heads*head_dim) — FP16/BF16/FP32 (GPU) / FP32 (CPU).
      
        //   K, V: (total_tokens_k, num_heads*head_dim) — same dtype, same packing.
      
        //   O:  (total_tokens_q, ...) forward output — currently unused (kept for API
      
        //       symmetry with flash_attention_backward).
      
        //   dO: (total_tokens_q, ...) upstream gradient, same dtype as Q/K/V.
      
        //   cu_seqlens_q, cu_seqlens_k: DEVICE pointers on GPU, host pointers on CPU,
      
        //       length batch_size + 1, INT32 prefix sums. Same convention as the
      
        //       forward.
      
        //   max_seqlen_q, max_seqlen_k: bound the longest per-sequence length — used
      
        //       by the GPU kernel for block/scratch sizing; the CPU impl ignores them.
      
        //   causal: per-sequence; only meaningful when per-sequence Lq == Lk.
      
        //   dQ: (total_tokens_q, ...); dK, dV: (total_tokens_k, ...) — OVERWRITTEN
      
        //       (resized + dtype-set), matching flash_attention_backward's contract.
      
        // Out-of-range positions (causal-excluded, or in a fully empty K sequence)
      
        // contribute nothing to any gradient. No cross-sequence attention.
      
        void flash_attention_varlen_backward(const Tensor& Q,
      
                                             const Tensor& K,
      
                                             const Tensor& V,
      
                                             const Tensor& O,
      
                                             const Tensor& dO,
      
                                             const int32_t* cu_seqlens_q,
      
                                             const int32_t* cu_seqlens_k,
      
                                             int batch_size,
      
                                             int max_seqlen_q,
      
                                             int max_seqlen_k,
      
                                             int num_heads,
      
                                             int head_dim,
      
                                             bool causal,
      
                                             Tensor& dQ,
      
                                             Tensor& dK,
      
                                             Tensor& dV);
      
        // Flash-attention with QKV and output projections fused at the boundary.
      
        // Projects X->Q, Ctx->K,V (or X->Q,K,V when Ctx is null), runs the tiled core,
      
        // then projects with Wo. FP16 throughout.
      
        //   X:   (Lq,D) FP16, query source.
      
        //   Ctx: (Lk,D_ctx) FP16 or null. Null => self-attention (Ctx<-X). D_ctx may
      
        //        differ from D (e.g. SD1.5 cross-attention).
      
        //   Wq, Wo: (D,D); Wk, Wv: (D,D_ctx) — FP16.
      
        //   bq, bk, bv, bo: optional (D,1) FP16 biases; null to skip.
      
        //   d_mask: optional length-Lk FP32 mask.  num_heads divides D.
      
        //   causal: see flash_attention_forward (typically with Ctx == null).
      
        //   O: (Lq,D) FP16, resized as needed.
      
        void flash_attention_qkvo_forward(const Tensor& X,
      
                                          const Tensor* Ctx,
      
                                          const Tensor& Wq, const Tensor* bq,
      
                                          const Tensor& Wk, const Tensor* bk,
      
                                          const Tensor& Wv, const Tensor* bv,
      
                                          const Tensor& Wo, const Tensor* bo,
      
                                          const float* d_mask,
      
                                          int num_heads,
      
                                          bool causal,
      
                                          Tensor& O);
      
        // Backward of flash_attention_qkvo_forward. Recompute-style: consumes no
      
        // forward caches — re-runs the attention math from the inputs, then reverses
      
        // it. FP16 storage, FP32 accumulation. All shape / dtype / Ctx-null /
      
        // rectangular-Wk-Wv / causal / optional-bias rules match the forward; pass the
      
        // same values.
      
        //   dO:   (Lq,D) FP16 upstream.
      
        //   dX:   (Lq,D) FP16 overwritten. For self-attention (Ctx null) dX absorbs
      
        //         the K/V-projection gradients too: dX = dQ.Wq + dK.Wk + dV.Wv.
      
        //   dCtx: (Lk,D_ctx) FP16 overwritten; must be null iff Ctx is null. For
      
        //         cross-attention dCtx = dK.Wk + dV.Wv.
      
        //   dWq, dWo: (D,D); dWk, dWv: (D,D_ctx) — FP16, accumulated (caller zeros).
      
        //   dbq, dbk, dbv, dbo: (D,1) FP16, accumulated iff the matching forward bias
      
        //         was non-null; pass null otherwise (the null/non-null symmetry must
      
        //         be exact — a mismatch is rejected).
      
        // Causal- and mask-excluded positions contribute nothing to any gradient.
      
        void flash_attention_qkvo_backward(
      
            const Tensor& X, const Tensor* Ctx,
      
            const Tensor& Wq, const Tensor* bq,
      
            const Tensor& Wk, const Tensor* bk,
      
            const Tensor& Wv, const Tensor* bv,
      
            const Tensor& Wo, const Tensor* bo,
      
            const float* d_mask,
      
            int num_heads,
      
            bool causal,
      
            const Tensor& dO,
      
            Tensor& dX, Tensor* dCtx,
      
            Tensor& dWq, Tensor* dbq,
      
            Tensor& dWk, Tensor* dbk,
      
            Tensor& dWv, Tensor* dbv,
      
            Tensor& dWo, Tensor* dbo);
      
        // Backward of flash_attention_forward — bare attention core, no projection
      
        // weights (what LoRA-style adapters need; projections are wrapped externally).
      
        // Recompute-based; FP16 storage, FP32 accumulation.
      
        //   Q: (Lq,D); K, V: (Lk,D) — pre-projected forward inputs, FP16.
      
        //   O: (Lq,D) forward output — currently unused (kept for API symmetry).
      
        //   dO: (Lq,D) FP16 upstream.
      
        //   d_mask: optional length-Lk FP32 mask (null for unmasked); positions with
      
        //           mask[k] <= 0.5 are dropped.
      
        //   num_heads divides D.  causal: match the forward (requires Lq == Lk).
      
        //   dQ: (Lq,D); dK, dV: (Lk,D) — FP16, overwritten (resized + dtype-set).
      
        // Causal- and mask-excluded positions contribute nothing to dQ/dK/dV.
      
        void flash_attention_backward(const Tensor& Q,
      
                                      const Tensor& K,
      
                                      const Tensor& V,
      
                                      const Tensor& O,
      
                                      const Tensor& dO,
      
                                      const float* d_mask,
      
                                      int num_heads,
      
                                      bool causal,
      
                                      Tensor& dQ,
      
                                      Tensor& dK,
      
                                      Tensor& dV);
      
        // Project a context tensor through Wk/Wv into the exact (Lk,D) FP16 K/V buffers
      
        // flash_attention_forward consumes. Used to precompute cross-attention K/V once
      
        // per generate() (text context is fixed across denoising steps). Numerically
      
        // identical to the K/V projection stage of flash_attention_qkvo_forward.
      
        //   ctx: (Lk,D_ctx) FP16.  Wk, Wv: (D,D_ctx) FP16.  bk, bv: (D,1) FP16 or null.
      
        //   K_out, V_out: (Lk,D) FP16, resized as needed.
      
        void flash_attention_project_kv(const Tensor& ctx,
      
                                        const Tensor& Wk, const Tensor* bk,
      
                                        const Tensor& Wv, const Tensor* bv,
      
                                        Tensor& K_out,
      
                                        Tensor& V_out);
      
        // Like flash_attention_qkvo_forward but K and V are already projected by the
      
        // caller (typically via flash_attention_project_kv). Projects X->Q with Wq/bq,
      
        // runs the tiled core against the supplied K/V, applies Wo/bo. FP16 throughout.
      
        //   X: (Lq,D); K, V: (Lk,D) — FP16.  Wq, Wo: (D,D) FP16; bq, bo: (D,1) FP16 or null.
      
        //   d_mask: optional length-Lk FP32 mask.  num_heads divides D.
      
        //   causal: see flash_attention_forward.  O: (Lq,D) FP16, resized as needed.
      
        void flash_attention_q_with_kv_cached_forward(const Tensor& X,
      
                                                      const Tensor& K,
      
                                                      const Tensor& V,
      
                                                      const Tensor& Wq, const Tensor* bq,
      
                                                      const Tensor& Wo, const Tensor* bo,
      
                                                      const float* d_mask,
      
                                                      int num_heads,
      
                                                      bool causal,
      
                                                      Tensor& O);
      
        // KV-cache append (FP16): copy K_new, V_new into rows [cur_len, cur_len+L_new)
      
        // of K_cache, V_cache.
      
        //   K_new, V_new: (L_new,D) FP16.  K_cache, V_cache: (L_max,D) FP16 — must be
      
        //   pre-sized (not resized); cur_len + L_new <= L_max.
      
        void kv_cache_append(const Tensor& K_new, const Tensor& V_new,
      
                             int cur_len, Tensor& K_cache, Tensor& V_cache);
      
        // Causal flash-attention against a partially-filled KV cache (FP16, fwd-only).
      
        // Runs the tiled core against rows [0, valid_len) of the caches. Query position
      
        // p_q = (valid_len - L_q) + i attends to cache positions [0, p_q].
      
        //   Q: (L_q, num_q_heads*head_dim) FP16 — L_q == 1 for token-by-token decode,
      
        //       L_q > 1 supported.
      
        //   K_cache, V_cache: (L_max, num_kv_heads*head_dim) FP16 — only rows
      
        //       [0, valid_len) are read.
      
        //   valid_len >= L_q.  num_kv_heads must divide num_q_heads (GQA); each KV head
      
        //       serves num_q_heads/num_kv_heads consecutive query heads. num_kv_heads ==
      
        //       num_q_heads is plain MHA. head_dim = Q.cols/num_q_heads must equal
      
        //       K_cache.cols/num_kv_heads.
      
        //   O: (L_q, num_q_heads*head_dim) FP16, resized as needed.
      
        // GQA is native on all three backends: the decode kernel maps query head h to KV
      
        // head h/(num_q_heads/num_kv_heads), reading the n_kv-wide cache directly — no
      
        // KV-head widening needed. num_kv_heads == num_q_heads is plain MHA.
      
        //
      
        // Gemma-2 extensions (both default to today's exact behaviour):
      
        //   attn_softcap > 0: tanh logit soft-capping. After the 1/sqrt(head_dim) scale
      
        //       and BEFORE the causal/window mask and the online-softmax max/exp, each
      
        //       raw score s is replaced by attn_softcap * tanh(s / attn_softcap).
      
        //       attn_softcap == 0 (default) disables it — bit-identical to before.
      
        //   window > 0: sliding-window causal masking. Query at absolute position p
      
        //       (p = valid_len - L_q + q) attends key j only when j <= p AND
      
        //       j > p - window, i.e. keys in [max(0, p-window+1), p]. window <= 0
      
        //       (default) is unbounded causal — bit-identical to before. Composes
      
        //       multiplicatively with the causal mask.
      
        void flash_attention_decode(const Tensor& Q,
      
                                   const Tensor& K_cache, const Tensor& V_cache,
      
                                   int valid_len, int num_q_heads, int num_kv_heads,
      
                                   Tensor& O,
      
                                   float attn_softcap = 0.0f, int window = 0);
      
        // Single-token decode attention over a FIXED-CAPACITY masked KV cache — the
      
        // CUDA-graph-capturable twin of flash_attention_decode. The kernel always
      
        // sees the full (L_max, ·) cache buffers, so its launch shape never changes
      
        // as generation advances; validity lives in `d_mask`, a device-resident
      
        // length-L_max FP32 key mask (1 valid / 0 invalid) the caller updates
      
        // between graph replays. Masked keys are dropped before the dot product and
      
        // softmax (their weights underflow to exact zeros), so with
      
        // mask = [1]*valid_len + [0]*(L_max-valid_len) the result is bit-identical
      
        // to flash_attention_decode(Q, K, V, valid_len, ...) at L_q == 1. Fully
      
        // masked key tiles are skipped without touching K/V.
      
        //   Q: (1, num_q_heads*head_dim) — single query row only.
      
        //   K_cache, V_cache: (L_max, num_kv_heads*head_dim); same GQA rules as
      
        //       flash_attention_decode.
      
        //   d_mask: device pointer, length L_max, FP32; must not be null and must
      
        //       mark at least one key valid.
      
        //   O: (1, num_q_heads*head_dim), resized as needed.
      
        // Gemma-2 extensions match flash_attention_decode (see above):
      
        //   attn_softcap > 0: tanh logit soft-capping applied to each raw score before
      
        //       the mask and softmax. attn_softcap == 0 (default) disables it.
      
        //   window > 0: sliding-window masking. The query's absolute position p is the
      
        //       highest valid key index (the last 1 in d_mask); a key j is kept only
      
        //       when it is valid AND j > p - window, i.e. j in [max(0,p-window+1), p].
      
        //       window <= 0 (default) keeps the full valid set — bit-identical to
      
        //       before. Combines multiplicatively with the validity mask.
      
        void flash_attention_decode_masked(const Tensor& Q,
      
                                           const Tensor& K_cache,
      
                                           const Tensor& V_cache,
      
                                           const float* d_mask,
      
                                           int num_q_heads, int num_kv_heads,
      
                                           Tensor& O,
      
                                           float attn_softcap = 0.0f, int window = 0);
      
        // Back-compat overload: num_kv_heads defaults to num_heads (plain MHA).
      
        1
        inline void flash_attention_decode(const Tensor& Q,
      
                                           const Tensor& K_cache, const Tensor& V_cache,
      
                                           int valid_len, int num_heads, Tensor& O,
      
                                           float attn_softcap = 0.0f, int window = 0) {
      
        2
            flash_attention_decode(Q, K_cache, V_cache, valid_len,
      
        1
                                   num_heads, /*num_kv_heads=*/num_heads, O,
      
        1
                                   attn_softcap, window);
      
        1
        }
      
        // ─── W8A16 flash-attention variants ────────────────────────────────────────
      
        //
      
        // Same composition as flash_attention_project_kv /
      
        // flash_attention_q_with_kv_cached_forward / flash_attention_qkvo_forward, but
      
        // every projection takes an INT8 weight + per-output-row FP32 scale instead of
      
        // an FP16 weight. The attention core stays FP16 (activations are never
      
        // quantised). Each W*_int8 is (D, in_dim) with a matching scale (D,1); biases
      
        // stay FP16 (D,1), optional. Masks, causal flag, and num_heads match the FP16
      
        // versions.
      
        void flash_attention_project_kv_int8w_fp16(const Tensor& ctx,
      
                                                   const Tensor& Wk_int8,
      
                                                   const Tensor& sk,
      
                                                   const Tensor* bk,
      
                                                   const Tensor& Wv_int8,
      
                                                   const Tensor& sv,
      
                                                   const Tensor* bv,
      
                                                   Tensor& K_out,
      
                                                   Tensor& V_out);
      
        void flash_attention_q_with_kv_cached_int8w_fp16(const Tensor& X,
      
                                                         const Tensor& K,
      
                                                         const Tensor& V,
      
                                                         const Tensor& Wq_int8,
      
                                                         const Tensor& sq,
      
                                                         const Tensor* bq,
      
                                                         const Tensor& Wo_int8,
      
                                                         const Tensor& so,
      
                                                         const Tensor* bo,
      
                                                         const float* d_mask,
      
                                                         int num_heads,
      
                                                         bool causal,
      
                                                         Tensor& O);
      
        void flash_attention_qkvo_int8w_fp16(const Tensor& X,
      
                                             const Tensor* Ctx,
      
                                             const Tensor& Wq_int8, const Tensor& sq, const Tensor* bq,
      
                                             const Tensor& Wk_int8, const Tensor& sk, const Tensor* bk,
      
                                             const Tensor& Wv_int8, const Tensor& sv, const Tensor* bv,
      
                                             const Tensor& Wo_int8, const Tensor& so, const Tensor* bo,
      
                                             const float* d_mask,
      
                                             int num_heads,
      
                                             bool causal,
      
                                             Tensor& O);
      
        }  // namespace brotensor

Line	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_headshead_dim); K, V: (Lk, num_kv_headshead_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 causal self-attention (FP32, inference-only) — the local
53		// attention of streaming neural codecs (e.g. Qwen3-TTS / Mimi) and the
54		// autoregressive decode step. Q, K, V already projected, (L, num_heads*head_dim).
55		// Always causal. The Lq queries occupy the last Lq positions of a length-Lk
56		// causal sequence (q_offset = Lk - Lq): query row r is at absolute position
57		// r + q_offset and attends keys [max(0, pos-window+1), pos]. window <= 0 means
58		// unbounded causal — identical to flash_attention_forward with causal=true.
59		// - Lq == Lk: self-attention (prefill / codec sliding window).
60		// - Lq < Lk: incremental decode of an Lq-token block over a K/V cache;
61		// Lq == 1 with window <= 0 attends every cached key (full cache attention),
62		// replacing a varlen call with no cu_seqlens upload. Requires Lk >= Lq.
63		// d_mask: optional length-Lk FP32 key mask (1 valid / 0 invalid), combined
64		// multiplicatively with the window; may be null.
65		// num_heads divides D. O: (Lq, num_heads*head_dim), resized as needed.
66		void flash_attention_windowed_forward(const Tensor& Q,
67		const Tensor& K,
68		const Tensor& V,
69		const float* d_mask,
70		int num_heads,
71		int window,
72		Tensor& O);
73
74
75		// Packed variable-length multi-head attention, forward only (Qwen3-VL window
76		// attention). All sequences in a batch live contiguously in one packed tensor;
77		// per-sequence boundaries come from `cu_seqlens_*` INT32 prefix-sum buffers of
78		// length `batch_size + 1`. Sequence b covers Q rows
79		// [cu_seqlens_q[b], cu_seqlens_q[b+1]) and attends to K/V rows
80		// [cu_seqlens_k[b], cu_seqlens_k[b+1]); no cross-sequence attention. Mirrors
81		// flash_attn_varlen_func semantics.
82		// Q: (total_tokens_q, num_heads * head_dim) — FP16/BF16/FP32 (GPU) / FP32 (CPU).
83		// K, V: (total_tokens_k, num_heads * head_dim) — same packing, num_heads
84		// matches Q (no GQA — that's a future cleanup).
85		// cu_seqlens_q, cu_seqlens_k: DEVICE pointers (CUDA/Metal device, raw host
86		// pointers on CPU), length batch_size + 1, INT32 prefix sums. Same
87		// convention as `const float* d_mask` elsewhere — caller owns the
88		// buffer; the op does not allocate or copy.
89		// max_seqlen_q, max_seqlen_k: bound the longest sequence's length — used by
90		// the GPU kernel for block sizing; the CPU impl ignores them.
91		// causal: if true, key k in sequence b attends to query q only when
92		// (k - cu_seqlens_k[b]) <= (q - cu_seqlens_q[b]); only meaningful when
93		// the per-sequence Q and K lengths match.
94		// O: (total_tokens_q, num_heads * head_dim) — overwritten, resized + dtype-set.
95		void flash_attention_varlen_forward(const Tensor& Q,
96		const Tensor& K,
97		const Tensor& V,
98		const int32_t* cu_seqlens_q,
99		const int32_t* cu_seqlens_k,
100		int batch_size,
101		int max_seqlen_q,
102		int max_seqlen_k,
103		int num_heads,
104		int head_dim,
105		bool causal,
106		Tensor& O);
107
108
109		// Backward of flash_attention_varlen_forward — packed variable-length attention
110		// over pre-projected Q/K/V (no projection weights, no biases — projections are
111		// handled by the caller's linear layer; bias gradients belong to that layer).
112		// Recompute-based: consumes no forward caches, re-runs the per-sequence
113		// softmax then reverses it. Matches flash_attention_backward's per-sequence
114		// math; only the cu_seqlens scatter/gather and per-sequence causal differ.
115		// Q: (total_tokens_q, num_heads*head_dim) — FP16/BF16/FP32 (GPU) / FP32 (CPU).
116		// K, V: (total_tokens_k, num_heads*head_dim) — same dtype, same packing.
117		// O: (total_tokens_q, ...) forward output — currently unused (kept for API
118		// symmetry with flash_attention_backward).
119		// dO: (total_tokens_q, ...) upstream gradient, same dtype as Q/K/V.
120		// cu_seqlens_q, cu_seqlens_k: DEVICE pointers on GPU, host pointers on CPU,
121		// length batch_size + 1, INT32 prefix sums. Same convention as the
122		// forward.
123		// max_seqlen_q, max_seqlen_k: bound the longest per-sequence length — used
124		// by the GPU kernel for block/scratch sizing; the CPU impl ignores them.
125		// causal: per-sequence; only meaningful when per-sequence Lq == Lk.
126		// dQ: (total_tokens_q, ...); dK, dV: (total_tokens_k, ...) — OVERWRITTEN
127		// (resized + dtype-set), matching flash_attention_backward's contract.
128		// Out-of-range positions (causal-excluded, or in a fully empty K sequence)
129		// contribute nothing to any gradient. No cross-sequence attention.
130		void flash_attention_varlen_backward(const Tensor& Q,
131		const Tensor& K,
132		const Tensor& V,
133		const Tensor& O,
134		const Tensor& dO,
135		const int32_t* cu_seqlens_q,
136		const int32_t* cu_seqlens_k,
137		int batch_size,
138		int max_seqlen_q,
139		int max_seqlen_k,
140		int num_heads,
141		int head_dim,
142		bool causal,
143		Tensor& dQ,
144		Tensor& dK,
145		Tensor& dV);
146
147
148		// Flash-attention with QKV and output projections fused at the boundary.
149		// Projects X->Q, Ctx->K,V (or X->Q,K,V when Ctx is null), runs the tiled core,
150		// then projects with Wo. FP16 throughout.
151		// X: (Lq,D) FP16, query source.
152		// Ctx: (Lk,D_ctx) FP16 or null. Null => self-attention (Ctx<-X). D_ctx may
153		// differ from D (e.g. SD1.5 cross-attention).
154		// Wq, Wo: (D,D); Wk, Wv: (D,D_ctx) — FP16.
155		// bq, bk, bv, bo: optional (D,1) FP16 biases; null to skip.
156		// d_mask: optional length-Lk FP32 mask. num_heads divides D.
157		// causal: see flash_attention_forward (typically with Ctx == null).
158		// O: (Lq,D) FP16, resized as needed.
159		void flash_attention_qkvo_forward(const Tensor& X,
160		const Tensor* Ctx,
161		const Tensor& Wq, const Tensor* bq,
162		const Tensor& Wk, const Tensor* bk,
163		const Tensor& Wv, const Tensor* bv,
164		const Tensor& Wo, const Tensor* bo,
165		const float* d_mask,
166		int num_heads,
167		bool causal,
168		Tensor& O);
169
170
171		// Backward of flash_attention_qkvo_forward. Recompute-style: consumes no
172		// forward caches — re-runs the attention math from the inputs, then reverses
173		// it. FP16 storage, FP32 accumulation. All shape / dtype / Ctx-null /
174		// rectangular-Wk-Wv / causal / optional-bias rules match the forward; pass the
175		// same values.
176		// dO: (Lq,D) FP16 upstream.
177		// dX: (Lq,D) FP16 overwritten. For self-attention (Ctx null) dX absorbs
178		// the K/V-projection gradients too: dX = dQ.Wq + dK.Wk + dV.Wv.
179		// dCtx: (Lk,D_ctx) FP16 overwritten; must be null iff Ctx is null. For
180		// cross-attention dCtx = dK.Wk + dV.Wv.
181		// dWq, dWo: (D,D); dWk, dWv: (D,D_ctx) — FP16, accumulated (caller zeros).
182		// dbq, dbk, dbv, dbo: (D,1) FP16, accumulated iff the matching forward bias
183		// was non-null; pass null otherwise (the null/non-null symmetry must
184		// be exact — a mismatch is rejected).
185		// Causal- and mask-excluded positions contribute nothing to any gradient.
186		void flash_attention_qkvo_backward(
187		const Tensor& X, const Tensor* Ctx,
188		const Tensor& Wq, const Tensor* bq,
189		const Tensor& Wk, const Tensor* bk,
190		const Tensor& Wv, const Tensor* bv,
191		const Tensor& Wo, const Tensor* bo,
192		const float* d_mask,
193		int num_heads,
194		bool causal,
195		const Tensor& dO,
196		Tensor& dX, Tensor* dCtx,
197		Tensor& dWq, Tensor* dbq,
198		Tensor& dWk, Tensor* dbk,
199		Tensor& dWv, Tensor* dbv,
200		Tensor& dWo, Tensor* dbo);
201
202
203		// Backward of flash_attention_forward — bare attention core, no projection
204		// weights (what LoRA-style adapters need; projections are wrapped externally).
205		// Recompute-based; FP16 storage, FP32 accumulation.
206		// Q: (Lq,D); K, V: (Lk,D) — pre-projected forward inputs, FP16.
207		// O: (Lq,D) forward output — currently unused (kept for API symmetry).
208		// dO: (Lq,D) FP16 upstream.
209		// d_mask: optional length-Lk FP32 mask (null for unmasked); positions with
210		// mask[k] <= 0.5 are dropped.
211		// num_heads divides D. causal: match the forward (requires Lq == Lk).
212		// dQ: (Lq,D); dK, dV: (Lk,D) — FP16, overwritten (resized + dtype-set).
213		// Causal- and mask-excluded positions contribute nothing to dQ/dK/dV.
214		void flash_attention_backward(const Tensor& Q,
215		const Tensor& K,
216		const Tensor& V,
217		const Tensor& O,
218		const Tensor& dO,
219		const float* d_mask,
220		int num_heads,
221		bool causal,
222		Tensor& dQ,
223		Tensor& dK,
224		Tensor& dV);
225
226
227		// Project a context tensor through Wk/Wv into the exact (Lk,D) FP16 K/V buffers
228		// flash_attention_forward consumes. Used to precompute cross-attention K/V once
229		// per generate() (text context is fixed across denoising steps). Numerically
230		// identical to the K/V projection stage of flash_attention_qkvo_forward.
231		// ctx: (Lk,D_ctx) FP16. Wk, Wv: (D,D_ctx) FP16. bk, bv: (D,1) FP16 or null.
232		// K_out, V_out: (Lk,D) FP16, resized as needed.
233		void flash_attention_project_kv(const Tensor& ctx,
234		const Tensor& Wk, const Tensor* bk,
235		const Tensor& Wv, const Tensor* bv,
236		Tensor& K_out,
237		Tensor& V_out);
238
239
240		// Like flash_attention_qkvo_forward but K and V are already projected by the
241		// caller (typically via flash_attention_project_kv). Projects X->Q with Wq/bq,
242		// runs the tiled core against the supplied K/V, applies Wo/bo. FP16 throughout.
243		// X: (Lq,D); K, V: (Lk,D) — FP16. Wq, Wo: (D,D) FP16; bq, bo: (D,1) FP16 or null.
244		// d_mask: optional length-Lk FP32 mask. num_heads divides D.
245		// causal: see flash_attention_forward. O: (Lq,D) FP16, resized as needed.
246		void flash_attention_q_with_kv_cached_forward(const Tensor& X,
247		const Tensor& K,
248		const Tensor& V,
249		const Tensor& Wq, const Tensor* bq,
250		const Tensor& Wo, const Tensor* bo,
251		const float* d_mask,
252		int num_heads,
253		bool causal,
254		Tensor& O);
255
256
257		// KV-cache append (FP16): copy K_new, V_new into rows [cur_len, cur_len+L_new)
258		// of K_cache, V_cache.
259		// K_new, V_new: (L_new,D) FP16. K_cache, V_cache: (L_max,D) FP16 — must be
260		// pre-sized (not resized); cur_len + L_new <= L_max.
261		void kv_cache_append(const Tensor& K_new, const Tensor& V_new,
262		int cur_len, Tensor& K_cache, Tensor& V_cache);
263
264
265		// Causal flash-attention against a partially-filled KV cache (FP16, fwd-only).
266		// Runs the tiled core against rows [0, valid_len) of the caches. Query position
267		// p_q = (valid_len - L_q) + i attends to cache positions [0, p_q].
268		// Q: (L_q, num_q_heads*head_dim) FP16 — L_q == 1 for token-by-token decode,
269		// L_q > 1 supported.
270		// K_cache, V_cache: (L_max, num_kv_heads*head_dim) FP16 — only rows
271		// [0, valid_len) are read.
272		// valid_len >= L_q. num_kv_heads must divide num_q_heads (GQA); each KV head
273		// serves num_q_heads/num_kv_heads consecutive query heads. num_kv_heads ==
274		// num_q_heads is plain MHA. head_dim = Q.cols/num_q_heads must equal
275		// K_cache.cols/num_kv_heads.
276		// O: (L_q, num_q_heads*head_dim) FP16, resized as needed.
277		// GQA is native on all three backends: the decode kernel maps query head h to KV
278		// head h/(num_q_heads/num_kv_heads), reading the n_kv-wide cache directly — no
279		// KV-head widening needed. num_kv_heads == num_q_heads is plain MHA.
280		//
281		// Gemma-2 extensions (both default to today's exact behaviour):
282		// attn_softcap > 0: tanh logit soft-capping. After the 1/sqrt(head_dim) scale
283		// and BEFORE the causal/window mask and the online-softmax max/exp, each
284		// raw score s is replaced by attn_softcap * tanh(s / attn_softcap).
285		// attn_softcap == 0 (default) disables it — bit-identical to before.
286		// window > 0: sliding-window causal masking. Query at absolute position p
287		// (p = valid_len - L_q + q) attends key j only when j <= p AND
288		// j > p - window, i.e. keys in [max(0, p-window+1), p]. window <= 0
289		// (default) is unbounded causal — bit-identical to before. Composes
290		// multiplicatively with the causal mask.
291		void flash_attention_decode(const Tensor& Q,
292		const Tensor& K_cache, const Tensor& V_cache,
293		int valid_len, int num_q_heads, int num_kv_heads,
294		Tensor& O,
295		float attn_softcap = 0.0f, int window = 0);
296
297
298		// Single-token decode attention over a FIXED-CAPACITY masked KV cache — the
299		// CUDA-graph-capturable twin of flash_attention_decode. The kernel always
300		// sees the full (L_max, ·) cache buffers, so its launch shape never changes
301		// as generation advances; validity lives in `d_mask`, a device-resident
302		// length-L_max FP32 key mask (1 valid / 0 invalid) the caller updates
303		// between graph replays. Masked keys are dropped before the dot product and
304		// softmax (their weights underflow to exact zeros), so with
305		// mask = [1]valid_len + [0](L_max-valid_len) the result is bit-identical
306		// to flash_attention_decode(Q, K, V, valid_len, ...) at L_q == 1. Fully
307		// masked key tiles are skipped without touching K/V.
308		// Q: (1, num_q_heads*head_dim) — single query row only.
309		// K_cache, V_cache: (L_max, num_kv_heads*head_dim); same GQA rules as
310		// flash_attention_decode.
311		// d_mask: device pointer, length L_max, FP32; must not be null and must
312		// mark at least one key valid.
313		// O: (1, num_q_heads*head_dim), resized as needed.
314		// Gemma-2 extensions match flash_attention_decode (see above):
315		// attn_softcap > 0: tanh logit soft-capping applied to each raw score before
316		// the mask and softmax. attn_softcap == 0 (default) disables it.
317		// window > 0: sliding-window masking. The query's absolute position p is the
318		// highest valid key index (the last 1 in d_mask); a key j is kept only
319		// when it is valid AND j > p - window, i.e. j in [max(0,p-window+1), p].
320		// window <= 0 (default) keeps the full valid set — bit-identical to
321		// before. Combines multiplicatively with the validity mask.
322		void flash_attention_decode_masked(const Tensor& Q,
323		const Tensor& K_cache,
324		const Tensor& V_cache,
325		const float* d_mask,
326		int num_q_heads, int num_kv_heads,
327		Tensor& O,
328		float attn_softcap = 0.0f, int window = 0);
329
330
331		// Back-compat overload: num_kv_heads defaults to num_heads (plain MHA).
332	1	inline void flash_attention_decode(const Tensor& Q,
333		const Tensor& K_cache, const Tensor& V_cache,
334		int valid_len, int num_heads, Tensor& O,
335		float attn_softcap = 0.0f, int window = 0) {
336	2	flash_attention_decode(Q, K_cache, V_cache, valid_len,
337	1	num_heads, /num_kv_heads=/num_heads, O,
338	1	attn_softcap, window);
339	1	}
340
341
342		// ─── W8A16 flash-attention variants ────────────────────────────────────────
343		//
344		// Same composition as flash_attention_project_kv /
345		// flash_attention_q_with_kv_cached_forward / flash_attention_qkvo_forward, but
346		// every projection takes an INT8 weight + per-output-row FP32 scale instead of
347		// an FP16 weight. The attention core stays FP16 (activations are never
348		// quantised). Each W*_int8 is (D, in_dim) with a matching scale (D,1); biases
349		// stay FP16 (D,1), optional. Masks, causal flag, and num_heads match the FP16
350		// versions.
351		void flash_attention_project_kv_int8w_fp16(const Tensor& ctx,
352		const Tensor& Wk_int8,
353		const Tensor& sk,
354		const Tensor* bk,
355		const Tensor& Wv_int8,
356		const Tensor& sv,
357		const Tensor* bv,
358		Tensor& K_out,
359		Tensor& V_out);
360
361
362		void flash_attention_q_with_kv_cached_int8w_fp16(const Tensor& X,
363		const Tensor& K,
364		const Tensor& V,
365		const Tensor& Wq_int8,
366		const Tensor& sq,
367		const Tensor* bq,
368		const Tensor& Wo_int8,
369		const Tensor& so,
370		const Tensor* bo,
371		const float* d_mask,
372		int num_heads,
373		bool causal,
374		Tensor& O);
375
376
377		void flash_attention_qkvo_int8w_fp16(const Tensor& X,
378		const Tensor* Ctx,
379		const Tensor& Wq_int8, const Tensor& sq, const Tensor* bq,
380		const Tensor& Wk_int8, const Tensor& sk, const Tensor* bk,
381		const Tensor& Wv_int8, const Tensor& sv, const Tensor* bv,
382		const Tensor& Wo_int8, const Tensor& so, const Tensor* bo,
383		const float* d_mask,
384		int num_heads,
385		bool causal,
386		Tensor& O);
387
388		} // namespace brotensor
389