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

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

	Coverage	Exec / Excl / Total
Lines:	100.0%	9 / 0 / 9
Functions:	100.0%	2 / 0 / 2
Branches:	-%	0 / 0 / 0

    include/brotensor/ops/attention.h
    
        Line
        Branch
        Exec
        Source
      
        #pragma once
      
        // brotensor ops/attention.h — Attention (non-flash): single-head, MHA, self/cross, bias, decomposed rel-pos, masks.
      
        #include "../tensor.h"
      
        #include <cstdint>
      
        namespace brotensor {
      
        // Build a slot-validity mask on-device. For k in [0, K):
      
        //   mask[k] = (x[offset + k*stride] > 0.5f) ? 1.0f : 0.0f
      
        // mask resized to (K, 1).
      
        void build_slot_mask(const Tensor& x, int offset, int K, int stride,
      
                             Tensor& mask);
      
        // Single-head scaled dot-product self-attention. Square (D,D) projections,
      
        // no biases.
      
        //   X: (N,D).  Wq, Wk, Wv, Wo: each (D,D).
      
        //   d_mask: optional length-N device mask (1 valid / 0 invalid); may be null.
      
        //           Invalid keys are excluded from the softmax denominator; invalid
      
        //           query rows produce zero output.
      
        //   O: (N,D) output, resized if mis-shaped.
      
        //   Backward caches (out-params): Q, K, V each (N,D); Attn (N,N) post-softmax
      
        //   weights; Y_pre_Wo (N,D) = Attn @ V before the output projection.
      
        void attention_forward(const Tensor& X,
      
                               const Tensor& Wq, const Tensor& Wk,
      
                               const Tensor& Wv, const Tensor& Wo,
      
                               const float* d_mask,
      
                               Tensor& Q, Tensor& K, Tensor& V,
      
                               Tensor& Attn, Tensor& Y_pre_Wo,
      
                               Tensor& O);
      
        // Backward of attention_forward.
      
        //   dO: (N,D) upstream.  X, Q, K, V, Attn, Y_pre_Wo: forward caches.
      
        //   Wq, Wk, Wv, Wo: (D,D) forward weights.  d_mask: as forward (or null).
      
        //   dX: (N,D) overwritten.  dWq, dWk, dWv, dWo: (D,D) accumulated — caller zeros.
      
        void attention_backward(const Tensor& dO,
      
                                const Tensor& X,
      
                                const Tensor& Q, const Tensor& K,
      
                                const Tensor& V, const Tensor& Attn,
      
                                const Tensor& Y_pre_Wo,
      
                                const Tensor& Wq, const Tensor& Wk,
      
                                const Tensor& Wv, const Tensor& Wo,
      
                                const float* d_mask,
      
                                Tensor& dX,
      
                                Tensor& dWq, Tensor& dWk,
      
                                Tensor& dWv, Tensor& dWo);
      
        // ─── Multi-head self-attention ─────────────────────────────────────────────
      
        // Multi-head scaled dot-product self-attention. Square (D,D) projections,
      
        // split into num_heads heads of head_dim = D / num_heads; num_heads must
      
        // divide D.
      
        //   X: (K,D).  Wq, Wk, Wv, Wo: each (D,D).
      
        //   bq, bk, bv, bo: optional length-D bias vectors (any shape with D
      
        //                   elements, FP32). Added row-wise after the matching
      
        //                   projection: Q/K/V get bq/bk/bv post-projection, O gets
      
        //                   bo post-Wo. Any of the four may be null to skip that
      
        //                   bias term.
      
        //   d_mask: optional length-K device mask (1 valid / 0 invalid); may be null.
      
        //           Same semantics as single-head attention.
      
        //   O: (K,D) output, resized if mis-shaped.
      
        //   Backward caches (out-params, resized if mis-shaped): Qh, Kh, Vh
      
        //   (num_heads*K, head_dim) with head h in rows [h*K, (h+1)*K); Attnh
      
        //   (num_heads*K, K) per-head softmax weights; Yconcat (K,D) pre-Wo concat
      
        //   (does NOT include bo — bo is folded into O directly).
      
        void mha_forward(const Tensor& X,
      
                         const Tensor& Wq, const Tensor& Wk,
      
                         const Tensor& Wv, const Tensor& Wo,
      
                         const Tensor* bq, const Tensor* bk,
      
                         const Tensor* bv, const Tensor* bo,
      
                         const float* d_mask,
      
                         int num_heads,
      
                         Tensor& Qh, Tensor& Kh, Tensor& Vh,
      
                         Tensor& Attnh, Tensor& Yconcat,
      
                         Tensor& O);
      
        // Bias-less convenience overload — forwards to the bias-aware mha_forward
      
        // with bq/bk/bv/bo == nullptr. Preserves the original call shape so existing
      
        // callers don't change.
      
        16
        inline void mha_forward(const Tensor& X,
      
                                const Tensor& Wq, const Tensor& Wk,
      
                                const Tensor& Wv, const Tensor& Wo,
      
                                const float* d_mask,
      
                                int num_heads,
      
                                Tensor& Qh, Tensor& Kh, Tensor& Vh,
      
                                Tensor& Attnh, Tensor& Yconcat,
      
                                Tensor& O) {
      
        32
            mha_forward(X, Wq, Wk, Wv, Wo,
      
                        nullptr, nullptr, nullptr, nullptr,
      
        16
                        d_mask, num_heads,
      
        16
                        Qh, Kh, Vh, Attnh, Yconcat, O);
      
        16
        }
      
        // Backward of mha_forward.
      
        //   dO: (K,D) upstream.  X, Qh, Kh, Vh, Attnh, Yconcat: forward caches.
      
        //   Wq, Wk, Wv, Wo: (D,D) forward weights.  d_mask: as forward (or null).
      
        //   num_heads must match forward.
      
        //   dX: (K,D) overwritten.  dWq, dWk, dWv, dWo: (D,D) accumulated — caller zeros.
      
        //   dbq, dbk, dbv, dbo: optional length-D bias gradients, accumulated
      
        //                       (caller zeros). Pass null to skip — must match the
      
        //                       null/non-null pattern of the forward biases.
      
        void mha_backward(const Tensor& dO,
      
                          const Tensor& X,
      
                          const Tensor& Qh, const Tensor& Kh,
      
                          const Tensor& Vh, const Tensor& Attnh,
      
                          const Tensor& Yconcat,
      
                          const Tensor& Wq, const Tensor& Wk,
      
                          const Tensor& Wv, const Tensor& Wo,
      
                          const float* d_mask,
      
                          int num_heads,
      
                          Tensor& dX,
      
                          Tensor& dWq, Tensor& dWk,
      
                          Tensor& dWv, Tensor& dWo,
      
                          Tensor* dbq = nullptr, Tensor* dbk = nullptr,
      
                          Tensor* dbv = nullptr, Tensor* dbo = nullptr);
      
        // ─── Cross / self attention + masks ────────────────────────────────────────
      
        // Causal mask helper: fills the length-L FP32 buffer for query row q,
      
        //   mask[k] = (k <= q) ? 1.0f : 0.0f
      
        // resized to (L,1) if mis-shaped. The attention kernels consume a per-row
      
        // length-Lk mask; for fully causal self-attention launch attention per query.
      
        void build_causal_mask_row(int L, int q, Tensor& mask);
      
        // Cross-attention: like mha_forward but K and V are projected from a separate
      
        // context tensor. Dispatched on X.dtype:
      
        //   FP16 — flash-attention inference path; caches not exposed (use
      
        //          cross_attention_forward_train if you need them).
      
        //   FP32 — training-aware path (allocates scratch caches internally).
      
        //   X:   (Lq,D)      query input.
      
        //   Ctx: (Lk,D_ctx)  key/value input; Lk, D_ctx may differ from Lq, D.
      
        //                    Ctx.dtype must match X.dtype.
      
        //   Wq, Wo: (D,D).  Wk, Wv: (D,D_ctx) (rectangular for cross-attention).
      
        //   d_mask: optional length-Lk FP32 mask (1 valid / 0 invalid); may be null.
      
        //   num_heads divides D.
      
        //   O: (Lq,D) output, same dtype as X, resized if mis-shaped.
      
        void cross_attention_forward(const Tensor& X,
      
                                     const Tensor& Ctx,
      
                                     const Tensor& Wq, const Tensor& Wk,
      
                                     const Tensor& Wv, const Tensor& Wo,
      
                                     const float* d_mask,
      
                                     int num_heads,
      
                                     Tensor& O);
      
        // Cross-attention with a head-averaged attention map and an optional
      
        // pre-softmax logit bias. FP16 only, FP32 accumulation, no backward. Same math
      
        // as cross_attention_forward, plus:
      
        //   * if attn_logit_bias is non-null it is added to the scaled QK^T scores
      
        //     before softmax, broadcast across heads;
      
        //   * AttnAvg receives the across-head average of the softmax weights.
      
        //   X: (Lq,D); Ctx: (Lk,D_ctx); Wq, Wo: (D,D); Wk, Wv: (D,D_ctx) — all FP16.
      
        //   d_mask: optional length-Lk FP32 mask (1 valid / 0 invalid); may be null.
      
        //   attn_logit_bias: optional (Lq,Lk) FP32 pre-softmax bias; may be null.
      
        //   num_heads divides D.
      
        //   O: (Lq,D) FP16; AttnAvg: (Lq,Lk) FP16 — both resized + dtype-set if needed.
      
        void cross_attention_forward_with_attn(const Tensor& X,
      
                                               const Tensor& Ctx,
      
                                               const Tensor& Wq, const Tensor& Wk,
      
                                               const Tensor& Wv, const Tensor& Wo,
      
                                               const float* d_mask,
      
                                               const Tensor* attn_logit_bias,
      
                                               int num_heads,
      
                                               Tensor& O,
      
                                               Tensor& AttnAvg);
      
        // FP32 training-side self-attention forward — thin wrapper over mha_forward
      
        // (the mha case with Ctx == X). All tensors FP32.
      
        //   X, O: (L,D).  Wq, Wk, Wv, Wo: (D,D).
      
        //   d_mask: optional length-L FP32 mask; may be null.  num_heads divides D.
      
        //   Caches (resized if mis-shaped): Qh, Kh, Vh (num_heads*L, D/num_heads);
      
        //   Attnh (num_heads*L, L); Yconcat (L,D).
      
        void self_attention_forward_train(const Tensor& X,
      
                                          const Tensor& Wq, const Tensor& Wk,
      
                                          const Tensor& Wv, const Tensor& Wo,
      
                                          const float* d_mask,
      
                                          int num_heads,
      
                                          Tensor& Qh, Tensor& Kh, Tensor& Vh,
      
                                          Tensor& Attnh, Tensor& Yconcat,
      
                                          Tensor& O);
      
        // FP32 training-side self-attention backward — thin wrapper over mha_backward.
      
        // All tensors FP32.
      
        //   dO: (L,D) upstream.  X, Qh, Kh, Vh, Attnh, Yconcat: forward caches.
      
        //   Wq, Wk, Wv, Wo: (D,D) forward weights.  d_mask: as forward (or null).
      
        //   num_heads must match forward.
      
        //   dX: (L,D) overwritten.  dWq, dWk, dWv, dWo: (D,D) accumulated — caller zeros.
      
        void self_attention_backward(const Tensor& dO,
      
                                     const Tensor& X,
      
                                     const Tensor& Qh, const Tensor& Kh,
      
                                     const Tensor& Vh, const Tensor& Attnh,
      
                                     const Tensor& Yconcat,
      
                                     const Tensor& Wq, const Tensor& Wk,
      
                                     const Tensor& Wv, const Tensor& Wo,
      
                                     const float* d_mask,
      
                                     int num_heads,
      
                                     Tensor& dX,
      
                                     Tensor& dWq, Tensor& dWk,
      
                                     Tensor& dWv, Tensor& dWo);
      
        // Per-text-token spatial moments of a cross-attention map. Given Attn (Lq,Lk)
      
        // with Lq = h_lat*w_lat a flattened row-major image-token grid (q = y*w_lat+x),
      
        // for each text token k:
      
        //   mass[k]       = sum_q Attn[q,k]
      
        //   centroid[k,0] = sum_q y(q)*Attn[q,k] / max(mass[k], 1e-8)   (y)
      
        //   centroid[k,1] = sum_q x(q)*Attn[q,k] / max(mass[k], 1e-8)   (x)
      
        //   (centroid set to (0,0) when mass[k] is ~ 0).
      
        //   Attn: (Lq,Lk) FP16.  mass: (Lk,1) FP32.  centroid: (Lk,2) FP32 [y,x].
      
        //   mass and centroid resized if mis-shaped. FP32 reductions over FP16 input.
      
        void attention_token_moments(const Tensor& Attn,
      
                                     int h_lat, int w_lat,
      
                                     Tensor& mass,
      
                                     Tensor& centroid);
      
        // FP32 training-side cross-attention forward. mha_forward math with a separate
      
        // Ctx for K/V and rectangular Wk/Wv. All tensors FP32.
      
        //   X: (Lq,D); Ctx: (Lk,D_ctx); Wq, Wo: (D,D); Wk, Wv: (D,D_ctx).
      
        //   d_mask: optional length-Lk FP32 mask; may be null.  num_heads divides D.
      
        //   Caches (resized if mis-shaped): Qh (num_heads*Lq, D/num_heads);
      
        //   Kh, Vh (num_heads*Lk, D/num_heads); Attnh (num_heads*Lq, Lk);
      
        //   Yconcat (Lq,D).  O: (Lq,D), resized if mis-shaped.
      
        void cross_attention_forward_train(const Tensor& X,
      
                                           const Tensor& Ctx,
      
                                           const Tensor& Wq, const Tensor& Wk,
      
                                           const Tensor& Wv, const Tensor& Wo,
      
                                           const float* d_mask,
      
                                           int num_heads,
      
                                           Tensor& Qh, Tensor& Kh, Tensor& Vh,
      
                                           Tensor& Attnh, Tensor& Yconcat,
      
                                           Tensor& O);
      
        // FP32 training-side cross-attention backward.
      
        //   dO: (Lq,D) upstream.  X, Ctx, Qh, Kh, Vh, Attnh, Yconcat: forward caches.
      
        //   Wq, Wo: (D,D); Wk, Wv: (D,D_ctx).  d_mask: as forward (or null).
      
        //   num_heads must match forward.
      
        //   dX: (Lq,D), dCtx: (Lk,D_ctx) — overwritten.
      
        //   dWq, dWo: (D,D); dWk, dWv: (D,D_ctx) — accumulated, caller zeros.
      
        void cross_attention_backward(const Tensor& dO,
      
                                      const Tensor& X,
      
                                      const Tensor& Ctx,
      
                                      const Tensor& Qh, const Tensor& Kh,
      
                                      const Tensor& Vh, const Tensor& Attnh,
      
                                      const Tensor& Yconcat,
      
                                      const Tensor& Wq, const Tensor& Wk,
      
                                      const Tensor& Wv, const Tensor& Wo,
      
                                      const float* d_mask,
      
                                      int num_heads,
      
                                      Tensor& dX,
      
                                      Tensor& dCtx,
      
                                      Tensor& dWq, Tensor& dWk,
      
                                      Tensor& dWv, Tensor& dWo);
      
        // FP16 self-attention — thin wrapper over the cross-attention kernel with
      
        // Ctx = X. Conventions as cross_attention_forward (FP16 path).
      
        //   X, O: (L,D) FP16.  Wq, Wk, Wv, Wo: (D,D) FP16.
      
        //   d_mask: optional length-L FP32 mask; may be null.  num_heads divides D.
      
        //   O resized if mis-shaped.
      
        void self_attention_forward(const Tensor& X,
      
                                    const Tensor& Wq, const Tensor& Wk,
      
                                    const Tensor& Wv, const Tensor& Wo,
      
                                    const float* d_mask,
      
                                    int num_heads,
      
                                    Tensor& O);
      
        // Multi-head self-attention with an optional additive pre-softmax bias. Per
      
        // head h:
      
        //   S[q,k] = scale*(Q_h[q].K_h[k]) + attn_bias[h*L+q, k]
      
        //   O = (softmax_k S) @ V_h, concatenated over heads, projected by Wo.
      
        // The additive bias is the primitive behind T5 relative-position bias and
      
        // ALiBi. `scale` multiplies the raw dot product BEFORE the bias: pass
      
        // 1/sqrt(head_dim) for standard attention, or 1.0 for T5.
      
        //   X, O: (L,D).  Wq, Wk, Wv, Wo: (D,D), same dtype as X.
      
        //   bq, bk, bv, bo: optional length-D projection biases (X.dtype), added
      
        //                   row-wise after the matching projection — bq/bk/bv to Q/K/V,
      
        //                   bo after Wo. Any may be null. This makes the op a full
      
        //                   biased MHA + additive-bias attention (the Swin window-attn
      
        //                   block: qkv/proj bias + relative-position bias as attn_bias).
      
        //   d_mask: optional length-L FP32 key mask (also gates padded query rows);
      
        //           may be null.
      
        //   attn_bias: optional (num_heads*L, L) FP32 — row h*L+q is head h query q's
      
        //              length-L bias. Null => plain scaled self-attention.
      
        //   num_heads divides D.  O resized + dtype-set to match X.
      
        // Dispatched on X.dtype (FP32/FP16/BF16); FP32 math; attn_bias is FP32 on every
      
        // backend. Scores are materialised (L,L) per head — for encoder-length seqs.
      
        void self_attention_bias_forward(const Tensor& X,
      
                                         const Tensor& Wq, const Tensor& Wk,
      
                                         const Tensor& Wv, const Tensor& Wo,
      
                                         const Tensor* bq, const Tensor* bk,
      
                                         const Tensor* bv, const Tensor* bo,
      
                                         const float* d_mask,
      
                                         const Tensor* attn_bias,
      
                                         int num_heads, float scale,
      
                                         Tensor& O);
      
        // Bias-less-projection convenience overload — preserves the original call shape
      
        // (no qkv/proj biases) by forwarding bq/bk/bv/bo == nullptr.
      
        33
        inline void self_attention_bias_forward(const Tensor& X,
      
                                                const Tensor& Wq, const Tensor& Wk,
      
                                                const Tensor& Wv, const Tensor& Wo,
      
                                                const float* d_mask,
      
                                                const Tensor* attn_bias,
      
                                                int num_heads, float scale,
      
                                                Tensor& O) {
      
        66
            self_attention_bias_forward(X, Wq, Wk, Wv, Wo,
      
                                        nullptr, nullptr, nullptr, nullptr,
      
        33
                                        d_mask, attn_bias, num_heads, scale, O);
      
        33
        }
      
        // Multi-head self-attention with a DECOMPOSED 2D relative-position bias — the
      
        // SAM / ViTDet image-encoder attention (segment_anything add_decomposed_rel_pos).
      
        // A token index t maps to grid coords (t/grid_w, t%grid_w) over a grid_h*grid_w
      
        // patch grid (so X.rows == grid_h*grid_w). Per head, with r_q the projected,
      
        // UNSCALED query:
      
        //   bias[q,k] = r_q . rel_pos_h[(qh-kh)+grid_h-1] + r_q . rel_pos_w[(qw-kw)+grid_w-1]
      
        //   S[q,k]    = scale*(Q[q].K[k]) + bias[q,k]   (scale multiplies the dot only)
      
        //   O = concat_h( softmax_k S @ V_h ) @ Wo
      
        // Unlike self_attention_bias_forward the bias is data-dependent (reads Q) and is
      
        // never materialised as (num_heads*L, L) — it's factored into length-grid_h and
      
        // length-grid_w terms. Windowed blocks call this per window (grid == window);
      
        // global blocks call it once over the full grid.
      
        //   X, O: (L, D).  Wq/Wk/Wv/Wo: (D, D).  bq/bk/bv/bo: optional (D,1), null to skip.
      
        //   rel_pos_h: (2*grid_h-1, head_dim).  rel_pos_w: (2*grid_w-1, head_dim).
      
        //   num_heads divides D.  scale: typically 1/sqrt(head_dim).
      
        // Dispatched on X.dtype; O resized + dtype-set to match X.
      
        void self_attention_decomposed_rel_pos_forward(
      
                const Tensor& X,
      
                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 Tensor& rel_pos_h, const Tensor& rel_pos_w,
      
                int num_heads, int grid_h, int grid_w, float scale,
      
                Tensor& O);
      
        // Windowed multi-head self-attention with a decomposed 2D relative-position
      
        // bias — the SAM/ViTDet *windowed* encoder block (segment_anything runs most
      
        // blocks over non-overlapping local windows, a few globally). Splits the
      
        // (grid_h, grid_w) token grid into window x window tiles and runs the
      
        // decomposed-rel-pos attention above INDEPENDENTLY within each tile, sharing
      
        // one set of weights and rel-pos tables. The bottom/right of the grid is
      
        // zero-padded up to a multiple of `window` (SAM's window_partition pad) and the
      
        // padding is cropped back off the output, so grid_h/grid_w need not be
      
        // multiples of `window`. For a grid that is exactly one window this is the
      
        // plain decomposed-rel-pos op.
      
        //   X, O: (grid_h*grid_w, D), token-major (row = h*grid_w + w).
      
        //   Wq/Wk/Wv/Wo: (D,D).  bq/bk/bv/bo: optional (D,1), null to skip.
      
        //   rel_pos_h, rel_pos_w: (2*window-1, head_dim) — sized for the window, not
      
        //                         the full grid.
      
        //   num_heads divides D.  scale: typically 1/sqrt(head_dim).
      
        // Dispatched on X.dtype; O resized + dtype-set to match X.
      
        void self_attention_decomposed_rel_pos_windowed_forward(
      
                const Tensor& X,
      
                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 Tensor& rel_pos_h, const Tensor& rel_pos_w,
      
                int num_heads, int grid_h, int grid_w, int window, float scale,
      
                Tensor& O);
      
        // W8A16 variant of self_attention_bias_forward — quantised T5-bias attention.
      
        // Identical math and semantics, but each projection weight is an INT8 (D,D)
      
        // matrix paired with an FP32 (D,1) per-output-row dequant scale (the
      
        // quantize_int8_per_row_host convention). Activations stay FP16; the attention
      
        // core is FP32 internally. GPU-only.
      
        //   X, O: (L,D) FP16.  Wq/Wk/Wv/Wo_int8: (D,D) INT8.  sq/sk/sv/so: (D,1) FP32.
      
        //   d_mask: optional length-L FP32 key mask; may be null.
      
        //   attn_bias: optional (num_heads*L, L) FP32 bias; may be null.
      
        //   num_heads divides D.  scale: as in self_attention_bias_forward.
      
        //   O resized as needed.
      
        void self_attention_bias_int8w_fp16(const Tensor& X,
      
                                            const Tensor& Wq_int8, const Tensor& sq,
      
                                            const Tensor& Wk_int8, const Tensor& sk,
      
                                            const Tensor& Wv_int8, const Tensor& sv,
      
                                            const Tensor& Wo_int8, const Tensor& so,
      
                                            const float* d_mask,
      
                                            const Tensor* attn_bias,
      
                                            int num_heads, float scale,
      
                                            Tensor& O);
      
        }  // namespace brotensor

