246 lines
10 KiB
C
246 lines
10 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* rwkvx.h — RWKV-X inference engine built on Foundry
|
|
*
|
|
* Pure C implementation of the RWKV-X hybrid model (RWKV-7 recurrent blocks
|
|
* interleaved with sparse attention blocks). Loads safetensors weights,
|
|
* runs inference, and generates text.
|
|
*
|
|
* Prefix: rx_
|
|
*/
|
|
|
|
#ifndef RWKVX_H
|
|
#define RWKVX_H
|
|
|
|
#include <stdint.h>
|
|
#include <stdbool.h>
|
|
#include <stdio.h>
|
|
#include "tensor/tensor.h"
|
|
#include "hybridmem/hybridmem.h"
|
|
|
|
#ifdef __cplusplus
|
|
extern "C" {
|
|
#endif
|
|
|
|
/* ================================================================
|
|
* Types
|
|
* ================================================================ */
|
|
|
|
/* Opaque types */
|
|
typedef struct rx_model rx_model;
|
|
typedef struct rx_state rx_state;
|
|
typedef struct rx_tokenizer rx_tokenizer;
|
|
|
|
/* Model configuration (extracted from weight shapes) */
|
|
typedef struct {
|
|
int n_embd; /* embedding dim: 768 for 0.2B */
|
|
int n_head; /* attention heads: 12 for 0.2B */
|
|
int head_size; /* per-head dim: always 64 */
|
|
int n_rwkv_layer; /* RWKV-7 recurrent layers: 12 for 0.2B */
|
|
int n_attn_layer; /* sparse attention layers: 4 for 0.2B */
|
|
int vocab_size; /* vocabulary size: 65536 */
|
|
} rx_config;
|
|
|
|
/* HybridMem configuration for disk-backed attention memory */
|
|
typedef struct {
|
|
const char *store_dir; /* Directory for persistent storage (NULL = disabled) */
|
|
int max_cache_len; /* KV cache size per attention layer (0 = use default 2048) */
|
|
int top_k; /* Number of entries to retrieve from disk (default: 5) */
|
|
float evict_fraction; /* Fraction of cache to evict when full (default: 0.5) */
|
|
} rx_hybridmem_config;
|
|
|
|
/* Sampling parameters */
|
|
typedef struct {
|
|
float temperature; /* 0 = greedy */
|
|
float top_p; /* nucleus sampling threshold */
|
|
int top_k; /* top-k filtering (0 = disabled) */
|
|
float frequency_penalty;
|
|
float presence_penalty;
|
|
} rx_sampling_params;
|
|
|
|
/* ================================================================
|
|
* Model lifecycle
|
|
* ================================================================ */
|
|
|
|
/* Load model from safetensors file. Returns NULL on failure. */
|
|
rx_model* rx_model_load(const char *safetensors_path);
|
|
|
|
/* Load model from safetensors and quantize all linear weights to INT8.
|
|
* Reduces memory by ~4x and improves bandwidth-bound inference speed.
|
|
* Output quality may differ slightly from f32 — check greedy decode. */
|
|
rx_model* rx_model_load_q8(const char *safetensors_path);
|
|
|
|
/* Free model and all associated weights. */
|
|
void rx_model_free(rx_model *model);
|
|
|
|
/* Get model configuration. */
|
|
rx_config rx_model_config(const rx_model *model);
|
|
|
|
/* ================================================================
|
|
* State management
|
|
* ================================================================ */
|
|
|
|
/* Create a new zeroed state for the given model. Returns NULL on failure. */
|
|
rx_state* rx_state_new(const rx_model *model);
|
|
|
|
/* Create a new state with HybridMem disk-backed memory enabled.
|
|
* When the KV cache fills, old entries are evicted to disk via fm_put().
|
|
* Before attention, relevant past context is retrieved via fm_query().
|
|
* Returns NULL on failure. */
|
|
rx_state* rx_state_new_with_hybridmem(const rx_model *model,
|
|
const rx_hybridmem_config *hm_config);
|
|
|
|
/* Get default HybridMem config. */
|
|
rx_hybridmem_config rx_hybridmem_config_defaults(void);
|
|
|
|
/* Free state (also frees HybridMem store if present). */
|
|
void rx_state_free(rx_state *state);
|
|
|
|
/* Flush HybridMem store to disk. No-op if HybridMem is disabled.
|
|
* Returns 0 on success, -1 on error. */
|
|
int rx_state_flush(rx_state *state);
|
|
|
|
/* Reset state to initial zeroed values. */
|
|
void rx_state_reset(rx_state *state);
|
|
|
|
/* ================================================================
|
|
* Instrumentation (optional, zero-cost when disabled)
|
|
* ================================================================ */
|
|
|
|
/* Enable per-head attention weight + entropy logging to a CSV file.
|
|
* Pass NULL to disable. Resets the step counter.
|
|
* CSV format: token_step,layer,head,ret_mass,cache_mass,entropy */
|
|
void rx_state_set_attn_log(rx_state *state, FILE *fp);
|
|
|
|
/* Enable embedding capture from attention layers.
|
|
* buffer must have space for max_entries * n_embd floats.
|
|
* layer_idx: which attention layer to capture from (-1 = first one encountered).
|
|
* Call rx_state_embed_capture_count() after inference to get count. */
|
|
void rx_state_set_embed_capture(rx_state *state, float *buffer,
|
|
int max_entries, int layer_idx);
|
|
|
|
/* Get number of embeddings captured so far. */
|
|
int rx_state_embed_capture_count(const rx_state *state);
|
|
|
|
/* Get the HybridMem store from a state (for direct querying).
|
|
* Returns NULL if HybridMem is disabled. */
|
|
fm_store* rx_state_get_hm_store(const rx_state *state);
|
|
|
|
/* ================================================================
|
|
* Inference
|
|
* ================================================================ */
|
|
|
|
/* Forward pass for a single token. Returns logits [vocab_size].
|
|
* Caller must tm_free() the returned tensor. */
|
|
tm_tensor* rx_forward(rx_model *model, uint32_t token, rx_state *state);
|
|
|
|
/* Forward pass for a batch of tokens (prefill). Returns logits for the
|
|
* last token only: [vocab_size]. Caller must tm_free(). */
|
|
tm_tensor* rx_forward_batch(rx_model *model, const uint32_t *tokens,
|
|
int n_tokens, rx_state *state);
|
|
|
|
/* ================================================================
|
|
* Training support
|
|
* ================================================================ */
|
|
|
|
/* Block callback for training forward pass.
|
|
* Called BEFORE and AFTER each attention block.
|
|
* phase: 0 = before attention block, 1 = after attention block
|
|
* attn_idx: which attention block (0-based)
|
|
* x: current activation vector [n_embd], can be read/modified
|
|
* user_data: opaque pointer passed through from rx_forward_ex
|
|
*
|
|
* Return 0 to continue normal processing, 1 to skip this block
|
|
* (caller handles the block processing instead). */
|
|
typedef int (*rx_block_callback)(int phase, int attn_idx, float *x,
|
|
void *user_data);
|
|
|
|
/* Extended forward pass with block-level callbacks.
|
|
* Same as rx_forward but calls cb before/after each attention block.
|
|
* If cb returns 1 for phase=0, the attention block forward is SKIPPED
|
|
* and the caller is expected to process x in-place.
|
|
* logits_out: if non-NULL, receives pointer to logits data [vocab_size]
|
|
* (only valid until next rx_forward call — not a new allocation).
|
|
* Returns logits tensor (caller must tm_free()). */
|
|
tm_tensor* rx_forward_ex(rx_model *model, uint32_t token, rx_state *state,
|
|
rx_block_callback cb, void *user_data);
|
|
|
|
/* Get raw k_cache and v_cache pointers for an attention layer.
|
|
* Useful for training backward pass. Returns cache_len. */
|
|
int rx_state_get_kv_cache(const rx_state *state, int attn_idx,
|
|
const float **k_cache_out,
|
|
const float **v_cache_out);
|
|
|
|
/* ================================================================
|
|
* Weight accessors (for training backward pass)
|
|
*
|
|
* These expose internal weight pointers needed by rx_train.c.
|
|
* All returned pointers are owned by the model — do NOT free them.
|
|
* ================================================================ */
|
|
|
|
/* Get attention block weight pointers for a given attention layer index.
|
|
* All weights are [in_dim, out_dim] row-major (pre-transposed for fast_matvec).
|
|
* Returns 0 on success, -1 on invalid index. */
|
|
typedef struct {
|
|
const float *ln1_w, *ln1_b; /* [C] pre-attn LayerNorm */
|
|
const float *ln2_w, *ln2_b; /* [C] pre-FFN LayerNorm */
|
|
const float *receptance_w; /* [C, C] Q projection */
|
|
const float *key_w; /* [C, C] K projection */
|
|
const float *value_w; /* [C, C] V projection */
|
|
const float *output_w; /* [C, C] output projection */
|
|
const float *ffn_x_k; /* [C] token shift lerp */
|
|
const float *ffn_key_w; /* [C, ffn_dim] FFN key */
|
|
const float *ffn_value_w; /* [ffn_dim, C] FFN value */
|
|
int ffn_dim; /* FFN intermediate dimension */
|
|
} rx_attn_weights;
|
|
|
|
int rx_model_get_attn_weights(const rx_model *model, int attn_idx,
|
|
rx_attn_weights *out);
|
|
|
|
/* Get head projection weight pointer. Returns [C, vocab] row-major.
|
|
* Returns NULL if model is NULL. */
|
|
const float* rx_model_get_head_weight(const rx_model *model);
|
|
|
|
/* Get final layer norm weights. Returns 0 on success. */
|
|
int rx_model_get_ln_out(const rx_model *model,
|
|
const float **w_out, const float **b_out);
|
|
|
|
/* ================================================================
|
|
* Tokenizer
|
|
* ================================================================ */
|
|
|
|
/* Load RWKV World tokenizer from vocab file. Returns NULL on failure. */
|
|
rx_tokenizer* rx_tokenizer_new(const char *vocab_path);
|
|
|
|
/* Free tokenizer. */
|
|
void rx_tokenizer_free(rx_tokenizer *tok);
|
|
|
|
/* Encode text into token IDs. Returns number of tokens written.
|
|
* tokens_out must have space for at least max_tokens entries. */
|
|
int rx_tokenizer_encode(const rx_tokenizer *tok, const char *text,
|
|
uint32_t *tokens_out, int max_tokens);
|
|
|
|
/* Decode token IDs into text. Returns malloc'd string (caller must free). */
|
|
char* rx_tokenizer_decode(const rx_tokenizer *tok, const uint32_t *tokens,
|
|
int n_tokens);
|
|
|
|
/* ================================================================
|
|
* Sampling
|
|
* ================================================================ */
|
|
|
|
/* Return default sampling parameters. */
|
|
rx_sampling_params rx_sampling_defaults(void);
|
|
|
|
/* Sample a token from logits using the given parameters.
|
|
* token_counts: per-token occurrence counts for penalties (can be NULL).
|
|
* vocab_size: size of the logits/token_counts arrays. */
|
|
uint32_t rx_sample(const tm_tensor *logits, rx_sampling_params *params,
|
|
const uint32_t *token_counts, int vocab_size);
|
|
|
|
#ifdef __cplusplus
|
|
}
|
|
#endif
|
|
|
|
#endif /* RWKVX_H */
|