File size: 1,967 Bytes
ce25ab4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | // mamba_bind.h — Mamba Binding Layer for SM89 sovereign kernels
// Author: Ahmad Ali Parr · Trust: Bel Esprit D'Accord Irrevocable Trust
#ifndef MAMBA_BIND_H
#define MAMBA_BIND_H
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cstdint>
struct MambaAttentionGEMMPipeline;
using PipelineHandle = MambaAttentionGEMMPipeline*;
struct MambaParams {
half* A_log; // [d_state, d_inner]
half* B; // [d_state, d_inner]
half* C; // [d_state, d_inner]
half* D; // [d_inner]
half* dt_proj; // [d_inner]
int d_inner, d_state, d_conv;
float dt_min, dt_max, dt_init_floor;
};
struct KVPageManager {
uint64_t* page_table;
half* kv_pages;
int* page_refcount;
int* free_list;
int free_head, free_tail, total_pages, page_size, head_dim;
};
struct GGUFQuantization {
enum Type { Q4_0, Q4_K, Q5_0, Q8_0 } type;
const void* packed_data;
const half* scales;
const half* mins;
int block_size;
int num_blocks;
};
struct MambaPerfMetrics {
float attn_ms, gemm_ms, mamba_ms, copy_ms, dequant_ms, total_ms;
float attn_tflops, gemm_tflops;
int64_t bytes_moved;
};
#ifdef __cplusplus
extern "C" {
#endif
PipelineHandle mamba_create_pipeline(int batch, int heads, int seq, int dim, int pages);
void mamba_destroy_pipeline(PipelineHandle h);
void mamba_execute(PipelineHandle h);
void mamba_execute_async(PipelineHandle h, int steps);
void mamba_launch_flash_attention(
const half* q, const half* k, const half* v, half* o,
const uint64_t* page_table,
int batch, int heads, int seq, int dim,
int page_size, int num_pages, cudaStream_t stream
);
void mamba_launch_dequant_gguf(
const GGUFQuantization* quant,
half* output,
cudaStream_t stream
);
MambaPerfMetrics mamba_benchmark(PipelineHandle h, int warmup, int iterations);
#ifdef __cplusplus
}
#endif
#endif // MAMBA_BIND_H
|