multihead_attention_module.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <optional>
#include <memory>
#include <vector>
 
#include "pulsatrix/linear_module.hpp"
#include "pulsatrix/module.hpp"
#include "pulsatrix/rms_norm_module.hpp"
#include "pulsatrix/rope_module.hpp"
#include "pulsatrix/softmax_module.hpp"
 
namespace pulsatrix {
 
class MultiHeadAttentionModule : public Module {
public:
    MultiHeadAttentionModule(int64_t d_model, int64_t num_heads, DeviceBackend* backend, bool use_rope = true,
                             bool use_qk_norm = false);
 
    [[nodiscard]] Tensor backward(const Tensor& grad_output) override;
 
    [[nodiscard]] OpType op_type() const override { return OpType::Attention; }
 
    [[nodiscard]] Tensor propagate_relevance(const Tensor& relevance_out, const LRPRuleConfig& config) override;
 
    [[nodiscard]] std::vector<NamedParamRef> named_parameters() override;
 
    void set_training(bool training) override;
 
    [[nodiscard]] int64_t d_model() const { return d_model_; }
    [[nodiscard]] int64_t num_heads() const { return num_heads_; }
    [[nodiscard]] int64_t head_dim() const { return head_dim_; }
    [[nodiscard]] bool uses_rope() const { return use_rope_; }
    [[nodiscard]] bool uses_qk_norm() const { return use_qk_norm_; }
 
    [[nodiscard]] LinearModule& q_proj() { return q_proj_; }
    [[nodiscard]] LinearModule& k_proj() { return k_proj_; }
    [[nodiscard]] LinearModule& v_proj() { return v_proj_; }
    [[nodiscard]] LinearModule& out_proj() { return out_proj_; }
    [[nodiscard]] RMSNormModule* q_norm() { return q_norm_.get(); }
    [[nodiscard]] RMSNormModule* k_norm() { return k_norm_.get(); }
 
    [[nodiscard]] const Tensor& last_attention_weights() const { return last_attn_; }
 
 
    [[nodiscard]] std::optional<DeviceType> compute_device() const override { return backend_->device(); }
 
protected:
    [[nodiscard]] Tensor forward_impl(const Tensor& input) override;
 
private:
    int64_t d_model_;
    int64_t num_heads_;
    int64_t head_dim_;
    bool use_rope_;
    bool use_qk_norm_;
    DeviceBackend* backend_;
 
    LinearModule q_proj_;
    LinearModule k_proj_;
    LinearModule v_proj_;
    LinearModule out_proj_;
    SoftmaxModule softmax_;
    // Held by pointer rather than std::optional purely so the "disabled" state costs nothing
    // and needs no move/copy of a Module subclass (Module declares a virtual destructor,
    // which suppresses implicit move construction -- std::optional's in-place paths would
    // work but the pointer is unambiguous).
    std::unique_ptr<RoPEModule> q_rope_;
    std::unique_ptr<RoPEModule> k_rope_;
    std::unique_ptr<RMSNormModule> q_norm_;
    std::unique_ptr<RMSNormModule> k_norm_;
 
    // Forward caches. Q/K/V are the post-QK-Norm, post-RoPE, head-split values -- exactly
    // the operands the two matmuls' backward and Eq. 15 rules need.
    int64_t last_N_ = 0;
    int64_t last_L_ = 0;
    Tensor last_q_;           
    Tensor last_k_;           
    Tensor last_v_;           
    Tensor last_scores_raw_;  
    Tensor last_attn_;        
    Tensor last_context_;     
    bool has_forwarded_ = false;
};
 
}  // namespace pulsatrix