grouped_query_attention_layer.h

Go to the documentation of this file.

// SPDX-License-Identifier: AGPL-3.0-only
// Copyright (C) 2026 Artificial Intelligence Techniques, S.L. and Roberto López González
 
#pragma once
 
#include "opennn/network/layers/layer.h"
#include "opennn/network/operators/operator.h"
 
namespace opennn
{
 
void rotary_build_tables(TensorView&, TensorView&, Index sequence_length, Index rotary_dim, float base);
void rotary_forward(const TensorView&, const TensorView&, const TensorView&,
                    TensorView&, Index head_dim, Index rotary_dim, Index position_offset);
void grouped_attention_forward(const TensorView& query, const TensorView& key, const TensorView& value,
                               TensorView& output, Index n_query_heads, Index n_kv_heads, Index head_dim,
                               bool causal, float scale, Index query_position_offset = 0,
                               float* decode_partials = nullptr, const int* position_device = nullptr);
Index grouped_attention_decode_scratch_floats(Index n_query_heads, Index head_dim);
void qk_rope_cache_append(const TensorView& qkv_row, const TensorView& q_norm_weight,
                          const TensorView& k_norm_weight, const TensorView& cos_table,
                          const TensorView& sin_table, TensorView& q_out,
                          TensorView& key_cache, TensorView& value_cache,
                          Index n_query_heads, Index n_kv_heads, Index head_dim,
                          float epsilon, const int* position_device);
void qk_norm_forward(const TensorView& input, const TensorView& weight, TensorView& output,
                     Index head_dim, float epsilon);
 
struct GroupedQueryAttentionOperator : Operator
{
    static constexpr size_t forward_scratch_slots_count = 9;
    struct GraphCache;
 
    GroupedQueryAttentionOperator();
    ~GroupedQueryAttentionOperator() override;
 
    void set(Index new_sequence_length, Index new_hidden,
             Index new_q_heads, Index new_kv_heads, Index new_head_dim,
             float new_rope_theta, float new_rms_epsilon, bool new_use_qk_norm);
 
    vector<TensorSpec> parameter_specs() const override;
    vector<TensorSpec> forward_scratch_specs() const;
    vector<SlotQuantization> parameter_quantization() const override;
    void link_parameters(span<const TensorView>) override;
    void link_parameter_scales(span<const TensorView>) override;
    void set_parameters_random() override;
 
    const Buffer* get_rope_tables() const noexcept { return rope_tables.get(); }
 
    void forward_propagate(ForwardPropagation&, size_t, ForwardPropagationMode) override;
    void back_propagate(ForwardPropagation&, BackPropagation&, size_t) const override;
 
private:
 
    friend class GroupedQueryAttention;
 
    enum ForwardSlot : size_t
    {
        Input,
        Query,
        Key,
        Value,
        RotatedQuery,
        RotatedKey,
        AttentionOutput,
        FusedQkv,
        DecodePartials,
        SequenceLengths,
        Output
    };
 
    Index q_dim()  const { return q_heads  * head_dim; }
    Index kv_dim() const { return kv_heads * head_dim; }
 
    void prepare_rope_tables(Device);
    void apply_attention(const TensorView& query, const TensorView& key,
                         const TensorView& value, TensorView& output,
                         bool causal, float scale, Index query_position_offset,
                         float* decode_partials = nullptr,
                         const int* position_device = nullptr);
    void attend_sequence_cpu(const TensorView& x_b,
                             TensorView& q_v, TensorView& k_v,
                             TensorView& v_target, TensorView& k_target,
                             TensorView& qr_v,
                             const TensorView& key_all, const TensorView& value_all,
                             TensorView& attn_v, TensorView& o_b,
                             const TensorView& cos_v, const TensorView& sin_v,
                             float scale, Index position_offset);
 
    void forward_gpu(TensorView& input, TensorView& output, Index batch, Index past,
                     Index query_capacity, const int* position_device,
                     vector<TensorView>& forward_slots, Buffer& kv_cache,
                     device::PinnedBuffer& pinned_storage);
 
    Index sequence_length = 0;
    Index hidden          = 0;
    Index q_heads         = 0;
    Index kv_heads        = 0;
    Index head_dim        = 0;
    float rope_theta      = 1000000.0f;
    float rms_epsilon     = 1.0e-6f;
    bool use_qk_norm      = true;
 
    TensorView q_proj, k_proj, v_proj, o_proj, q_norm, k_norm;
    TensorView q_scale, k_scale, v_scale, o_scale, qkv_scale;
    bool qkv_fused = false;
 
    shared_ptr<const Buffer> rope_tables;
    unique_ptr<GraphCache> graph_cache;
};
 
class GroupedQueryAttention final : public Layer
{
public:
 
    GroupedQueryAttention(const Shape& = Shape({0, 0}),
                          Index q_heads = 0, Index kv_heads = 0, Index head_dim = 0,
                          float rope_theta = 1000000.0f, float rms_epsilon = 1.0e-6f,
                          bool use_qk_norm = true,
                          const string& = "grouped_query_attention");
 
    Shape get_input_shape() const noexcept override { return { sequence_length, hidden }; }
    Shape get_output_shape() const override { return { sequence_length, hidden }; }
 
    Index get_sequence_length() const { return sequence_length; }
    Index get_hidden() const { return hidden; }
    Index get_q_heads() const { return q_heads; }
    Index get_kv_heads() const { return kv_heads; }
    Index get_head_dim() const { return head_dim; }
    bool  get_use_qk_norm() const { return use_qk_norm; }
 
    vector<TensorSpec> get_forward_specs(Index) const override;
 
    using Layer::is_forward_slot_inference_elidable;
    bool is_forward_slot_inference_elidable(size_t, Device, Index) const noexcept;
    bool uses_compact_inference() const noexcept;
 
    ForwardSlotKind get_forward_slot_kind(size_t slot) const override
    {
        return slot != Input
            && slot <= GroupedQueryAttentionOperator::forward_scratch_slots_count
            ? ForwardSlotKind::Transient
            : ForwardSlotKind::Pooled;
    }
 
    void set(const Shape&, Index, Index, Index, float, float, bool new_use_qk_norm, const string&);
    bool accepts_input_rank(Index rank) const override { return is_one_of(rank, 2, 3); }
    bool uses_sequence_position() const noexcept override { return true; }
 
    void apply_input_shape(const Shape&) override;
 
    void on_compute_dtype_changed() override { attention.compute_dtype = get_compute_dtype(); }
 
    void read_JSON_body(const Json*) override;
    void write_JSON_body(JsonWriter&) const override;
 
private:
 
    friend class Qwen3;
    void enable_compact_inference();
 
    Index sequence_length = 0;
    Index hidden          = 0;
    Index q_heads         = 0;
    Index kv_heads        = 0;
    Index head_dim        = 0;
    float rope_theta      = 1000000.0f;
    float rms_epsilon     = 1.0e-6f;
    bool  use_qk_norm     = true;
 
    GroupedQueryAttentionOperator attention;
};
 
}