#pragma once
#include "ctranslate2/layers/common.h"
#include "ctranslate2/padder.h"
namespace ctranslate2 {
namespace layers {
StorageView make_relative_positions(dim_t queries_length,
dim_t keys_length,
dim_t max_position);
class RotaryEmbeddings;
class Alibi;
class AttentionLayer : public Layer
{
public:
AttentionLayer(const models::Model& model,
const std::string& scope,
dim_t num_heads,
bool self_attention,
bool pre_norm = true,
bool is_decoder = false,
Alibi* alibi = nullptr,
bool is_flash_attn = false);
virtual ~AttentionLayer() {};
DataType output_type() const override;
dim_t output_size() const override;
virtual void operator()(const StorageView& queries,
const StorageView& values,
const StorageView* values_lengths,
StorageView& output,
StorageView* cached_keys = nullptr,
StorageView* cached_values = nullptr,
StorageView* attention = nullptr,
const Padder* queries_padder = nullptr,
const Padder* values_padder = nullptr,
bool return_normalized_attention = true,
StorageView* position_bias = nullptr,
dim_t offset = 0) const = 0;
virtual bool has_positional_embeddings() const = 0;
bool multi_query() const {
return _multi_query;
}
static StorageView prepare_length_mask(const StorageView& lengths,
const dim_t num_heads,
const dim_t num_queries,
const bool mask_future = false,
const bool multi_query = false);
protected:
const bool _tensor_parallel;
const dim_t _num_heads;
const bool _self_attention;
const bool _is_decoder;
const std::vector<Dense> _linear;
const dim_t _d_model;
const dim_t _d_head;
const bool _pre_norm;
const std::unique_ptr<const LayerNorm> _layer_norm;
const std::unique_ptr<RotaryEmbeddings> _rotary_embeddings;
Alibi* _alibi;
const float _queries_scale;
const bool _multi_query;
const dim_t _num_heads_kv;
const dim_t _sliding_window;
};
enum class RotaryScalingType {
None = -1,
Linear,
Su,
Llama3,
};
class RotaryEmbeddings {
public:
RotaryEmbeddings(const dim_t dim = 0,
const bool interleave = true,
const RotaryScalingType scaling_type = RotaryScalingType::None,
const float scaling_factor = 1,
const float base = 10000,
const dim_t num_initial_positions = 2048,
const StorageView* long_scaling_factor = nullptr,
const StorageView* short_scaling_factor = nullptr,
const float low_freq_factor = 1.0,
const float high_freq_factor = 4.0,
const dim_t original_max_position_embeddings = 0,
const dim_t max_position_embeddings = 0,
const bool transpose = true);
void apply(StorageView& x, const dim_t offset = 0, bool fa2 = false);
StorageView& get_cos_half() {
return *_cos_half;
}
StorageView& get_sin_half() {
return *_sin_half;
}
bool get_interleave() const {
return _interleave;
}
private:
void initialize(const dim_t num_positions,
const dim_t dim,
const Device device,
const DataType dtype);
const dim_t _dim;
const bool _interleave;
const RotaryScalingType _scaling_type;
const float _scaling_factor;
const float _base;
const dim_t _num_initial_positions;
std::unique_ptr<StorageView> _rotary_scaling_long_factor;
std::unique_ptr<StorageView> _rotary_scaling_short_factor;
const float _rotary_low_freq_factor;
const float _rotary_high_freq_factor;
const dim_t _original_max_position_embeddings;
const dim_t _max_position_embeddings;
const ops::Rotary _rotary_op;
const bool _transpose;
StorageView _sin;
StorageView _cos;
std::unique_ptr<StorageView> _sin_half;
std::unique_ptr<StorageView> _cos_half;
};
class Alibi {
public:
Alibi(const bool use_positive_positions = false, const bool scale_alibi = false, const dim_t num_initial_positions = 2048);
void apply(StorageView& x, const float scale = 1);
private:
const bool _use_positive_positions;
const dim_t _num_initial_positions;
const bool _scale_alibi;
const ops::AlibiAdd _alibi_op;
StorageView _alibi;
};
}
}