TinyLLMInference.h

Go to the documentation of this file.

 
#pragma once
 
#include "Assets/TinyLLMAsset.h"  // For TinyLLMConfig
#include <stdint.h>
#include <stddef.h>
 
#ifdef __cplusplus
extern "C" {
#endif
 
void tinyllm_rmsnorm(float* o, float* x, float* weight, int size);
 
void tinyllm_softmax(float* x, int size);
 
void tinyllm_matmul(float* xout, float* x, float* w, int n, int d);
 
float* tinyllm_forward(
    float* x, float* xb, float* xb2, float* hb, float* hb2,
    float* q, float* key_cache, float* value_cache, float* att, float* logits,
    float* token_emb, float* rms_att, float* wq, float* wk, float* wv, float* wo,
    float* rms_ffn, float* w1, float* w2, float* w3, float* rms_final, float* wcls,
    const TinyLLMConfig* config,
    int token, int pos
);
 
size_t tinyllm_calc_runstate_size(const TinyLLMConfig* config, int max_seq_len);
 
size_t tinyllm_calc_weights_size(const TinyLLMConfig* config);
 
#ifdef __cplusplus
}
#endif