File size: 5,965 Bytes
51250cb | 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 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 | #pragma once
#ifndef __cplusplus
#include <stdatomic.h>
#endif
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include "internal/metal.h"
struct gptoss_tokenizer {
#ifndef __cplusplus
atomic_uint_least64_t ref_count;
#else
uint_least64_t ref_count;
#endif
void* mapping_ptr;
size_t mapping_size;
const char* regex_ptr;
const char* tokens_ptr;
uint32_t num_text_tokens;
uint32_t num_special_tokens;
uint32_t special_token_id[gptoss_special_token_max - 1];
};
struct gptoss_model {
#ifndef __cplusplus
atomic_uint_least64_t ref_count;
#else
uint_least64_t ref_count;
#endif
struct gptoss_tokenizer* tokenizer;
void* mapping_ptr;
size_t mapping_size;
uint32_t context_length;
uint32_t num_blocks;
uint32_t num_experts;
uint32_t num_active_experts;
uint32_t embedding_dim;
uint32_t mlp_dim;
float swiglu_limit;
uint32_t head_dim;
uint32_t num_heads;
uint32_t num_kv_heads;
uint32_t attention_window;
float rope_theta;
float interpolation_scale;
float yarn_offset;
float yarn_scale;
float yarn_multiplier;
float rmsnorm_epsilon;
uint32_t vocabulary_size;
bool lock_memory;
size_t weights_size;
size_t allocation_size;
// Metal objects
struct gptoss_metal_device device;
size_t max_threadgroups;
struct gptoss_metal_command_queue command_queue;
struct gptoss_metal_library library;
struct gptoss_metal_function bf16_f32_embeddings_fn;
struct gptoss_metal_function f32_bf16w_rmsnorm_fn;
struct gptoss_metal_function f32_bf16w_matmul_fn;
struct gptoss_metal_function f32_bf16w_matmul_qkv_fn;
struct gptoss_metal_function f32_bf16w_dense_matmul_qkv_fn;
struct gptoss_metal_function f32_bf16w_dense_matmul_attn_output_fn;
struct gptoss_metal_function f32_bf16w_dense_matmul_mlp_gate_fn;
struct gptoss_metal_function f32_bf16w_unembedding_fn;
struct gptoss_metal_function f32_rope_fn;
struct gptoss_metal_function f32_mf4w_moe_matmul_swiglu_fn;
struct gptoss_metal_function f32_mf4w_moe_matmul_fn;
struct gptoss_metal_function f32_accumulate_e4_fn;
struct gptoss_metal_function f32_scatter_e4_fn;
struct gptoss_metal_function f32_mf4w_moe_dense_matmul_swiglu_fn;
struct gptoss_metal_function f32_mf4w_moe_dense_matmul_fn;
struct gptoss_metal_function f32_gather_and_accumulate_e4_fn;
struct gptoss_metal_function f32_expert_routing_metadata_fn;
struct gptoss_metal_function f32_topk_softmax_e32_k4_fn;
struct gptoss_metal_function f32_topk_softmax_e128_k4_fn;
struct gptoss_metal_function f32_sdpa_q8_d64_fn;
struct gptoss_metal_function f32_softmax_fn;
struct gptoss_metal_function f32_sample_fn;
size_t per_block_shared_weights_size;
size_t per_expert_block_weight_size;
size_t embeddings_threadgroup_size;
size_t attn_qkv_threadgroup_size;
size_t attn_out_threadgroup_size;
size_t mlp_gate_threadgroup_size;
size_t mlp_swiglu_threadgroup_size;
size_t mlp_out_threadgroup_size;
size_t mlp_acc_threadgroup_size;
size_t unembedding_threadgroup_size;
size_t attn_rmsnorm_gain_offset;
size_t attn_qkv_weight_offset;
size_t attn_qkv_bias_offset;
size_t attn_sdpa_sink_offset;
size_t attn_out_weight_offset;
size_t attn_out_bias_offset;
size_t mlp_rmsnorm_gain_offset;
size_t mlp_gate_weight_offset;
size_t mlp_gate_bias_offset;
size_t mlp_swiglu_scale_offset;
size_t mlp_swiglu_bias_offset;
size_t mlp_out_block_offset;
size_t mlp_out_scale_offset;
size_t mlp_out_bias_offset;
size_t rmsnorm_weight_offset;
size_t unembedding_weight_offset;
// Buffer with non-MoE weights. Includes MoE gates, embeddings/unembeddings.
struct gptoss_metal_buffer shared_weight_buffer;
// num_blocks per-block buffers with MoE weights to follow.
struct gptoss_metal_buffer block_weight_buffers[];
};
#define GPTOSS_DEFAULT_BATCH_SIZE 128
struct gptoss_context {
#ifndef __cplusplus
atomic_uint_least64_t ref_count;
#else
uint_least64_t ref_count;
#endif
struct gptoss_model* model;
// Number of tokens processed in the context.
size_t num_tokens;
// Number of tokens in the KV cache.
size_t num_kv_tokens;
// Length of the context.
size_t max_tokens;
// Maximum number of tokens that can be processed in a single batch.
// Activation buffers are allocated with this size.
size_t max_batch_tokens;
size_t kvcache_size;
size_t allocation_size;
// Activation buffers.
// TODO: merge into a single buffer.
struct gptoss_metal_buffer residual_activation_buffer; // Residual stream
struct gptoss_metal_buffer rmsnorm_activation_buffer; // Both attention & MLP RMSNorm output
struct gptoss_metal_buffer qkv_activation_buffer; // QKV projection output
struct gptoss_metal_buffer sdpa_activation_buffer; // SDPA output
struct gptoss_metal_buffer gate_activation_buffer; // MoE gating output
struct gptoss_metal_buffer expert_activation_buffer; // MoE expert predictions
struct gptoss_metal_buffer expert_offset_buffer; // MoE expert histograms cumsum
struct gptoss_metal_buffer token_to_expert_routing_buffer; // MoE token to expert routing
struct gptoss_metal_buffer swiglu_input_buffer; // MLP+SwiGLU input for prefill.
struct gptoss_metal_buffer swiglu_activation_buffer; // MLP+SwiGLU output
struct gptoss_metal_buffer moe_activation_buffer; // MoE MLP output (per-active expert)
// Input/output buffers.
struct gptoss_metal_buffer control_buffer;
struct gptoss_metal_buffer token_buffer; // uint32 token IDs
struct gptoss_metal_buffer score_buffer; // unembedding outputs
struct gptoss_metal_buffer prob_buffer;
struct gptoss_metal_buffer sum_buffer;
struct gptoss_metal_buffer argmax_buffer;
struct gptoss_metal_buffer kvcache_buffer;
};
|