spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#pragma once34#include "nn/config.h"5#include "nn/linear.h"6#include "ops/ops.h"78#include <cmath>9#include <memory>1011namespace forge::nn {1213// Interface so linear-attention / sliding-window variants can slot in14// without touching TransformerBlock.15class AttentionBase : public Module {16public:17    virtual ~AttentionBase() = default;18    // x: [B, T, C] → [B, T, C]19    virtual Var forward(const Var& x) const = 0;20};2122// Causal multi-head self-attention with GQA and (optional) RoPE.23// Per-layer flavour comes from `layer_idx` + config: NoPE layers skip RoPE24// (SmolLM3), global layers may use a different rope theta (Gemma 3), and25// head_dim may be decoupled from d_model/n_heads (Qwen3). Optional QKV bias26// (Qwen2.5, on q/k/v only — llama convention keeps wo bias-free).27class CausalSelfAttention : public AttentionBase {28public:29    CausalSelfAttention(const ModelConfig& cfg, float proj_std, std::mt19937_64& rng,30                        int64_t layer_idx = 0)31        : n_heads_(cfg.n_heads),32          n_kv_heads_(cfg.n_kv_heads),33          head_dim_(cfg.head_dim()),34          use_rope_(cfg.layer_uses_rope(layer_idx)),35          qk_norm_(cfg.qk_norm),36          norm_eps_(cfg.norm_eps),37          window_(cfg.layer_is_global(layer_idx) ? 0 : cfg.sliding_window),38          attn_softcap_(cfg.attn_softcap) {39        const float std = 0.02f;40        const int64_t C = cfg.d_model;41        const int64_t Cq = cfg.n_heads * head_dim_;42        const int64_t Ckv = cfg.n_kv_heads * head_dim_;43        const bool bias = cfg.attention_bias;44        wq_ = std::make_unique<Linear>(C, Cq, bias, std, rng);45        wk_ = std::make_unique<Linear>(C, Ckv, bias, std, rng);46        wv_ = std::make_unique<Linear>(C, Ckv, bias, std, rng);47        wo_ = std::make_unique<Linear>(Cq, C, false, proj_std, rng); // residual projection48        const ops::QuantMode qm = quant_mode_from(cfg.quant);49        wq_->set_quant(qm);50        wk_->set_quant(qm);51        wv_->set_quant(qm);52        wo_->set_quant(qm);53        absorb("wq", *wq_);54        absorb("wk", *wk_);55        absorb("wv", *wv_);56        absorb("wo", *wo_);57        if (qk_norm_) {58            q_norm_w_ = register_param("q_norm.weight",59                                       Tensor::full({head_dim_}, 1.0f));60            k_norm_w_ = register_param("k_norm.weight",61                                       Tensor::full({head_dim_}, 1.0f));62        }63        if (use_rope_) {64            const bool global = cfg.layer_is_global(layer_idx);65            const float theta = (global && cfg.rope_theta_global > 0.0f)66                                    ? cfg.rope_theta_global67                                    : cfg.rope_theta;68            freqs_ = ops::rope_freqs(head_dim_, theta, cfg.rope_scale_factor,69                                     cfg.rope_scale_low, cfg.rope_scale_high,70                                     cfg.rope_scale_orig_ctx);71        }72    }7374    Var forward(const Var& x) const override {75        const int64_t B = x.value().size(0), T = x.value().size(1), C = x.value().size(2);76        const int64_t Cq = n_heads_ * head_dim_;77        const int64_t Ckv = n_kv_heads_ * head_dim_;7879        Var x2d = x.reshaped({B * T, C});80        Var q = wq_->forward(x2d).reshaped({B, T, Cq});81        Var k = wk_->forward(x2d).reshaped({B, T, Ckv});82        Var v = wv_->forward(x2d).reshaped({B, T, Ckv});8384        if (qk_norm_) {85            // Per-head RMSNorm before RoPE (Qwen3/Gemma3): normalize over86            // head_dim by viewing heads as rows.87            q = ops::rmsnorm(q.reshaped({B * T * n_heads_, head_dim_}), q_norm_w_,88                             norm_eps_).reshaped({B, T, Cq});89            k = ops::rmsnorm(k.reshaped({B * T * n_kv_heads_, head_dim_}), k_norm_w_,90                             norm_eps_).reshaped({B, T, Ckv});91        }9293        if (use_rope_) {94            q = ops::rope(q, n_heads_, freqs_, 0);95            k = ops::rope(k, n_kv_heads_, freqs_, 0);96        }9798        const float scale = 1.0f / std::sqrt(float(head_dim_));99        Var o = ops::attention(q, k, v, n_heads_, n_kv_heads_, /*causal=*/true, scale,100                               window_, attn_softcap_);101        return wo_->forward(o.reshaped({B * T, Cq})).reshaped({B, T, C});102    }103104private:105    int64_t n_heads_, n_kv_heads_, head_dim_;106    bool use_rope_;107    bool qk_norm_;108    float norm_eps_;109    int64_t window_;     // 0 = full attention (global layer or SWA off)110    float attn_softcap_;111    std::unique_ptr<Linear> wq_, wk_, wv_, wo_;112    Var q_norm_w_, k_norm_w_; // defined only when qk_norm113    Tensor freqs_;            // [head_dim/2], defined only when use_rope114};115116} // namespace forge::nn117