123 lines
4.1 KiB
C
123 lines
4.1 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
#ifndef FOUNDRY_ATTENTION_H
|
|
#define FOUNDRY_ATTENTION_H
|
|
|
|
#include <stdint.h>
|
|
#include <stdbool.h>
|
|
#include <stddef.h>
|
|
#include "../tensor/tensor.h"
|
|
|
|
#ifdef __cplusplus
|
|
extern "C" {
|
|
#endif
|
|
|
|
/*
|
|
* Efficient attention implementations.
|
|
* Replaces FlashAttention.
|
|
*/
|
|
|
|
/* ============================================================= */
|
|
/* Types */
|
|
/* ============================================================= */
|
|
|
|
/* Attention configuration */
|
|
typedef struct {
|
|
int num_heads;
|
|
int head_dim;
|
|
float dropout;
|
|
bool causal; /* Causal masking for autoregressive */
|
|
float scale; /* Attention scale (default: 1/sqrt(head_dim)) */
|
|
} fa_config;
|
|
|
|
/* Attention context for caching */
|
|
typedef struct fa_context fa_context;
|
|
|
|
/* ============================================================= */
|
|
/* Standard Attention */
|
|
/* ============================================================= */
|
|
|
|
/* Scaled dot-product attention: softmax(QK^T / sqrt(d)) * V
|
|
* Q, K, V: (batch, seq_len, num_heads, head_dim)
|
|
* Returns: (batch, seq_len, num_heads, head_dim)
|
|
*/
|
|
tm_tensor* fa_attention(tm_tensor* Q, tm_tensor* K, tm_tensor* V, fa_config* config);
|
|
|
|
/* Attention with explicit mask
|
|
* mask: (batch, 1, seq_len, seq_len) or (1, 1, seq_len, seq_len)
|
|
* -inf values are masked out
|
|
*/
|
|
tm_tensor* fa_attention_mask(tm_tensor* Q, tm_tensor* K, tm_tensor* V,
|
|
tm_tensor* mask, fa_config* config);
|
|
|
|
/* ============================================================= */
|
|
/* Multi-Head Attention */
|
|
/* ============================================================= */
|
|
|
|
/* Multi-head attention with projections
|
|
* x: (batch, seq_len, embed_dim)
|
|
* Wq, Wk, Wv: (embed_dim, num_heads * head_dim)
|
|
* Wo: (num_heads * head_dim, embed_dim)
|
|
*/
|
|
tm_tensor* fa_multi_head(tm_tensor* x, tm_tensor* Wq, tm_tensor* Wk, tm_tensor* Wv,
|
|
tm_tensor* Wo, fa_config* config);
|
|
|
|
/* Multi-head self-attention (Q=K=V from same source) */
|
|
tm_tensor* fa_self_attention(tm_tensor* x, tm_tensor* Wqkv, tm_tensor* Wo,
|
|
fa_config* config);
|
|
|
|
/* ============================================================= */
|
|
/* KV Cache (for inference) */
|
|
/* ============================================================= */
|
|
|
|
/* Create KV cache context */
|
|
fa_context* fa_context_new(int batch_size, int max_seq_len, int num_heads,
|
|
int head_dim, int num_layers);
|
|
|
|
/* Free context */
|
|
void fa_context_free(fa_context* ctx);
|
|
|
|
/* Attention with KV cache (for autoregressive generation)
|
|
* Updates cache in-place, returns attention output
|
|
*/
|
|
tm_tensor* fa_attention_cached(tm_tensor* Q, tm_tensor* K, tm_tensor* V,
|
|
fa_context* ctx, int layer_idx, int pos,
|
|
fa_config* config);
|
|
|
|
/* Clear cache */
|
|
void fa_context_clear(fa_context* ctx);
|
|
|
|
/* Get current sequence length in cache */
|
|
int fa_context_seq_len(fa_context* ctx);
|
|
|
|
/* ============================================================= */
|
|
/* Rotary Position Embedding (RoPE) */
|
|
/* ============================================================= */
|
|
|
|
/* Apply rotary position embedding to Q and K
|
|
* Modifies tensors in-place
|
|
*/
|
|
void fa_rope(tm_tensor* Q, tm_tensor* K, int start_pos, int head_dim);
|
|
|
|
/* Precompute RoPE frequencies */
|
|
tm_tensor* fa_rope_freqs(int max_seq_len, int head_dim);
|
|
|
|
/* ============================================================= */
|
|
/* Utility */
|
|
/* ============================================================= */
|
|
|
|
/* Default config */
|
|
fa_config fa_config_default(int num_heads, int head_dim);
|
|
|
|
/* Create causal mask for sequence length */
|
|
tm_tensor* fa_causal_mask(int seq_len);
|
|
|
|
/* Print attention info */
|
|
void fa_print_config(fa_config* config);
|
|
|
|
#ifdef __cplusplus
|
|
}
|
|
#endif
|
|
|
|
#endif /* FOUNDRY_ATTENTION_H */
|