diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index bf0a6fd..edcd5d6 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -56,23 +56,26 @@ The parser normalizes integer metadata arrays to `u32` and string arrays to borrowable slices because those are the two array forms used by the model configuration and tokenizer. -## Supported model layout - -The implementation targets one architecture and validates it at startup: - -- hidden width: 2048 -- layers: 24 -- attention layers: 2, 6, 10, 14, 18, 21 -- query heads: 32 -- KV heads: 8 -- head width: 64 -- dense FF width: 7168 for layers 0 and 1 -- MoE FF width: 1792 -- experts: 32, top 4 selected -- short-convolution cache: 3 taps -- vocabulary: 128,000 - -All 256 required tensor names and shapes are checked before inference. +## Supported model layouts + +The loader detects and validates two LFM2.5 configurations at startup: + +| Property | LFM2.5-2.6B | LFM2.5-8B-A1B | +|---|---:|---:| +| GGUF architecture | `lfm2` | `lfm2moe` | +| Layers | 30 | 24 | +| Attention layers | 2, 5, 9, 13, 17, 21, 24, 27 | 2, 6, 10, 14, 18, 21 | +| Dense FF width | 10,752 in every layer | 7,168 in layers 0 and 1 | +| MoE FF width | — | 1,792 | +| Experts | — | 32, top 4 selected | +| RoPE base | 10,000,000 | 5,000,000 | +| Required tensors | 266 | 256 | + +Both models use a hidden width of 2,048, 32 query heads, 8 KV heads, 64-wide +heads, a three-tap short-convolution cache, and a 128,000-token vocabulary. All +required tensor names and shapes are checked before inference. Cache storage is +sized from the detected runtime layer count so agents can use either model +without a separate API. ## Forward pass @@ -84,7 +87,7 @@ For a token at absolute position `p`: - execute grouped-query attention or gated short convolution; - add the operator residual; - RMSNorm again; - - execute dense SwiGLU or sparse top-4 MoE; + - execute dense SwiGLU or, for the 8B-A1B model, sparse top-4 MoE; - add the FFN residual; 3. apply final RMSNorm; 4. project through the tied token embedding to 128,000 logits. @@ -150,9 +153,15 @@ layer because no other CPU is supported. Output rows are independent. A process-lifetime `std.Thread` pool assigns contiguous row chunks to workers from a shared counter, allowing faster cores to claim more work while preserving mmap locality within each chunk and -avoiding per-operation fork/join overhead. macOS dispatch semaphores coordinate -jobs; `dispatch_apply_f` remains a fallback if the pool cannot be initialized. -The limit defaults to the logical CPU count and can be changed before inference. +avoiding per-operation fork/join overhead. Helpers briefly spin on a generation +counter between decode projections, then park on macOS dispatch semaphores when +inference goes idle. The submitting thread also spins on completion, removing +the per-projection wake/park round trip from sustained generation. +`dispatch_apply_f` remains a fallback if the pool cannot be initialized. + +Prompt batches default to the logical CPU count. Memory-bound single-token +matvecs retain scheduler headroom on machines with more than 12 logical CPUs; +an explicit thread limit overrides both defaults. ### Accelerate @@ -228,7 +237,7 @@ The fast suite covers: - sampling; - tool schema/call parsing; - tokenizer byte mapping and scanner behavior; -- the derived 256-tensor schedule. +- the derived 256- and 266-tensor schedules. `./zigw build e2e -Doptimize=ReleaseFast` loads the full GGUF and checks that a -greedy continuation for the Denmark-capital prompt contains “Copenhagen”. +greedy continuation for the France-capital prompt contains “Paris”. diff --git a/CHANGELOG.md b/CHANGELOG.md index e4cd342..3b461b0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,9 +7,13 @@ The project follows [Semantic Versioning](https://semver.org/). Before version ## Unreleased +## 0.2.0 - 2026-08-11 + ### Added - Public CLI and Zig library for LFM2.5-8B-A1B inference. +- Native LFM2.5-2.6B GGUF inference with runtime model detection, dense + 30-layer scheduling, and model-sized caches. - Interactive chat, raw generation, token streaming, agent cloning, and tool calling. - Apple Silicon CPU kernels using ARM64 `sdot` and Accelerate. @@ -21,6 +25,9 @@ The project follows [Semantic Versioning](https://semver.org/). Before version gate/up projections to reuse activation quantization and worker dispatches. - Use native ARM64 FP16 conversion, vectorize Q8 activation quantization, and dynamically distribute decode rows across asymmetric Apple CPU cores. +- Keep decode workers hot across adjacent projections with generation-based + dispatch, park them after a short idle window, and reserve scheduler headroom + for memory-bound single-token matvecs on high-core-count machines. ## 0.1.0 - 2026-07-25 diff --git a/README.md b/README.md index 0128190..c100fad 100644 --- a/README.md +++ b/README.md @@ -11,13 +11,12 @@ Apple's Accelerate framework. On the reference M5 Max system, the current runtime reaches **220.6 prefill tokens/s** and **95.67 decode tokens/s** with the 5.2 GB Q4_K_M model described in [Reference performance](#reference-performance). -The current release is specialized for -[Liquid AI's LFM2.5-8B-A1B](https://huggingface.co/LiquidAI/LFM2.5-8B-A1B-GGUF), -but the project is built from reusable runtime pieces: a bounds-checked GGUF -loader, tokenizer, packed quantized kernels, persistent CPU worker pool, -attention and convolution caches, sampling, and an embeddable Agent API. The -goal is to grow that foundation to support additional model architectures -without giving up architecture-specific performance. +The current release supports Liquid AI's dense +[LFM2.5-2.6B](https://huggingface.co/LiquidAI/LFM2.5-2.6B-GGUF) and sparse MoE +[LFM2.5-8B-A1B](https://huggingface.co/LiquidAI/LFM2.5-8B-A1B-GGUF). Both use +the reusable runtime pieces: a bounds-checked GGUF loader, tokenizer, packed +quantized kernels, persistent CPU worker pool, attention and convolution +caches, sampling, and an embeddable Agent API. The runtime includes a command-line interface for generation and chat, plus a Zig library for streaming tokens, cloning agent state, and running tool-calling @@ -28,16 +27,16 @@ loops. | Supported today | Not supported today | |---|---| | Apple Silicon macOS | Intel Macs, Linux, and Windows | -| LFM2.5-8B-A1B GGUF files with the expected tensor layout | Other model architectures | +| LFM2.5-2.6B and LFM2.5-8B-A1B GGUF files with the expected tensor layouts | Other model architectures | | Q4_K and Q6_K quantized weights | Arbitrary GGUF quantization formats | | CPU inference through ARM64 and Accelerate | Metal, MLX, MPS, CUDA, or Vulkan | | Pinned Zig stable (0.16.0), with no third-party Zig packages | Other Zig versions | -Model support is strict by design: the loader validates the architecture, -dimensions, layer schedule, and all 256 required tensor names and shapes before -inference. Adding another model family will require an explicit model -configuration and forward path, while reusing the platform, storage, kernel, -execution, and session layers already in place. +Model support is strict by design: the loader detects the dense or MoE +architecture and validates its dimensions, layer schedule, and all 266 or 256 +required tensor names and shapes before inference. Adding another model family +requires an explicit model configuration and forward path while reusing the +platform, storage, kernel, execution, and session layers already in place. The project is independent and is not affiliated with or endorsed by Liquid AI. See [ARCHITECTURE.md](ARCHITECTURE.md) for a detailed implementation tour. @@ -55,22 +54,23 @@ matching system installation, and requires the version declared in `.zigversion` ### 2. Download the model -The official Q4_K_M model is approximately 5.2 GB: +The official LFM2.5-2.6B Q4_K_M model is approximately 1.67 GB: ```sh mkdir -p "$HOME/.cache/models" curl -fL \ - "https://huggingface.co/LiquidAI/LFM2.5-8B-A1B-GGUF/resolve/main/LFM2.5-8B-A1B-Q4_K_M.gguf" \ - -o "$HOME/.cache/models/LFM2.5-8B-A1B-Q4_K_M.gguf" + "https://huggingface.co/LiquidAI/LFM2.5-2.6B-GGUF/resolve/main/LFM2.5-2.6B-Q4_K_M.gguf" \ + -o "$HOME/.cache/models/LFM2.5-2.6B-Q4_K_M.gguf" -export LFM_WEIGHTS_FILE="$HOME/.cache/models/LFM2.5-8B-A1B-Q4_K_M.gguf" +export LFM_WEIGHTS_FILE="$HOME/.cache/models/LFM2.5-2.6B-Q4_K_M.gguf" ``` -The compatible +For the larger sparse model, download +[`LFM2.5-8B-A1B-Q4_K_M.gguf`](https://huggingface.co/LiquidAI/LFM2.5-8B-A1B-GGUF/blob/main/LFM2.5-8B-A1B-Q4_K_M.gguf) +and point `LFM_WEIGHTS_FILE` at it. The compatible [Heretic Q4_K_M variant](https://huggingface.co/FadedRedStar/LFM2.5-8B-A1B-heretic-GGUF) -has also passed the model-backed smoke test. Point `LFM_WEIGHTS_FILE` at either -model. The official weights are distributed under the -[LFM Open License 1.0](https://huggingface.co/LiquidAI/LFM2.5-8B-A1B-GGUF/blob/main/LICENSE). +has also passed the model-backed smoke test. The official weights are +distributed under the [LFM Open License 1.0](https://huggingface.co/LiquidAI/LFM2.5-2.6B-GGUF/blob/main/LICENSE). ### 3. Build and run @@ -294,7 +294,7 @@ hooks, and error handling. ## Implementation -The current LFM2.5 inference path is specialized for this model and platform: +The current LFM2.5 inference paths are specialized for these models and platform: - GGUF weights remain in a read-only private mmap and are demand-loaded by macOS. @@ -380,12 +380,28 @@ benchmark's Copenhagen sanity check. Absolute throughput varies with temperature, background load, prompt length, and growing attention context. +### Dense-model decode optimization + +On 2026-08-11, three `MAX_NEW=256` runs of the same 233-token benchmark prompt +with the official `LFM2.5-2.6B-Q4_K_M.gguf` measured the worker-dispatch change: + +| Build | Average prefill | Average decode | +|---|---:|---:| +| Before generation-based worker dispatch | 201.0 tok/s | 68.33 tok/s | +| Current default scheduling | 199.8 tok/s | 90.33 tok/s | + +This is a 32.2% decode-throughput improvement with effectively unchanged +prefill throughput. Prompt batches still use all 18 logical CPUs on the test +machine; memory-bound decode automatically uses 12 worker partitions. Passing +`--num-threads` continues to override the automatic choice. + ## Memory use -The mapped Q4_K_M weights occupy approximately 5.2 GB. Runtime memory also -includes the KV cache, convolution state, tokenizer tables, a small set of -pre-decoded F32 tensors, and transient forward-pass arenas. A typical session -uses roughly 6-8 GB, primarily depending on context length. +The mapped Q4_K_M weights occupy approximately 1.67 GB for LFM2.5-2.6B or +5.2 GB for LFM2.5-8B-A1B. Runtime memory also includes the KV cache, +convolution state, tokenizer tables, a small set of pre-decoded F32 tensors, +and transient forward-pass arenas. Total use primarily depends on the selected +model and context length. ## Origins and license diff --git a/build.zig b/build.zig index fd8ae1b..7410530 100644 --- a/build.zig +++ b/build.zig @@ -16,7 +16,7 @@ pub fn build(b: *std.Build) void { .name = "lfm", .root_module = lfm, .linkage = .static, - .version = .{ .major = 0, .minor = 1, .patch = 0 }, + .version = .{ .major = 0, .minor = 2, .patch = 0 }, }); b.installArtifact(library); diff --git a/build.zig.zon b/build.zig.zon index 9bc5205..801207c 100644 --- a/build.zig.zon +++ b/build.zig.zon @@ -1,6 +1,6 @@ .{ .name = .lfm_zig, - .version = "0.1.0", + .version = "0.2.0", .fingerprint = 0xc39470334c87d45d, .minimum_zig_version = "0.16.0", .paths = .{ diff --git a/src/agent.zig b/src/agent.zig index 17c71e7..87e35b9 100644 --- a/src/agent.zig +++ b/src/agent.zig @@ -106,7 +106,7 @@ pub const Agent = struct { return .{ .allocator = allocator, .model = model, - .cache = try Cache.init(allocator), + .cache = try model.initCache(allocator), .sampler = Sampler.recommended(allocator), }; } diff --git a/src/cache.zig b/src/cache.zig index a20f470..6f472e5 100644 --- a/src/cache.zig +++ b/src/cache.zig @@ -3,20 +3,27 @@ const config = @import("config.zig"); pub const Cache = struct { allocator: std.mem.Allocator, - k: [config.n_layers]std.ArrayList(f32), - v: [config.n_layers]std.ArrayList(f32), - conv: [config.n_layers]std.ArrayList(f32), + layer_count: usize, + k: [config.max_layers]std.ArrayList(f32), + v: [config.max_layers]std.ArrayList(f32), + conv: [config.max_layers]std.ArrayList(f32), pos: usize = 0, pub fn init(allocator: std.mem.Allocator) !Cache { + return initForLayers(allocator, config.n_layers); + } + + pub fn initForLayers(allocator: std.mem.Allocator, layer_count: usize) !Cache { + if (layer_count > config.max_layers) return error.InvalidLayerCount; var cache = Cache{ .allocator = allocator, + .layer_count = layer_count, .k = @splat(.empty), .v = @splat(.empty), .conv = @splat(.empty), }; errdefer cache.deinit(); - for (&cache.conv) |*state| { + for (cache.conv[0..layer_count]) |*state| { try state.resize(allocator, config.hidden * (config.conv_l_cache - 1)); @memset(state.items, 0); } @@ -33,13 +40,14 @@ pub const Cache = struct { pub fn clone(self: *const Cache, allocator: std.mem.Allocator) !Cache { var result = Cache{ .allocator = allocator, + .layer_count = self.layer_count, .k = @splat(.empty), .v = @splat(.empty), .conv = @splat(.empty), .pos = self.pos, }; errdefer result.deinit(); - for (0..config.n_layers) |layer| { + for (0..self.layer_count) |layer| { try result.k[layer].appendSlice(allocator, self.k[layer].items); try result.v[layer].appendSlice(allocator, self.v[layer].items); try result.conv[layer].appendSlice(allocator, self.conv[layer].items); @@ -49,13 +57,13 @@ pub const Cache = struct { pub fn kvLen(self: *const Cache) usize { var longest: usize = 0; - for (self.k) |values| longest = @max(longest, values.items.len); + for (self.k[0..self.layer_count]) |values| longest = @max(longest, values.items.len); return longest / config.kv_dim; } pub fn evictFront(self: *Cache, positions: usize) void { const drop = positions * config.kv_dim; - for (0..config.n_layers) |layer| { + for (0..self.layer_count) |layer| { if (self.k[layer].items.len < drop) continue; std.mem.copyForwards( f32, @@ -90,3 +98,18 @@ test "KV eviction slides attention layers" { cache.evictFront(1); try testing.expectEqual(@as(usize, 2), cache.kvLen()); } + +test "cache supports runtime layer counts" { + const testing = std.testing; + try testing.expectError( + error.InvalidLayerCount, + Cache.initForLayers(testing.allocator, config.max_layers + 1), + ); + var cache = try Cache.initForLayers(testing.allocator, config.lfm2_5_2_6b.n_layers); + defer cache.deinit(); + try testing.expectEqual(config.lfm2_5_2_6b.n_layers, cache.layer_count); + try testing.expectEqual( + config.hidden * (config.conv_l_cache - 1), + cache.conv[config.lfm2_5_2_6b.n_layers - 1].items.len, + ); +} diff --git a/src/config.zig b/src/config.zig index 765b76b..396904a 100644 --- a/src/config.zig +++ b/src/config.zig @@ -1,33 +1,99 @@ const std = @import("std"); const gguf = @import("gguf.zig"); -pub const arch = "lfm2moe"; +// These dimensions are shared by both currently supported LFM2.5 checkpoints. pub const hidden: usize = 2048; -pub const n_layers: usize = 24; pub const vocab: usize = 128_000; pub const n_heads: usize = 32; pub const n_kv_heads: usize = 8; pub const head_dim: usize = hidden / n_heads; pub const kv_dim: usize = n_kv_heads * head_dim; -pub const dense_ff: usize = 7168; +pub const conv_l_cache: usize = 3; +pub const rms_eps: f32 = 1e-5; + +// MoE dimensions used by LFM2.5-8B-A1B. pub const moe_ff: usize = 1792; pub const n_experts: usize = 32; pub const n_experts_used: usize = 4; -pub const n_dense_layers: usize = 2; -pub const conv_l_cache: usize = 3; -pub const rope_theta: f32 = 5_000_000; -pub const rms_eps: f32 = 1e-5; -pub const attention_layers = [_]usize{ 2, 6, 10, 14, 18, 21 }; -pub fn isAttention(layer: usize) bool { - for (attention_layers) |candidate| { - if (candidate == layer) return true; +pub const max_layers: usize = 30; + +pub const ModelKind = enum { + lfm2_5_8b_a1b, + lfm2_5_2_6b, +}; + +pub const ModelConfig = struct { + kind: ModelKind, + name: []const u8, + architecture: []const u8, + n_layers: usize, + dense_ff: usize, + n_dense_layers: usize, + rope_theta: f32, + attention_layers: []const usize, + + pub fn isAttention(self: ModelConfig, layer: usize) bool { + for (self.attention_layers) |candidate| { + if (candidate == layer) return true; + } + return false; + } + + pub fn isDenseFfn(self: ModelConfig, layer: usize) bool { + return layer < self.n_dense_layers; + } + + pub fn tensorCount(self: ModelConfig) usize { + var count: usize = 2; + for (0..self.n_layers) |layer| { + count += 2; + count += if (self.isAttention(layer)) 6 else 3; + count += if (self.isDenseFfn(layer)) 3 else 5; + } + return count; } - return false; +}; + +const attention_layers_8b = [_]usize{ 2, 6, 10, 14, 18, 21 }; +const attention_layers_2_6b = [_]usize{ 2, 5, 9, 13, 17, 21, 24, 27 }; + +pub const lfm2_5_8b_a1b = ModelConfig{ + .kind = .lfm2_5_8b_a1b, + .name = "LFM2.5-8B-A1B", + .architecture = "lfm2moe", + .n_layers = 24, + .dense_ff = 7168, + .n_dense_layers = 2, + .rope_theta = 5_000_000, + .attention_layers = &attention_layers_8b, +}; + +pub const lfm2_5_2_6b = ModelConfig{ + .kind = .lfm2_5_2_6b, + .name = "LFM2.5-2.6B", + .architecture = "lfm2", + .n_layers = 30, + .dense_ff = 10_752, + .n_dense_layers = 30, + .rope_theta = 10_000_000, + .attention_layers = &attention_layers_2_6b, +}; + +// Preserve source compatibility for callers that use the original 8B constants. +pub const arch = lfm2_5_8b_a1b.architecture; +pub const n_layers = lfm2_5_8b_a1b.n_layers; +pub const dense_ff = lfm2_5_8b_a1b.dense_ff; +pub const n_dense_layers = lfm2_5_8b_a1b.n_dense_layers; +pub const rope_theta = lfm2_5_8b_a1b.rope_theta; +pub const attention_layers = attention_layers_8b; + +pub fn isAttention(layer: usize) bool { + return lfm2_5_8b_a1b.isAttention(layer); } pub fn isDenseFfn(layer: usize) bool { - return layer < n_dense_layers; + return lfm2_5_8b_a1b.isDenseFfn(layer); } pub const ValidationError = error{ @@ -37,33 +103,57 @@ pub const ValidationError = error{ ScheduleMismatch, }; -pub fn validate(file: *const gguf.GgufFile) ValidationError!void { +pub fn detect(file: *const gguf.GgufFile) ValidationError!ModelConfig { const got_arch = file.getStr("general.architecture") orelse return error.MissingMetadata; - if (!std.mem.eql(u8, got_arch, arch)) return error.ArchitectureMismatch; - - try expectU32(file, "lfm2moe.block_count", n_layers); - try expectU32(file, "lfm2moe.embedding_length", hidden); - try expectU32(file, "lfm2moe.vocab_size", vocab); - try expectU32(file, "lfm2moe.attention.head_count", n_heads); - try expectU32(file, "lfm2moe.expert_count", n_experts); - try expectU32(file, "lfm2moe.expert_used_count", n_experts_used); - try expectU32(file, "lfm2moe.feed_forward_length", dense_ff); - try expectU32(file, "lfm2moe.expert_feed_forward_length", moe_ff); - try expectU32(file, "lfm2moe.leading_dense_block_count", n_dense_layers); - try expectU32(file, "lfm2moe.shortconv.l_cache", conv_l_cache); - - const theta = file.getF32("lfm2moe.rope.freq_base") orelse return error.MissingMetadata; - if (@abs(theta - rope_theta) > 1) return error.DimensionMismatch; - const eps = file.getF32("lfm2moe.attention.layer_norm_rms_epsilon") orelse + const model_config = if (std.mem.eql(u8, got_arch, lfm2_5_8b_a1b.architecture)) + lfm2_5_8b_a1b + else if (std.mem.eql(u8, got_arch, lfm2_5_2_6b.architecture)) + lfm2_5_2_6b + else + return error.ArchitectureMismatch; + try validateModel(file, model_config); + return model_config; +} + +pub fn validate(file: *const gguf.GgufFile) ValidationError!void { + _ = try detect(file); +} + +fn validateModel(file: *const gguf.GgufFile, model_config: ModelConfig) ValidationError!void { + var key_buffer: [96]u8 = undefined; + try expectU32(file, metadataKey(&key_buffer, model_config, "block_count"), model_config.n_layers); + try expectU32(file, metadataKey(&key_buffer, model_config, "embedding_length"), hidden); + try expectU32(file, metadataKey(&key_buffer, model_config, "vocab_size"), vocab); + try expectU32(file, metadataKey(&key_buffer, model_config, "attention.head_count"), n_heads); + try expectU32(file, metadataKey(&key_buffer, model_config, "feed_forward_length"), model_config.dense_ff); + try expectU32(file, metadataKey(&key_buffer, model_config, "shortconv.l_cache"), conv_l_cache); + + if (model_config.kind == .lfm2_5_8b_a1b) { + try expectU32(file, "lfm2moe.expert_count", n_experts); + try expectU32(file, "lfm2moe.expert_used_count", n_experts_used); + try expectU32(file, "lfm2moe.expert_feed_forward_length", moe_ff); + try expectU32(file, "lfm2moe.leading_dense_block_count", model_config.n_dense_layers); + } + + const theta = file.getF32(metadataKey(&key_buffer, model_config, "rope.freq_base")) orelse return error.MissingMetadata; + if (@abs(theta - model_config.rope_theta) > 1) return error.DimensionMismatch; + const eps = file.getF32(metadataKey( + &key_buffer, + model_config, + "attention.layer_norm_rms_epsilon", + )) orelse return error.MissingMetadata; if (@abs(eps - rms_eps) > 1e-9) return error.DimensionMismatch; - const kv = file.getU32Array("lfm2moe.attention.head_count_kv") orelse - return error.MissingMetadata; - if (kv.len != n_layers) return error.ScheduleMismatch; - for (kv, 0..) |heads, layer| { - const expected: u32 = if (isAttention(layer)) n_kv_heads else 0; - if (heads != expected) return error.ScheduleMismatch; + const kv_key = metadataKey(&key_buffer, model_config, "attention.head_count_kv"); + if (file.getU32Array(kv_key)) |kv| { + if (kv.len != model_config.n_layers) return error.ScheduleMismatch; + for (kv, 0..) |heads, layer| { + const expected: u32 = if (model_config.isAttention(layer)) n_kv_heads else 0; + if (heads != expected) return error.ScheduleMismatch; + } + } else { + try expectU32(file, kv_key, n_kv_heads); } } @@ -72,13 +162,23 @@ fn expectU32(file: *const gguf.GgufFile, key: []const u8, expected: usize) Valid if (got != expected) return error.DimensionMismatch; } -test "schedule and dimensions" { +fn metadataKey(buffer: []u8, model_config: ModelConfig, suffix: []const u8) []const u8 { + return std.fmt.bufPrint(buffer, "{s}.{s}", .{ model_config.architecture, suffix }) catch unreachable; +} + +test "supported schedules and dimensions" { const testing = std.testing; - try testing.expect(isAttention(2)); - try testing.expect(isAttention(21)); - try testing.expect(!isAttention(0)); - try testing.expect(isDenseFfn(0)); - try testing.expect(!isDenseFfn(2)); + try testing.expect(lfm2_5_8b_a1b.isAttention(2)); + try testing.expect(lfm2_5_8b_a1b.isAttention(21)); + try testing.expect(!lfm2_5_8b_a1b.isAttention(0)); + try testing.expect(lfm2_5_8b_a1b.isDenseFfn(0)); + try testing.expect(!lfm2_5_8b_a1b.isDenseFfn(2)); + try testing.expect(lfm2_5_2_6b.isAttention(5)); + try testing.expect(lfm2_5_2_6b.isAttention(27)); + try testing.expect(!lfm2_5_2_6b.isAttention(6)); + try testing.expect(lfm2_5_2_6b.isDenseFfn(29)); + try testing.expectEqual(@as(usize, 256), lfm2_5_8b_a1b.tensorCount()); + try testing.expectEqual(@as(usize, 266), lfm2_5_2_6b.tensorCount()); try testing.expectEqual(hidden, head_dim * n_heads); try testing.expectEqual(@as(usize, 512), kv_dim); } diff --git a/src/kernels/matmul.zig b/src/kernels/matmul.zig index 5581fc5..556997f 100644 --- a/src/kernels/matmul.zig +++ b/src/kernels/matmul.zig @@ -27,16 +27,20 @@ extern "c" fn dispatch_semaphore_signal(semaphore: *DispatchSemaphore) isize; const max_pool_helpers: usize = 63; const dispatch_time_forever = std.math.maxInt(u64); +const worker_spin_iterations: usize = 5_000; const CpuPool = struct { // 0 = uninitialized, 1 = initializing, 2 = ready, 3 = unavailable. state: std.atomic.Value(u8) = std.atomic.Value(u8).init(0), helper_count: usize = 0, wake: [max_pool_helpers]?*DispatchSemaphore = @splat(null), - completed: ?*DispatchSemaphore = null, + sleeping: [max_pool_helpers]std.atomic.Value(u8) = + @splat(std.atomic.Value(u8).init(0)), submit_lock: ?*DispatchSemaphore = null, context: ?*anyopaque = null, work: ?DispatchWork = null, + active_helpers: usize = 0, + generation: std.atomic.Value(usize) = std.atomic.Value(usize).init(0), remaining: std.atomic.Value(usize) = std.atomic.Value(usize).init(0), }; @@ -53,12 +57,32 @@ fn logicalCpuCount() usize { fn cpuPoolWorker(pool: *CpuPool, worker_index: usize) void { const wake = pool.wake[worker_index].?; + // The first submission can race a newly spawned helper reaching this loop. + // Starting from generation zero ensures that helper still observes it. + var generation: usize = 0; while (true) { - _ = dispatch_semaphore_wait(wake, dispatch_time_forever); - pool.work.?(pool.context, worker_index + 1); - if (pool.remaining.fetchSub(1, .acq_rel) == 1) { - _ = dispatch_semaphore_signal(pool.completed.?); + var spins: usize = 0; + while (pool.generation.load(.acquire) == generation and + spins < worker_spin_iterations) : (spins += 1) + { + std.atomic.spinLoopHint(); + } + + if (pool.generation.load(.acquire) == generation) { + pool.sleeping[worker_index].store(1, .release); + if (pool.generation.load(.acquire) == generation) { + _ = dispatch_semaphore_wait(wake, dispatch_time_forever); + } + pool.sleeping[worker_index].store(0, .release); } + + const next_generation = pool.generation.load(.acquire); + if (next_generation == generation) continue; + generation = next_generation; + if (worker_index >= pool.active_helpers) continue; + + pool.work.?(pool.context, worker_index + 1); + _ = pool.remaining.fetchSub(1, .acq_rel); } } @@ -81,9 +105,8 @@ fn initializeCpuPool() bool { } } - cpu_pool.completed = dispatch_semaphore_create(0); cpu_pool.submit_lock = dispatch_semaphore_create(1); - if (cpu_pool.completed == null or cpu_pool.submit_lock == null) { + if (cpu_pool.submit_lock == null) { cpu_pool.state.store(3, .release); return false; } @@ -130,12 +153,16 @@ fn runParallel( _ = dispatch_semaphore_wait(cpu_pool.submit_lock.?, dispatch_time_forever); cpu_pool.context = context; cpu_pool.work = work; + cpu_pool.active_helpers = workers - 1; cpu_pool.remaining.store(workers - 1, .release); + _ = cpu_pool.generation.fetchAdd(1, .release); for (0..workers - 1) |worker_index| { - _ = dispatch_semaphore_signal(cpu_pool.wake[worker_index].?); + if (cpu_pool.sleeping[worker_index].load(.acquire) != 0) { + _ = dispatch_semaphore_signal(cpu_pool.wake[worker_index].?); + } } work(context, 0); - _ = dispatch_semaphore_wait(cpu_pool.completed.?, dispatch_time_forever); + while (cpu_pool.remaining.load(.acquire) != 0) std.atomic.spinLoopHint(); _ = dispatch_semaphore_signal(cpu_pool.submit_lock.?); } @@ -161,6 +188,18 @@ fn workerCount(rows: usize) usize { return @max(1, @min(limit, rows)); } +fn decodeWorkerCount(rows: usize) usize { + const configured = thread_limit.load(.monotonic); + if (configured != 0) return @max(1, @min(configured, rows)); + + const cpus = logicalCpuCount(); + // Quantized single-token matvecs saturate memory bandwidth before every + // logical CPU is busy. Leave scheduler headroom on larger machines while + // retaining all cores on smaller systems. + const limit = if (cpus > 12) @max(12, cpus - cpus / 3) else cpus; + return @max(1, @min(limit, rows)); +} + pub fn dot(a: []const f32, b: []const f32) f32 { return accelerate.dot(a, b); } @@ -696,7 +735,7 @@ pub fn matvec( .rows = n_out, .q8 = &q8, .out = y, - .workers = if (n_out < parallel_min_rows) 1 else workerCount(n_out), + .workers = if (n_out < parallel_min_rows) 1 else decodeWorkerCount(n_out), }; runParallel(&context, context.workers, matvecWorker); return; @@ -926,7 +965,7 @@ pub fn matvecFused( .row_bytes = row_bytes, .total_rows = total_rows, .out = out, - .workers = if (total_rows < parallel_min_rows) 1 else workerCount(total_rows), + .workers = if (total_rows < parallel_min_rows) 1 else decodeWorkerCount(total_rows), }; runParallel(&context, context.workers, fusedWorker); } @@ -1115,10 +1154,13 @@ test "persistent quantized row workers match the serial result" { } var x: [dequant.qk_k]f32 = @splat(1); var out: [rows]f32 = @splat(0); - setThreadLimit(4); defer setThreadLimit(0); - try matvec(testing.allocator, .q4_k, weights, dequant.qk_k, rows, &x, &out); - for (out) |value| try testing.expectApproxEqAbs(@as(f32, 256), value, 1e-4); + for ([_]usize{ 4, 2, 4, 1, 3 }) |workers| { + @memset(&out, 0); + setThreadLimit(workers); + try matvec(testing.allocator, .q4_k, weights, dequant.qk_k, rows, &x, &out); + for (out) |value| try testing.expectApproxEqAbs(@as(f32, 256), value, 1e-4); + } } test "tiled quantized batches match per-token matrix vectors" { diff --git a/src/main.zig b/src/main.zig index 4927204..fd22006 100644 --- a/src/main.zig +++ b/src/main.zig @@ -62,7 +62,7 @@ pub fn main(init: std.process.Init) !void { fn usage(writer: *std.Io.Writer) !void { try writer.writeAll( - \\lfm.zig — Apple Silicon CPU inference for LFM2.5-8B-A1B + \\lfm.zig — Apple Silicon CPU inference for LFM2.5 \\ \\usage: \\ lfm generate [opts] ... diff --git a/src/model.zig b/src/model.zig index 2039008..3458eb7 100644 --- a/src/model.zig +++ b/src/model.zig @@ -1,4 +1,4 @@ -//! LFM2.5-8B-A1B model loading and cached CPU forward pass. +//! LFM2.5 model loading and cached CPU forward pass. const std = @import("std"); const accelerate = @import("accelerate.zig"); @@ -20,9 +20,11 @@ const Tokenizer = @import("tokenizer.zig").Tokenizer; pub const Error = error{ MissingTensor, + InvalidTensorCount, InvalidTensorShape, UnsupportedTensorType, TokenOutOfRange, + CacheMismatch, } || gguf_module.Error || config.ValidationError || @import("tokenizer.zig").Error || matmul_kernel.Error || dequant.Error || std.mem.Allocator.Error; @@ -32,12 +34,13 @@ pub const Model = struct { by_name: std.StringHashMap(usize), f32_cache: std.StringHashMap([]f32), tokenizer_value: Tokenizer, + model_config: config.ModelConfig, pub fn load(allocator: std.mem.Allocator, path: []const u8) Error!Model { var file = try GgufFile.open(allocator, path); var file_is_local = true; errdefer if (file_is_local) file.deinit(); - try config.validate(&file); + const model_config = try config.detect(&file); var by_name = std.StringHashMap(usize).init(allocator); var index_is_local = true; errdefer if (index_is_local) by_name.deinit(); @@ -54,6 +57,7 @@ pub const Model = struct { .by_name = by_name, .f32_cache = std.StringHashMap([]f32).init(allocator), .tokenizer_value = tokenizer_instance, + .model_config = model_config, }; file_is_local = false; index_is_local = false; @@ -106,6 +110,14 @@ pub const Model = struct { return &self.tokenizer_value; } + pub fn configuration(self: *const Model) config.ModelConfig { + return self.model_config; + } + + pub fn initCache(self: *const Model, allocator: std.mem.Allocator) !Cache { + return Cache.initForLayers(allocator, self.model_config.n_layers); + } + /// Returned logits are allocated with the model allocator; the caller owns them. pub fn forwardStep( self: *const Model, @@ -113,6 +125,7 @@ pub const Model = struct { token: u32, cache: *Cache, ) Error![]f32 { + try self.validateCache(cache); var arena = std.heap.ArenaAllocator.init(self.allocator); defer arena.deinit(); const temporary = arena.allocator(); @@ -122,6 +135,7 @@ pub const Model = struct { } pub fn runLayers(self: *const Model, token: u32, cache: *Cache) Error!void { + try self.validateCache(cache); var arena = std.heap.ArenaAllocator.init(self.allocator); defer arena.deinit(); _ = try self.runLayersArena(arena.allocator(), token, cache); @@ -130,6 +144,7 @@ pub const Model = struct { /// Runs a token-major prompt batch through all layers while preserving the /// same causal KV and convolution-cache updates as the scalar decode path. pub fn runLayersBatch(self: *const Model, tokens: []const u32, cache: *Cache) Error!void { + try self.validateCache(cache); if (tokens.len == 0) return; if (tokens.len == 1) return self.runLayers(tokens[0], cache); @@ -141,7 +156,7 @@ pub const Model = struct { try self.embedToken(token, hidden[index * config.hidden ..][0..config.hidden]); } - for (0..config.n_layers) |layer| { + for (0..self.model_config.n_layers) |layer| { var layer_arena = std.heap.ArenaAllocator.init(self.allocator); defer layer_arena.deinit(); const allocator = layer_arena.allocator(); @@ -153,7 +168,7 @@ pub const Model = struct { tensorName(&name_buffer, layer, "attn_norm.weight"), token_count, ); - const operator = if (config.isAttention(layer)) + const operator = if (self.model_config.isAttention(layer)) try self.attentionBatch( allocator, layer, @@ -172,7 +187,7 @@ pub const Model = struct { tensorName(&name_buffer, layer, "ffn_norm.weight"), token_count, ); - const ffn = if (config.isDenseFfn(layer)) + const ffn = if (self.model_config.isDenseFfn(layer)) try self.denseFfnBatch(allocator, layer, ffn_normalized, token_count) else try self.moeFfnBatch(allocator, layer, ffn_normalized, token_count); @@ -317,12 +332,14 @@ pub const Model = struct { config.n_heads, q_gain, position_zero + token, + self.model_config.rope_theta, ); normRopeHeads( k[token * config.kv_dim ..][0..config.kv_dim], config.n_kv_heads, k_gain, position_zero + token, + self.model_config.rope_theta, ); } @@ -371,14 +388,15 @@ pub const Model = struct { token_count: usize, ) Error![]f32 { var name_buffer: [64]u8 = undefined; - const gate = try allocator.alloc(f32, token_count * config.dense_ff); - const up = try allocator.alloc(f32, token_count * config.dense_ff); + const dense_ff = self.model_config.dense_ff; + const gate = try allocator.alloc(f32, token_count * dense_ff); + const up = try allocator.alloc(f32, token_count * dense_ff); try self.matmulNamedBatch( allocator, tensorName(&name_buffer, layer, "ffn_gate.weight"), x, config.hidden, - config.dense_ff, + dense_ff, token_count, gate, ); @@ -387,16 +405,16 @@ pub const Model = struct { tensorName(&name_buffer, layer, "ffn_up.weight"), x, config.hidden, - config.dense_ff, + dense_ff, token_count, up, ); - const activated = try allocator.alloc(f32, token_count * config.dense_ff); + const activated = try allocator.alloc(f32, token_count * dense_ff); for (0..token_count) |token| { activation.swiglu( - gate[token * config.dense_ff ..][0..config.dense_ff], - up[token * config.dense_ff ..][0..config.dense_ff], - activated[token * config.dense_ff ..][0..config.dense_ff], + gate[token * dense_ff ..][0..dense_ff], + up[token * dense_ff ..][0..dense_ff], + activated[token * dense_ff ..][0..dense_ff], ); } const out = try allocator.alloc(f32, token_count * config.hidden); @@ -404,7 +422,7 @@ pub const Model = struct { allocator, tensorName(&name_buffer, layer, "ffn_down.weight"), activated, - config.dense_ff, + dense_ff, config.hidden, token_count, out, @@ -581,11 +599,11 @@ pub const Model = struct { const position = cache.pos; const hidden = try allocator.alloc(f32, config.hidden); try self.embedToken(token, hidden); - for (0..config.n_layers) |layer| { + for (0..self.model_config.n_layers) |layer| { var name_buffer: [64]u8 = undefined; const attn_norm_name = tensorName(&name_buffer, layer, "attn_norm.weight"); const normalized = try self.normOne(allocator, hidden, attn_norm_name); - const operator = if (config.isAttention(layer)) + const operator = if (self.model_config.isAttention(layer)) try self.attentionStep(allocator, layer, normalized, position, cache) else try self.convStep(allocator, layer, normalized, cache); @@ -593,7 +611,7 @@ pub const Model = struct { const ffn_norm_name = tensorName(&name_buffer, layer, "ffn_norm.weight"); const ffn_normalized = try self.normOne(allocator, hidden, ffn_norm_name); - const ffn = if (config.isDenseFfn(layer)) + const ffn = if (self.model_config.isDenseFfn(layer)) try self.denseFfn(allocator, layer, ffn_normalized) else try self.moeFfn(allocator, layer, ffn_normalized); @@ -744,12 +762,14 @@ pub const Model = struct { config.n_heads, self.floatTensor(tensorName(&name_buffer, layer, "attn_q_norm.weight")), position, + self.model_config.rope_theta, ); normRopeHeads( k, config.n_kv_heads, self.floatTensor(tensorName(&name_buffer, layer, "attn_k_norm.weight")), position, + self.model_config.rope_theta, ); try cache.k[layer].appendSlice(cache.allocator, k); try cache.v[layer].appendSlice(cache.allocator, v); @@ -791,15 +811,16 @@ pub const Model = struct { return error.MissingTensor; var input_q8 = try matmul_kernel.quantizeQ8(allocator, x); defer input_q8.deinit(); - const gate_up = try allocator.alloc(f32, 2 * config.dense_ff); - const gate = gate_up[0..config.dense_ff]; - const up = gate_up[config.dense_ff..]; + const dense_ff = self.model_config.dense_ff; + const gate_up = try allocator.alloc(f32, 2 * dense_ff); + const gate = gate_up[0..dense_ff]; + const up = gate_up[dense_ff..]; const jobs = [_]matmul_kernel.FusedJob{ .{ .dtype = gate_tensor.ggml_type, .weights = self.data(gate_tensor), .n_in = config.hidden, - .n_out = config.dense_ff, + .n_out = dense_ff, .x = x, .q8 = if (usesQ8(gate_tensor.ggml_type)) &input_q8 else null, }, @@ -807,20 +828,20 @@ pub const Model = struct { .dtype = up_tensor.ggml_type, .weights = self.data(up_tensor), .n_in = config.hidden, - .n_out = config.dense_ff, + .n_out = dense_ff, .x = x, .q8 = if (usesQ8(up_tensor.ggml_type)) &input_q8 else null, }, }; try matmul_kernel.matvecFused(allocator, &jobs, gate_up); - const activated = try allocator.alloc(f32, config.dense_ff); + const activated = try allocator.alloc(f32, dense_ff); activation.swiglu(gate, up, activated); const out = try allocator.alloc(f32, config.hidden); try self.matvecNamed( allocator, tensorName(&name_buffer, layer, "ffn_down.weight"), activated, - config.dense_ff, + dense_ff, config.hidden, out, ); @@ -1010,10 +1031,17 @@ pub const Model = struct { return out; } + fn validateCache(self: *const Model, cache: *const Cache) Error!void { + if (cache.layer_count != self.model_config.n_layers) return error.CacheMismatch; + } + fn checkTensors(self: *const Model) Error!void { + if (self.gguf.tensors.len != self.model_config.tensorCount()) { + return error.InvalidTensorCount; + } try self.checkShape("token_embd.weight", &.{ config.hidden, config.vocab }); try self.checkShape("token_embd_norm.weight", &.{config.hidden}); - for (0..config.n_layers) |layer| { + for (0..self.model_config.n_layers) |layer| { var name_buffer: [64]u8 = undefined; try self.checkShape( tensorName(&name_buffer, layer, "attn_norm.weight"), @@ -1023,7 +1051,7 @@ pub const Model = struct { tensorName(&name_buffer, layer, "ffn_norm.weight"), &.{config.hidden}, ); - if (config.isAttention(layer)) { + if (self.model_config.isAttention(layer)) { try self.checkShape( tensorName(&name_buffer, layer, "attn_q.weight"), &.{ config.hidden, config.hidden }, @@ -1062,18 +1090,18 @@ pub const Model = struct { &.{ config.hidden, config.hidden }, ); } - if (config.isDenseFfn(layer)) { + if (self.model_config.isDenseFfn(layer)) { try self.checkShape( tensorName(&name_buffer, layer, "ffn_gate.weight"), - &.{ config.hidden, config.dense_ff }, + &.{ config.hidden, self.model_config.dense_ff }, ); try self.checkShape( tensorName(&name_buffer, layer, "ffn_up.weight"), - &.{ config.hidden, config.dense_ff }, + &.{ config.hidden, self.model_config.dense_ff }, ); try self.checkShape( tensorName(&name_buffer, layer, "ffn_down.weight"), - &.{ config.dense_ff, config.hidden }, + &.{ self.model_config.dense_ff, config.hidden }, ); } else { try self.checkShape( @@ -1113,13 +1141,19 @@ fn tensorName(buffer: []u8, layer: usize, suffix: []const u8) []const u8 { return std.fmt.bufPrint(buffer, "blk.{d}.{s}", .{ layer, suffix }) catch unreachable; } -fn normRopeHeads(buffer: []f32, heads: usize, gain: []const f32, position: usize) void { +fn normRopeHeads( + buffer: []f32, + heads: usize, + gain: []const f32, + position: usize, + rope_theta: f32, +) void { var normalized: [config.head_dim]f32 = undefined; for (0..heads) |head_index| { const head = buffer[head_index * config.head_dim ..][0..config.head_dim]; rmsnorm(head, gain, config.rms_eps, &normalized); @memcpy(head, &normalized); - ropeNeox(head, position, config.rope_theta); + ropeNeox(head, position, rope_theta); } } @@ -1132,12 +1166,7 @@ fn usesQ8(dtype: GgmlType) bool { return dtype == .q4_k or dtype == .q6_k; } -test "tensor schedule has the expected count" { - var count: usize = 2; - for (0..config.n_layers) |layer| { - count += 2; - count += if (config.isAttention(layer)) 6 else 3; - count += if (config.isDenseFfn(layer)) 3 else 5; - } - try std.testing.expectEqual(@as(usize, 256), count); +test "tensor schedules have the expected counts" { + try std.testing.expectEqual(@as(usize, 256), config.lfm2_5_8b_a1b.tensorCount()); + try std.testing.expectEqual(@as(usize, 266), config.lfm2_5_2_6b.tensorCount()); }