Line	Exec	Source
1		#pragma once
2
3		// brotensor ops/attention.h — Attention (non-flash): single-head, MHA, self/cross, bias, decomposed rel-pos, masks.
4
5		#include "../tensor.h"
6		#include <cstdint>
7
8		namespace brotensor {
9
10
11		// Build a slot-validity mask on-device. For k in [0, K):
12		// mask[k] = (x[offset + k*stride] > 0.5f) ? 1.0f : 0.0f
13		// mask resized to (K, 1).
14		void build_slot_mask(const Tensor& x, int offset, int K, int stride,
15		Tensor& mask);
16
17
18		// Single-head scaled dot-product self-attention. Square (D,D) projections,
19		// no biases.
20		// X: (N,D). Wq, Wk, Wv, Wo: each (D,D).
21		// d_mask: optional length-N device mask (1 valid / 0 invalid); may be null.
22		// Invalid keys are excluded from the softmax denominator; invalid
23		// query rows produce zero output.
24		// O: (N,D) output, resized if mis-shaped.
25		// Backward caches (out-params): Q, K, V each (N,D); Attn (N,N) post-softmax
26		// weights; Y_pre_Wo (N,D) = Attn @ V before the output projection.
27		void attention_forward(const Tensor& X,
28		const Tensor& Wq, const Tensor& Wk,
29		const Tensor& Wv, const Tensor& Wo,
30		const float* d_mask,
31		Tensor& Q, Tensor& K, Tensor& V,
32		Tensor& Attn, Tensor& Y_pre_Wo,
33		Tensor& O);
34
35
36		// Backward of attention_forward.
37		// dO: (N,D) upstream. X, Q, K, V, Attn, Y_pre_Wo: forward caches.
38		// Wq, Wk, Wv, Wo: (D,D) forward weights. d_mask: as forward (or null).
39		// dX: (N,D) overwritten. dWq, dWk, dWv, dWo: (D,D) accumulated — caller zeros.
40		void attention_backward(const Tensor& dO,
41		const Tensor& X,
42		const Tensor& Q, const Tensor& K,
43		const Tensor& V, const Tensor& Attn,
44		const Tensor& Y_pre_Wo,
45		const Tensor& Wq, const Tensor& Wk,
46		const Tensor& Wv, const Tensor& Wo,
47		const float* d_mask,
48		Tensor& dX,
49		Tensor& dWq, Tensor& dWk,
50		Tensor& dWv, Tensor& dWo);
51
52
53		// ─── Multi-head self-attention ─────────────────────────────────────────────
54
55		// Multi-head scaled dot-product self-attention. Square (D,D) projections,
56		// split into num_heads heads of head_dim = D / num_heads; num_heads must
57		// divide D.
58		// X: (K,D). Wq, Wk, Wv, Wo: each (D,D).
59		// bq, bk, bv, bo: optional length-D bias vectors (any shape with D
60		// elements, FP32). Added row-wise after the matching
61		// projection: Q/K/V get bq/bk/bv post-projection, O gets
62		// bo post-Wo. Any of the four may be null to skip that
63		// bias term.
64		// d_mask: optional length-K device mask (1 valid / 0 invalid); may be null.
65		// Same semantics as single-head attention.
66		// O: (K,D) output, resized if mis-shaped.
67		// Backward caches (out-params, resized if mis-shaped): Qh, Kh, Vh
68		// (num_headsK, head_dim) with head h in rows [hK, (h+1)*K); Attnh
69		// (num_heads*K, K) per-head softmax weights; Yconcat (K,D) pre-Wo concat
70		// (does NOT include bo — bo is folded into O directly).
71		void mha_forward(const Tensor& X,
72		const Tensor& Wq, const Tensor& Wk,
73		const Tensor& Wv, const Tensor& Wo,
74		const Tensor* bq, const Tensor* bk,
75		const Tensor* bv, const Tensor* bo,
76		const float* d_mask,
77		int num_heads,
78		Tensor& Qh, Tensor& Kh, Tensor& Vh,
79		Tensor& Attnh, Tensor& Yconcat,
80		Tensor& O);
81
82
83		// Bias-less convenience overload — forwards to the bias-aware mha_forward
84		// with bq/bk/bv/bo == nullptr. Preserves the original call shape so existing
85		// callers don't change.
86	16	inline void mha_forward(const Tensor& X,
87		const Tensor& Wq, const Tensor& Wk,
88		const Tensor& Wv, const Tensor& Wo,
89		const float* d_mask,
90		int num_heads,
91		Tensor& Qh, Tensor& Kh, Tensor& Vh,
92		Tensor& Attnh, Tensor& Yconcat,
93		Tensor& O) {
94	32	mha_forward(X, Wq, Wk, Wv, Wo,
95		nullptr, nullptr, nullptr, nullptr,
96	16	d_mask, num_heads,
97	16	Qh, Kh, Vh, Attnh, Yconcat, O);
98	16	}
99
100
101		// Backward of mha_forward.
102		// dO: (K,D) upstream. X, Qh, Kh, Vh, Attnh, Yconcat: forward caches.
103		// Wq, Wk, Wv, Wo: (D,D) forward weights. d_mask: as forward (or null).
104		// num_heads must match forward.
105		// dX: (K,D) overwritten. dWq, dWk, dWv, dWo: (D,D) accumulated — caller zeros.
106		// dbq, dbk, dbv, dbo: optional length-D bias gradients, accumulated
107		// (caller zeros). Pass null to skip — must match the
108		// null/non-null pattern of the forward biases.
109		void mha_backward(const Tensor& dO,
110		const Tensor& X,
111		const Tensor& Qh, const Tensor& Kh,
112		const Tensor& Vh, const Tensor& Attnh,
113		const Tensor& Yconcat,
114		const Tensor& Wq, const Tensor& Wk,
115		const Tensor& Wv, const Tensor& Wo,
116		const float* d_mask,
117		int num_heads,
118		Tensor& dX,
119		Tensor& dWq, Tensor& dWk,
120		Tensor& dWv, Tensor& dWo,
121		Tensor* dbq = nullptr, Tensor* dbk = nullptr,
122		Tensor* dbv = nullptr, Tensor* dbo = nullptr);
123
124
125		// ─── Cross / self attention + masks ────────────────────────────────────────
126
127		// Causal mask helper: fills the length-L FP32 buffer for query row q,
128		// mask[k] = (k <= q) ? 1.0f : 0.0f
129		// resized to (L,1) if mis-shaped. The attention kernels consume a per-row
130		// length-Lk mask; for fully causal self-attention launch attention per query.
131		void build_causal_mask_row(int L, int q, Tensor& mask);
132
133
134		// Cross-attention: like mha_forward but K and V are projected from a separate
135		// context tensor. Dispatched on X.dtype:
136		// FP16 — flash-attention inference path; caches not exposed (use
137		// cross_attention_forward_train if you need them).
138		// FP32 — training-aware path (allocates scratch caches internally).
139		// X: (Lq,D) query input.
140		// Ctx: (Lk,D_ctx) key/value input; Lk, D_ctx may differ from Lq, D.
141		// Ctx.dtype must match X.dtype.
142		// Wq, Wo: (D,D). Wk, Wv: (D,D_ctx) (rectangular for cross-attention).
143		// d_mask: optional length-Lk FP32 mask (1 valid / 0 invalid); may be null.
144		// num_heads divides D.
145		// O: (Lq,D) output, same dtype as X, resized if mis-shaped.
146		void cross_attention_forward(const Tensor& X,
147		const Tensor& Ctx,
148		const Tensor& Wq, const Tensor& Wk,
149		const Tensor& Wv, const Tensor& Wo,
150		const float* d_mask,
151		int num_heads,
152		Tensor& O);
153
154
155		// Cross-attention with a head-averaged attention map and an optional
156		// pre-softmax logit bias. FP16 only, FP32 accumulation, no backward. Same math
157		// as cross_attention_forward, plus:
158		// * if attn_logit_bias is non-null it is added to the scaled QK^T scores
159		// before softmax, broadcast across heads;
160		// * AttnAvg receives the across-head average of the softmax weights.
161		// X: (Lq,D); Ctx: (Lk,D_ctx); Wq, Wo: (D,D); Wk, Wv: (D,D_ctx) — all FP16.
162		// d_mask: optional length-Lk FP32 mask (1 valid / 0 invalid); may be null.
163		// attn_logit_bias: optional (Lq,Lk) FP32 pre-softmax bias; may be null.
164		// num_heads divides D.
165		// O: (Lq,D) FP16; AttnAvg: (Lq,Lk) FP16 — both resized + dtype-set if needed.
166		void cross_attention_forward_with_attn(const Tensor& X,
167		const Tensor& Ctx,
168		const Tensor& Wq, const Tensor& Wk,
169		const Tensor& Wv, const Tensor& Wo,
170		const float* d_mask,
171		const Tensor* attn_logit_bias,
172		int num_heads,
173		Tensor& O,
174		Tensor& AttnAvg);
175
176
177		// FP32 training-side self-attention forward — thin wrapper over mha_forward
178		// (the mha case with Ctx == X). All tensors FP32.
179		// X, O: (L,D). Wq, Wk, Wv, Wo: (D,D).
180		// d_mask: optional length-L FP32 mask; may be null. num_heads divides D.
181		// Caches (resized if mis-shaped): Qh, Kh, Vh (num_heads*L, D/num_heads);
182		// Attnh (num_heads*L, L); Yconcat (L,D).
183		void self_attention_forward_train(const Tensor& X,
184		const Tensor& Wq, const Tensor& Wk,
185		const Tensor& Wv, const Tensor& Wo,
186		const float* d_mask,
187		int num_heads,
188		Tensor& Qh, Tensor& Kh, Tensor& Vh,
189		Tensor& Attnh, Tensor& Yconcat,
190		Tensor& O);
191
192
193		// FP32 training-side self-attention backward — thin wrapper over mha_backward.
194		// All tensors FP32.
195		// dO: (L,D) upstream. X, Qh, Kh, Vh, Attnh, Yconcat: forward caches.
196		// Wq, Wk, Wv, Wo: (D,D) forward weights. d_mask: as forward (or null).
197		// num_heads must match forward.
198		// dX: (L,D) overwritten. dWq, dWk, dWv, dWo: (D,D) accumulated — caller zeros.
199		void self_attention_backward(const Tensor& dO,
200		const Tensor& X,
201		const Tensor& Qh, const Tensor& Kh,
202		const Tensor& Vh, const Tensor& Attnh,
203		const Tensor& Yconcat,
204		const Tensor& Wq, const Tensor& Wk,
205		const Tensor& Wv, const Tensor& Wo,
206		const float* d_mask,
207		int num_heads,
208		Tensor& dX,
209		Tensor& dWq, Tensor& dWk,
210		Tensor& dWv, Tensor& dWo);
211
212
213		// Per-text-token spatial moments of a cross-attention map. Given Attn (Lq,Lk)
214		// with Lq = h_latw_lat a flattened row-major image-token grid (q = yw_lat+x),
215		// for each text token k:
216		// mass[k] = sum_q Attn[q,k]
217		// centroid[k,0] = sum_q y(q)*Attn[q,k] / max(mass[k], 1e-8) (y)
218		// centroid[k,1] = sum_q x(q)*Attn[q,k] / max(mass[k], 1e-8) (x)
219		// (centroid set to (0,0) when mass[k] is ~ 0).
220		// Attn: (Lq,Lk) FP16. mass: (Lk,1) FP32. centroid: (Lk,2) FP32 [y,x].
221		// mass and centroid resized if mis-shaped. FP32 reductions over FP16 input.
222		void attention_token_moments(const Tensor& Attn,
223		int h_lat, int w_lat,
224		Tensor& mass,
225		Tensor& centroid);
226
227
228		// FP32 training-side cross-attention forward. mha_forward math with a separate
229		// Ctx for K/V and rectangular Wk/Wv. All tensors FP32.
230		// X: (Lq,D); Ctx: (Lk,D_ctx); Wq, Wo: (D,D); Wk, Wv: (D,D_ctx).
231		// d_mask: optional length-Lk FP32 mask; may be null. num_heads divides D.
232		// Caches (resized if mis-shaped): Qh (num_heads*Lq, D/num_heads);
233		// Kh, Vh (num_headsLk, D/num_heads); Attnh (num_headsLq, Lk);
234		// Yconcat (Lq,D). O: (Lq,D), resized if mis-shaped.
235		void cross_attention_forward_train(const Tensor& X,
236		const Tensor& Ctx,
237		const Tensor& Wq, const Tensor& Wk,
238		const Tensor& Wv, const Tensor& Wo,
239		const float* d_mask,
240		int num_heads,
241		Tensor& Qh, Tensor& Kh, Tensor& Vh,
242		Tensor& Attnh, Tensor& Yconcat,
243		Tensor& O);
244
245
246		// FP32 training-side cross-attention backward.
247		// dO: (Lq,D) upstream. X, Ctx, Qh, Kh, Vh, Attnh, Yconcat: forward caches.
248		// Wq, Wo: (D,D); Wk, Wv: (D,D_ctx). d_mask: as forward (or null).
249		// num_heads must match forward.
250		// dX: (Lq,D), dCtx: (Lk,D_ctx) — overwritten.
251		// dWq, dWo: (D,D); dWk, dWv: (D,D_ctx) — accumulated, caller zeros.
252		void cross_attention_backward(const Tensor& dO,
253		const Tensor& X,
254		const Tensor& Ctx,
255		const Tensor& Qh, const Tensor& Kh,
256		const Tensor& Vh, const Tensor& Attnh,
257		const Tensor& Yconcat,
258		const Tensor& Wq, const Tensor& Wk,
259		const Tensor& Wv, const Tensor& Wo,
260		const float* d_mask,
261		int num_heads,
262		Tensor& dX,
263		Tensor& dCtx,
264		Tensor& dWq, Tensor& dWk,
265		Tensor& dWv, Tensor& dWo);
266
267
268		// FP16 self-attention — thin wrapper over the cross-attention kernel with
269		// Ctx = X. Conventions as cross_attention_forward (FP16 path).
270		// X, O: (L,D) FP16. Wq, Wk, Wv, Wo: (D,D) FP16.
271		// d_mask: optional length-L FP32 mask; may be null. num_heads divides D.
272		// O resized if mis-shaped.
273		void self_attention_forward(const Tensor& X,
274		const Tensor& Wq, const Tensor& Wk,
275		const Tensor& Wv, const Tensor& Wo,
276		const float* d_mask,
277		int num_heads,
278		Tensor& O);
279
280
281		// Multi-head self-attention with an optional additive pre-softmax bias. Per
282		// head h:
283		// S[q,k] = scale(Q_h[q].K_h[k]) + attn_bias[hL+q, k]
284		// O = (softmax_k S) @ V_h, concatenated over heads, projected by Wo.
285		// The additive bias is the primitive behind T5 relative-position bias and
286		// ALiBi. `scale` multiplies the raw dot product BEFORE the bias: pass
287		// 1/sqrt(head_dim) for standard attention, or 1.0 for T5.
288		// X, O: (L,D). Wq, Wk, Wv, Wo: (D,D), same dtype as X.
289		// bq, bk, bv, bo: optional length-D projection biases (X.dtype), added
290		// row-wise after the matching projection — bq/bk/bv to Q/K/V,
291		// bo after Wo. Any may be null. This makes the op a full
292		// biased MHA + additive-bias attention (the Swin window-attn
293		// block: qkv/proj bias + relative-position bias as attn_bias).
294		// d_mask: optional length-L FP32 key mask (also gates padded query rows);
295		// may be null.
296		// attn_bias: optional (num_headsL, L) FP32 — row hL+q is head h query q's
297		// length-L bias. Null => plain scaled self-attention.
298		// num_heads divides D. O resized + dtype-set to match X.
299		// Dispatched on X.dtype (FP32/FP16/BF16); FP32 math; attn_bias is FP32 on every
300		// backend. Scores are materialised (L,L) per head — for encoder-length seqs.
301		void self_attention_bias_forward(const Tensor& X,
302		const Tensor& Wq, const Tensor& Wk,
303		const Tensor& Wv, const Tensor& Wo,
304		const Tensor* bq, const Tensor* bk,
305		const Tensor* bv, const Tensor* bo,
306		const float* d_mask,
307		const Tensor* attn_bias,
308		int num_heads, float scale,
309		Tensor& O);
310
311		// Bias-less-projection convenience overload — preserves the original call shape
312		// (no qkv/proj biases) by forwarding bq/bk/bv/bo == nullptr.
313	33	inline void self_attention_bias_forward(const Tensor& X,
314		const Tensor& Wq, const Tensor& Wk,
315		const Tensor& Wv, const Tensor& Wo,
316		const float* d_mask,
317		const Tensor* attn_bias,
318		int num_heads, float scale,
319		Tensor& O) {
320	66	self_attention_bias_forward(X, Wq, Wk, Wv, Wo,
321		nullptr, nullptr, nullptr, nullptr,
322	33	d_mask, attn_bias, num_heads, scale, O);
323	33	}
324
325
326		// Multi-head self-attention with a DECOMPOSED 2D relative-position bias — the
327		// SAM / ViTDet image-encoder attention (segment_anything add_decomposed_rel_pos).
328		// A token index t maps to grid coords (t/grid_w, t%grid_w) over a grid_h*grid_w
329		// patch grid (so X.rows == grid_h*grid_w). Per head, with r_q the projected,
330		// UNSCALED query:
331		// bias[q,k] = r_q . rel_pos_h[(qh-kh)+grid_h-1] + r_q . rel_pos_w[(qw-kw)+grid_w-1]
332		// S[q,k] = scale*(Q[q].K[k]) + bias[q,k] (scale multiplies the dot only)
333		// O = concat_h( softmax_k S @ V_h ) @ Wo
334		// Unlike self_attention_bias_forward the bias is data-dependent (reads Q) and is
335		// never materialised as (num_heads*L, L) — it's factored into length-grid_h and
336		// length-grid_w terms. Windowed blocks call this per window (grid == window);
337		// global blocks call it once over the full grid.
338		// X, O: (L, D). Wq/Wk/Wv/Wo: (D, D). bq/bk/bv/bo: optional (D,1), null to skip.
339		// rel_pos_h: (2grid_h-1, head_dim). rel_pos_w: (2grid_w-1, head_dim).
340		// num_heads divides D. scale: typically 1/sqrt(head_dim).
341		// Dispatched on X.dtype; O resized + dtype-set to match X.
342		void self_attention_decomposed_rel_pos_forward(
343		const Tensor& X,
344		const Tensor& Wq, const Tensor* bq,
345		const Tensor& Wk, const Tensor* bk,
346		const Tensor& Wv, const Tensor* bv,
347		const Tensor& Wo, const Tensor* bo,
348		const Tensor& rel_pos_h, const Tensor& rel_pos_w,
349		int num_heads, int grid_h, int grid_w, float scale,
350		Tensor& O);
351
352
353		// Windowed multi-head self-attention with a decomposed 2D relative-position
354		// bias — the SAM/ViTDet windowed encoder block (segment_anything runs most
355		// blocks over non-overlapping local windows, a few globally). Splits the
356		// (grid_h, grid_w) token grid into window x window tiles and runs the
357		// decomposed-rel-pos attention above INDEPENDENTLY within each tile, sharing
358		// one set of weights and rel-pos tables. The bottom/right of the grid is
359		// zero-padded up to a multiple of `window` (SAM's window_partition pad) and the
360		// padding is cropped back off the output, so grid_h/grid_w need not be
361		// multiples of `window`. For a grid that is exactly one window this is the
362		// plain decomposed-rel-pos op.
363		// X, O: (grid_hgrid_w, D), token-major (row = hgrid_w + w).
364		// Wq/Wk/Wv/Wo: (D,D). bq/bk/bv/bo: optional (D,1), null to skip.
365		// rel_pos_h, rel_pos_w: (2*window-1, head_dim) — sized for the window, not
366		// the full grid.
367		// num_heads divides D. scale: typically 1/sqrt(head_dim).
368		// Dispatched on X.dtype; O resized + dtype-set to match X.
369		void self_attention_decomposed_rel_pos_windowed_forward(
370		const Tensor& X,
371		const Tensor& Wq, const Tensor* bq,
372		const Tensor& Wk, const Tensor* bk,
373		const Tensor& Wv, const Tensor* bv,
374		const Tensor& Wo, const Tensor* bo,
375		const Tensor& rel_pos_h, const Tensor& rel_pos_w,
376		int num_heads, int grid_h, int grid_w, int window, float scale,
377		Tensor& O);
378
379
380		// W8A16 variant of self_attention_bias_forward — quantised T5-bias attention.
381		// Identical math and semantics, but each projection weight is an INT8 (D,D)
382		// matrix paired with an FP32 (D,1) per-output-row dequant scale (the
383		// quantize_int8_per_row_host convention). Activations stay FP16; the attention
384		// core is FP32 internally. GPU-only.
385		// X, O: (L,D) FP16. Wq/Wk/Wv/Wo_int8: (D,D) INT8. sq/sk/sv/so: (D,1) FP32.
386		// d_mask: optional length-L FP32 key mask; may be null.
387		// attn_bias: optional (num_heads*L, L) FP32 bias; may be null.
388		// num_heads divides D. scale: as in self_attention_bias_forward.
389		// O resized as needed.
390		void self_attention_bias_int8w_fp16(const Tensor& X,
391		const Tensor& Wq_int8, const Tensor& sq,
392		const Tensor& Wk_int8, const Tensor& sk,
393		const Tensor& Wv_int8, const Tensor& sv,
394		const Tensor& Wo_int8, const Tensor& so,
395		const float* d_mask,
396		const Tensor* attn_bias,
397		int num_heads, float scale,
398		Tensor& O);
399
400		} // namespace brotensor
401