Something went wrong. Try again.
Anonymize your writing style. Zig WASM engine detects authorship markers, fine-tuned LLM rewrites to remove them. Runs entirely in-browser. fantasma.qstorage.quilibrium.com
wasm privacy qwen zig
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203const std = @import("std");
// ============================================================// Model configuration constants (Qwen3.5-0.8B)// ============================================================pub const Config = struct { pub const vocab_size: u32 = 248320; pub const hidden_size: u32 = 1024; pub const num_layers: u32 = 24; pub const rms_norm_eps: f32 = 1e-6;
// Full attention config pub const num_attention_heads: u32 = 8; // query heads (actual, not including gate) pub const num_kv_heads: u32 = 2; pub const head_dim: u32 = 256; pub const attn_output_dim: u32 = num_attention_heads * head_dim; // 2048
// RoPE pub const partial_rotary_factor: f32 = 0.25; pub const rotary_dim: u32 = @intFromFloat(@as(f32, head_dim) * partial_rotary_factor); // 64 pub const rope_theta: f32 = 10_000_000.0; pub const mrope_section: [3]u32 = .{ 11, 11, 10 }; // total = 32 = rotary_dim/2
// GatedDeltaNet (linear attention) config pub const linear_num_heads: u32 = 16; pub const linear_key_head_dim: u32 = 128; pub const linear_value_head_dim: u32 = 128; pub const linear_key_dim: u32 = linear_num_heads * linear_key_head_dim; // 2048 pub const linear_value_dim: u32 = linear_num_heads * linear_value_head_dim; // 2048 pub const linear_conv_dim: u32 = linear_key_dim * 2 + linear_value_dim; // 6144 pub const linear_conv_kernel: u32 = 4;
// MLP pub const intermediate_size: u32 = 3584;
// Special tokens pub const eos_token_id: u32 = 248046; // <|im_end|> pub const im_start_id: u32 = 248045; pub const im_end_id: u32 = 248046;
// Layer types: false = linear_attention, true = full_attention pub const layer_is_full_attn: [num_layers]bool = .{ false, false, false, true, // 0-3 false, false, false, true, // 4-7 false, false, false, true, // 8-11 false, false, false, true, // 12-15 false, false, false, true, // 16-19 false, false, false, true, // 20-23 };};
// ============================================================// Cache for autoregressive generation// ============================================================
/// KV cache for full attention layers./// Stores key and value vectors for all past positions.pub const KVCache = struct { /// [max_seq_len][num_kv_heads * head_dim] keys: []f32, /// [max_seq_len][num_kv_heads * head_dim] values: []f32, len: u32, max_len: u32,
pub fn init(allocator: std.mem.Allocator, max_seq_len: u32) !KVCache { const kv_dim = Config.num_kv_heads * Config.head_dim; return .{ .keys = try allocator.alloc(f32, max_seq_len * kv_dim), .values = try allocator.alloc(f32, max_seq_len * kv_dim), .len = 0, .max_len = max_seq_len, }; }
pub fn deinit(self: *KVCache, allocator: std.mem.Allocator) void { allocator.free(self.keys); allocator.free(self.values); }
/// Append a new key/value pair for the current position. pub fn append(self: *KVCache, key: []const f32, value: []const f32) void { const kv_dim = Config.num_kv_heads * Config.head_dim; const offset = self.len * kv_dim; @memcpy(self.keys[offset .. offset + kv_dim], key); @memcpy(self.values[offset .. offset + kv_dim], value); self.len += 1; }
/// Get the key vector at a given position. pub fn getKey(self: *const KVCache, pos: u32) []const f32 { const kv_dim = Config.num_kv_heads * Config.head_dim; const offset = pos * kv_dim; return self.keys[offset .. offset + kv_dim]; }
/// Get the value vector at a given position. pub fn getValue(self: *const KVCache, pos: u32) []const f32 { const kv_dim = Config.num_kv_heads * Config.head_dim; const offset = pos * kv_dim; return self.values[offset .. offset + kv_dim]; }
pub fn reset(self: *KVCache) void { self.len = 0; }};
/// Recurrent state for GatedDeltaNet (linear attention) layers./// Stores the recurrent state matrix S[num_heads][key_dim][value_dim]/// and the conv1d state buffer.pub const DeltaNetState = struct { /// Recurrent state: [num_heads * key_dim * value_dim] recurrent: []f32, /// Conv1d state: [conv_dim * (kernel_size - 1)] conv_state: []f32,
pub fn init(allocator: std.mem.Allocator) !DeltaNetState { const state_size = Config.linear_num_heads * Config.linear_key_head_dim * Config.linear_value_head_dim; const conv_size = Config.linear_conv_dim * (Config.linear_conv_kernel - 1); const recurrent = try allocator.alloc(f32, state_size); const conv_state = try allocator.alloc(f32, conv_size); @memset(recurrent, 0.0); @memset(conv_state, 0.0); return .{ .recurrent = recurrent, .conv_state = conv_state, }; }
pub fn deinit(self: *DeltaNetState, allocator: std.mem.Allocator) void { allocator.free(self.recurrent); allocator.free(self.conv_state); }
pub fn reset(self: *DeltaNetState) void { @memset(self.recurrent, 0.0); @memset(self.conv_state, 0.0); }
/// Get a mutable slice of the recurrent state for a given head. /// Returns [key_dim][value_dim] as a flat slice of length key_dim * value_dim. pub fn getHeadState(self: *DeltaNetState, head: u32) []f32 { const head_state_size = Config.linear_key_head_dim * Config.linear_value_head_dim; const offset = head * head_state_size; return self.recurrent[offset .. offset + head_state_size]; }};
/// Combined cache for all layers.pub const ModelCache = struct { kv_caches: [Config.num_layers]?KVCache, delta_states: [Config.num_layers]?DeltaNetState,
pub fn init(allocator: std.mem.Allocator, max_seq_len: u32) !ModelCache { var cache = ModelCache{ .kv_caches = .{null} ** Config.num_layers, .delta_states = .{null} ** Config.num_layers, };
for (0..Config.num_layers) |i| { if (Config.layer_is_full_attn[i]) { cache.kv_caches[i] = try KVCache.init(allocator, max_seq_len); } else { cache.delta_states[i] = try DeltaNetState.init(allocator); } }
return cache; }
pub fn deinit(self: *ModelCache, allocator: std.mem.Allocator) void { for (0..Config.num_layers) |i| { if (self.kv_caches[i]) |*kv| kv.deinit(allocator); if (self.delta_states[i]) |*ds| ds.deinit(allocator); } }
pub fn reset(self: *ModelCache) void { for (0..Config.num_layers) |i| { if (self.kv_caches[i]) |*kv| kv.reset(); if (self.delta_states[i]) |*ds| ds.reset(); } }};
test "kv cache append and retrieve" { const allocator = std.testing.allocator; var kv = try KVCache.init(allocator, 16); defer kv.deinit(allocator);
var key: [Config.num_kv_heads * Config.head_dim]f32 = undefined; var value: [Config.num_kv_heads * Config.head_dim]f32 = undefined; @memset(&key, 1.0); @memset(&value, 2.0);
kv.append(&key, &value); try std.testing.expectEqual(@as(u32, 1), kv.len);
const retrieved_k = kv.getKey(0); try std.testing.expectEqual(@as(f32, 1.0), retrieved_k[0]);}