From 715e2323561c72c99953cef4450384f02a943093 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sun, 1 Nov 2026 11:09:40 +0100 Subject: [PATCH 01/14] lfm3: the backends share the rope frequencies and the layer schedule Rotary.inverseFrequencies() builds RoPE's inverse frequencies as the CUDA backend does: theta to the 2i/128 in float64, rounded to float32, one over it in float32, zero past the rotated pairs. A test pins both of LFM3's settings (theta 1e6 partial 0.5, theta 1e4 partial 1) to torch 2.14's float32 rows bit for bit; a float32 pow is off by 1 to 3 ulps in 16 of 32 and 28 of 64 entries. count() gives the layers of a kind in the admitted schedule, LayerType.window() the keys a layer's attention reads back (the sliding window, or 0 for all of them), and ring() the keys a sliding ring holds when a pass writes its rows before the first of them reads its window. A test checks 13 conv, 6 full and 5 sliding layers and both sizes. --- zig/src/families/lfm3/config.zig | 46 +++++++++++++++++++++++++-- zig/src/families/lfm3/config_test.zig | 35 ++++++++++++++++++++ 2 files changed, 79 insertions(+), 2 deletions(-) diff --git a/zig/src/families/lfm3/config.zig b/zig/src/families/lfm3/config.zig index 84ac2ae5b5..3cd647ec7e 100644 --- a/zig/src/families/lfm3/config.zig +++ b/zig/src/families/lfm3/config.zig @@ -12,8 +12,38 @@ pub const head_dim = 128; pub const sliding_window = 256; pub const eps: f32 = 1e-5; pub const Kind = enum { text, vision_language }; -pub const LayerType = enum { conv, full_attention, sliding_attention }; -pub const Rotary = struct { theta: f32, partial: f32 }; +pub const LayerType = enum { + conv, + full_attention, + sliding_attention, + + /// Keys before a row its attention reads: the sliding window, or 0 for every one (a conv layer reads none). + pub fn window(t: LayerType) usize { + return if (t == .sliding_attention) sliding_window else 0; + } +}; + +pub const Rotary = struct { + theta: f32, + partial: f32, + + /// RoPE's inverse frequencies as torch's float32 rows: 1 / theta^(2i / 128) rotated, zero past them (cos 1, sin 0). + pub fn inverseFrequencies(rot: Rotary) [head_dim / 2]f32 { + var out: [head_dim / 2]f32 = undefined; + const angles: usize = @intFromFloat(rot.partial * @as(f32, head_dim) / 2); + for (&out, 0..) |*v, i| { + if (i >= angles) { + v.* = 0; + continue; + } + const e: f32 = @as(f32, @floatFromInt(2 * i)) / @as(f32, head_dim); + const pw: f32 = @floatCast(std.math.pow(f64, rot.theta, e)); + v.* = 1.0 / pw; + } + return out; + } +}; + pub const schedule: [layers]LayerType = .{ .conv, .conv, .full_attention, .conv, .sliding_attention, .conv, .full_attention, .conv, .sliding_attention, .conv, .full_attention, .conv, @@ -21,6 +51,18 @@ pub const schedule: [layers]LayerType = .{ .full_attention, .conv, .sliding_attention, .conv, .full_attention, .conv, }; +/// Layers of `kind` in the admitted schedule. +pub fn count(kind: LayerType) usize { + var n: usize = 0; + for (schedule) |t| n += @intFromBool(t == kind); + return n; +} + +/// Keys a sliding layer's ring holds when `rows` rows are written before the first of them reads its window back. +pub fn ring(rows: usize) usize { + return sliding_window + rows; +} + pub const Width = struct { bits: u8, group: u16 }; pub const Quant = struct { global: Width, diff --git a/zig/src/families/lfm3/config_test.zig b/zig/src/families/lfm3/config_test.zig index f1a1df00ce..136223dbd5 100644 --- a/zig/src/families/lfm3/config_test.zig +++ b/zig/src/families/lfm3/config_test.zig @@ -19,6 +19,41 @@ test "text and VL preserve their distinct context and actual mixed embedding wid try std.testing.expectEqual(@as(f32, 1), t.sliding_rotary.partial); } +test "RoPE: both settings' inverse frequencies bit for bit torch 2.14's float32 rows, zero past the rotated pairs" { + const full = [32]u32{ + 0x3f800000, 0x3f4e4bad, 0x3f263de0, 0x3f05f6ef, 0x3ed7e89b, 0x3eadfcff, 0x3e8c3504, 0x3e61f835, + 0x3e361887, 0x3e12bd91, 0x3dec7fd6, 0x3dbe94c6, 0x3d99940d, 0x3d778513, 0x3d47763f, 0x3d20bc1d, + 0x3d0186e3, 0x3cd0c1a8, 0x3ca8398b, 0x3c879008, 0x3c5a7bf2, 0x3c301052, 0x3c0de12d, 0x3be4aa46, + 0x3bb8449c, 0x3b947dae, 0x3b6f520e, 0x3b40dac5, 0x3b1b690d, 0x3afa78f0, 0x3ac9d75c, 0x3aa2a6f7, + }; + const sliding = [64]u32{ + 0x3f800000, 0x3f5dafd7, 0x3f3ff911, 0x3f263de0, 0x3f0ff59a, 0x3ef953cf, 0x3ed7e89b, 0x3ebaf81b, + 0x3ea1e89b, 0x3e8c3504, 0x3e72d423, 0x3e5247ed, 0x3e361887, 0x3e1db040, 0x3e088d77, 0x3dec7fd6, + 0x3dcccccd, 0x3db15978, 0x3d99940d, 0x3d84fe4d, 0x3d6655c2, 0x3d47763f, 0x3d2cba15, 0x3d159348, + 0x3d0186e3, 0x3ce054d2, 0x3cc2434f, 0x3ca8398b, 0x3c91ad39, 0x3c7c4d33, 0x3c5a7bf2, 0x3c3d3311, + 0x3c23d70a, 0x3c0de12d, 0x3bf5b9b0, 0x3bd4ca15, 0x3bb8449c, 0x3b9f91cc, 0x3b8a2e77, 0x3b6f520e, + 0x3b4f3e38, 0x3b33770f, 0x3b1b690d, 0x3b06946f, 0x3ae91528, 0x3ac9d75c, 0x3aaec98e, 0x3a975c0e, + 0x3a83126f, 0x3a6301e2, 0x3a44948c, 0x3a2a3b44, 0x3a136a16, 0x39ff4fac, 0x39dd1725, 0x39bf74d7, + 0x39a5cb60, 0x398f9272, 0x3978a815, 0x395753e4, 0x393a7753, 0x39217916, 0x390bd472, 0x38f22ce2, + }; + const f = (c.Rotary{ .theta = 1e6, .partial = 0.5 }).inverseFrequencies(); + const s = (c.Rotary{ .theta = 1e4, .partial = 1 }).inverseFrequencies(); + for (full, f[0..32]) |want, v| try std.testing.expectEqual(want, @as(u32, @bitCast(v))); + for (f[32..]) |v| try std.testing.expectEqual(@as(f32, 0), v); + for (sliding, s) |want, v| try std.testing.expectEqual(want, @as(u32, @bitCast(v))); +} + +test "the schedule: 13 conv, 6 full and 5 sliding layers; only a sliding layer has a window, its ring the window and the rows" { + try std.testing.expectEqual(@as(usize, 13), c.count(.conv)); + try std.testing.expectEqual(@as(usize, 6), c.count(.full_attention)); + try std.testing.expectEqual(@as(usize, 5), c.count(.sliding_attention)); + try std.testing.expectEqual(@as(usize, 256), c.LayerType.sliding_attention.window()); + try std.testing.expectEqual(@as(usize, 0), c.LayerType.full_attention.window()); + try std.testing.expectEqual(@as(usize, 0), c.LayerType.conv.window()); + try std.testing.expectEqual(@as(usize, 256 + 16), c.ring(16)); + try std.testing.expectEqual(@as(usize, 1280), c.ring(1024)); +} + test "6 and 8 bit affine variants use the same config parser" { inline for (.{ "6", "8" }) |bits| { const replacement = "\"bits\": " ++ bits; From a398e74df8577194d0769f77a9eec8a8b667c5a0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 30 Oct 2026 15:30:37 +0100 Subject: [PATCH 02/14] xpu: a server without DPAS leaves the Qwen3.8-27B runner and Gemma 4 E4B out, and nem_prow gets a 4 x 4 tile image for Gen12LP The Qwen3.8-27B runner and Gemma 4 E4B need DPAS, 2D block loads and nGEN, so `zig build xpu-native -Dxpu-device=tgllp` could not build. With -Dxpu-dpas=false the server is built without both, lists only the families it built, and refuses qwen3_5 and gemma4 naming DPAS; a bmg build is unchanged and still carries them. nem_prow.cl's PR and PO become overridable; the new image nem_prow_t4 (-DPR=4 -DPO=4) spills about 25 registers on Gen12LP where the 8 x 8 pdense spills about 210. The 14 existing bmg images stay byte-equal to main. --- zig/build/xpu.zig | 25 ++++++++++++++++++++----- zig/kernels/xpu/nem_prow.cl | 6 +++++- zig/src/native/xpu.zig | 18 +++++++++++++----- zig/src/native/xpu_nem.zig | 8 ++++---- zig/src/xpu/kernels.zig | 2 ++ 5 files changed, 44 insertions(+), 15 deletions(-) diff --git a/zig/build/xpu.zig b/zig/build/xpu.zig index e39bea9b48..8211729a02 100644 --- a/zig/build/xpu.zig +++ b/zig/build/xpu.zig @@ -139,11 +139,20 @@ pub fn steps(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.buil const lanes = b.createModule(.{ .root_source_file = b.path("zig/src/core/lanes/lanes.zig"), .target = target, .optimize = optimize, .link_libc = true }); const nem = nemotronModules(b, target, optimize, ocloc, device, dpas, ze_mod, lanes); const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} }); - const gemma4 = b.createModule(.{ .root_source_file = b.path("zig/src/families/gemma4/xpu_engine.zig"), .target = target, .optimize = optimize, .link_libc = true }); - gemma4.addImport("xpu_gemma4_kernels", embedded_g4); - gemma4.addImport("core", core); - gemma4.addImport("ze", ze_mod); - const engines = b.createModule(.{ .root_source_file = b.path("zig/src/native/xpu.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "engine_api", .module = api }, .{ .name = "xpu_qwen27", .module = qwen27 }, .{ .name = "lanes", .module = lanes }, .{ .name = "xpu", .module = nem.xpu }, .{ .name = "nemotron_xpu", .module = nem.nemotron }, .{ .name = "gemma4_xpu", .module = gemma4 } } }); + const engines = b.createModule(.{ .root_source_file = b.path("zig/src/native/xpu.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "engine_api", .module = api }, .{ .name = "lanes", .module = lanes }, .{ .name = "xpu", .module = nem.xpu }, .{ .name = "nemotron_xpu", .module = nem.nemotron } } }); + // the Qwen3.8-27B runner and Gemma 4 E4B (DPAS, nGEN) only with DPAS: a Gen12LP server serves Nemotron + const engine_options = b.addOptions(); + engine_options.addOption(bool, "qwen27", dpas); + engine_options.addOption(bool, "gemma4", dpas); + engines.addOptions("xpu_engine_options", engine_options); + if (dpas) { + engines.addImport("xpu_qwen27", qwen27); + const gemma4 = b.createModule(.{ .root_source_file = b.path("zig/src/families/gemma4/xpu_engine.zig"), .target = target, .optimize = optimize, .link_libc = true }); + gemma4.addImport("xpu_gemma4_kernels", embedded_g4); + gemma4.addImport("core", core); + gemma4.addImport("ze", ze_mod); + engines.addImport("gemma4_xpu", gemma4); + } const tokenizer = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true }); const template = b.createModule(.{ .root_source_file = b.path("zig/src/core/template/template.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true }); const checkpoint_cli = b.createModule(.{ .root_source_file = b.path("zig/src/cli/cli.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true, .imports = &.{ .{ .name = "native_engines", .module = engines }, exlr8_build.cliImport(b, target, .ReleaseSafe, core), .{ .name = "core", .module = core } } }); @@ -181,6 +190,7 @@ const kernels = [_]Kernel{ .{ .name = "nem_rows_r4", .src = "nem_rows", .options = cl3 ++ " -DMAXR=4" }, .{ .name = "nem_rows_r8", .src = "nem_rows", .options = cl3 ++ " -DMAXR=8" }, .{ .name = "nem_prow", .options = cl3 }, + .{ .name = "nem_prow_t4", .src = "nem_prow", .options = cl3 ++ " -DPR=4 -DPO=4" }, .{ .name = "nem_attn_dec", .options = cl3 }, }; @@ -289,6 +299,11 @@ pub fn hostTests(b: *std.Build, step: *std.Build.Step) void { gemma_native.addImport("ze", ze); gemma_native.addImport("xpu_gemma4_kernels", empty); sources.addImport("gemma4_xpu", gemma_native); + // the host test checks every family's admission, so it builds with both DPAS engines + const source_options = b.addOptions(); + source_options.addOption(bool, "qwen27", true); + source_options.addOption(bool, "gemma4", true); + sources.addOptions("xpu_engine_options", source_options); const source_tests = b.addRunArtifact(b.addTest(.{ .root_module = sources, .filters = &.{"draft sources:"} })); b.step("test-slide-xpu-host", "Unsupported XPU learning refuses before device or checkpoint access").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = sources, .filters = &.{"slide admission:"} })).step); b.step("test-draft-xpu-host", "XPU draft source metadata without devices or kernel compilation").dependOn(&source_tests.step); diff --git a/zig/kernels/xpu/nem_prow.cl b/zig/kernels/xpu/nem_prow.cl index 7cc36c1a92..9337e1f02a 100644 --- a/zig/kernels/xpu/nem_prow.cl +++ b/zig/kernels/xpu/nem_prow.cl @@ -1,7 +1,11 @@ // Nemotron-H prompt rows (any count) with the decode step's arithmetic: each output sums as the one-token path does. #define SG 16 -#define PR 8 // prompt rows a sub-group holds +#ifndef PR +#define PR 8 // prompt rows a sub-group holds (nem_prow_t4: 4, fewer registers for Gen12LP) +#endif +#ifndef PO #define PO 8 // output rows a sub-group holds: one load and one sum of a word's inputs serve all of them +#endif #define PSG (16 / PO) // sub-groups a work-group: 16 output rows inline float bf(ushort v) { return as_float((uint)v << 16); } inline ushort to_bf(float f) { diff --git a/zig/src/native/xpu.zig b/zig/src/native/xpu.zig index 935793b2a8..6ec6e1ebf4 100644 --- a/zig/src/native/xpu.zig +++ b/zig/src/native/xpu.zig @@ -1,7 +1,9 @@ //! XPU engines (Intel Xe2 over Level Zero): Qwen3.8-27B behind the lane host; a lone stream uses the backend's driver. const std = @import("std"); const api = @import("engine_api"); -const xq = @import("xpu_qwen27"); +/// The Qwen3.8-27B runner needs DPAS and nGEN: a build without DPAS (-Dxpu-dpas=false, Gen12LP) leaves it out. +const has_qwen27 = @import("xpu_engine_options").qwen27; +const xq = if (has_qwen27) @import("xpu_qwen27") else struct {}; const Drafter = xq.drafter.Drafter; const lanes = @import("lanes"); const sampling = lanes.sampling; @@ -11,7 +13,9 @@ const admit = @import("xpu_admit.zig"); const Allocator = std.mem.Allocator; pub const backends: []const []const u8 = &.{"xpu"}; -pub const families: []const api.Family = &.{ .{ .model_type = "qwen3_5", .formats = &.{ "mlx-q4g64", "gptq-int4-g128" } }, .{ .model_type = "nemotron_h", .formats = &.{"mlx-q4g64"} }, .{ .model_type = "gemma4", .formats = &.{"mlx-q4g64"} } }; +const qwen27_family: []const api.Family = if (has_qwen27) &.{.{ .model_type = "qwen3_5", .formats = &.{ "mlx-q4g64", "gptq-int4-g128" } }} else &.{}; +const gemma4_family: []const api.Family = if (nem.has_gemma4) &.{.{ .model_type = "gemma4", .formats = &.{"mlx-q4g64"} }} else &.{}; +pub const families: []const api.Family = qwen27_family ++ &[_]api.Family{.{ .model_type = "nemotron_h", .formats = &.{"mlx-q4g64"} }} ++ gemma4_family; pub fn chip(_: Allocator) ?[]const u8 { return "Intel Xe2 (Level Zero)"; @@ -175,10 +179,14 @@ pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]c return null; } if (nem.serves(o.model_type)) return nem.open(a, gpa, io, o, problem); - if (!std.mem.eql(u8, o.model_type, "qwen3_5")) { + if (comptime has_qwen27) { + if (std.mem.eql(u8, o.model_type, "qwen3_5")) return openQwen27(a, gpa, io, o, problem); problem.* = try std.fmt.allocPrint(a, "the XPU engine serves Qwen3.8-27B (qwen3_5, MLX affine 4/64 or symmetric GPTQ 4/128) and Nemotron-H (nemotron_h, MLX 4/64), not {s}", .{o.model_type}); - return null; - } + } else problem.* = try std.fmt.allocPrint(a, "this XPU build (no DPAS) serves Nemotron-H (nemotron_h, MLX 4/64), not {s}: Qwen3.8-27B and Gemma 4 E4B need DPAS (-Dxpu-dpas=true)", .{o.model_type}); + return null; +} + +fn openQwen27(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { const context = xq.limits.context(o.context orelse 32768) catch { problem.* = try std.fmt.allocPrint(a, "the XPU engine holds 1 to {d} tokens a stream, not {d}", .{ xq.limits.MAX_CONTEXT, o.context.? }); return null; diff --git a/zig/src/native/xpu_nem.zig b/zig/src/native/xpu_nem.zig index c14f29fcb5..61344fbb54 100644 --- a/zig/src/native/xpu_nem.zig +++ b/zig/src/native/xpu_nem.zig @@ -6,11 +6,11 @@ const lanes = @import("lanes"); const lane = @import("xpu_lane.zig"); const Allocator = std.mem.Allocator; +/// Gemma 4 E4B's GEMMs are nGEN DPAS kernels: a build without DPAS (-Dxpu-dpas=false, Gen12LP) leaves it out. +pub const has_gemma4 = @import("xpu_engine_options").gemma4; + /// The families this host serves, each a file with its table row, loader and prompt window. -const specs = .{ - @import("xpu_nemotron.zig"), - @import("xpu_gemma4.zig"), -}; +const specs = if (has_gemma4) .{ @import("xpu_nemotron.zig"), @import("xpu_gemma4.zig") } else .{@import("xpu_nemotron.zig")}; /// Prompt plus reply tokens when --context names none (and the model's window is larger). const default_context: i64 = 32768; diff --git a/zig/src/xpu/kernels.zig b/zig/src/xpu/kernels.zig index 14b3df7017..04fb4d0b34 100644 --- a/zig/src/xpu/kernels.zig +++ b/zig/src/xpu/kernels.zig @@ -24,4 +24,6 @@ pub const nem_rows: []const u8 = if (available) &Blob("spv_nem_rows").bytes else /// nem_rows built for windows of at most 1, 2, 4 and 8 rows (fewer registers a lane). pub const nem_rows_r: [4][]const u8 = if (available) .{ &Blob("spv_nem_rows_r1").bytes, &Blob("spv_nem_rows_r2").bytes, &Blob("spv_nem_rows_r4").bytes, &Blob("spv_nem_rows_r8").bytes } else .{ &.{}, &.{}, &.{}, &.{} }; pub const nem_prow: []const u8 = if (available) &Blob("spv_nem_prow").bytes else &.{}; +/// nem_prow with 4 x 4 tiles (prompt rows x outputs a sub-group): no register spills to speak of on Gen12LP. +pub const nem_prow_t4: []const u8 = if (available) &Blob("spv_nem_prow_t4").bytes else &.{}; pub const nem_attn_dec: []const u8 = if (available) &Blob("spv_nem_attn_dec").bytes else &.{}; From e0f045f0029fd0a3c8f4e24b7a690cbe7ea41b30 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 30 Oct 2026 17:52:15 +0100 Subject: [PATCH 03/14] xpu: the Gen12LP prompt image takes 8 x 4 tiles (nem_prow_gen12) A sweep of nem_prow's PR and PO over {2, 4, 8} on Iris Xe (LFM3-1.7B prefill, 1,024 tokens, windows of 32, 64 and 128 rows) puts 8 prompt rows by 4 outputs a sub-group first: about 8 % above 4 x 4, while PO 8 spills (8 x 8: about 210 registers, 4 x 8: about 575) and runs 2 to 10 times slower. The image is renamed for what it is for; every pdense output sums as before, so the bits do not change. The 14 existing bmg images stay byte-equal to main. --- zig/build/xpu.zig | 2 +- zig/kernels/xpu/nem_prow.cl | 4 ++-- zig/src/xpu/kernels.zig | 4 ++-- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/zig/build/xpu.zig b/zig/build/xpu.zig index 8211729a02..1444ade410 100644 --- a/zig/build/xpu.zig +++ b/zig/build/xpu.zig @@ -190,7 +190,7 @@ const kernels = [_]Kernel{ .{ .name = "nem_rows_r4", .src = "nem_rows", .options = cl3 ++ " -DMAXR=4" }, .{ .name = "nem_rows_r8", .src = "nem_rows", .options = cl3 ++ " -DMAXR=8" }, .{ .name = "nem_prow", .options = cl3 }, - .{ .name = "nem_prow_t4", .src = "nem_prow", .options = cl3 ++ " -DPR=4 -DPO=4" }, + .{ .name = "nem_prow_gen12", .src = "nem_prow", .options = cl3 ++ " -DPO=4" }, .{ .name = "nem_attn_dec", .options = cl3 }, }; diff --git a/zig/kernels/xpu/nem_prow.cl b/zig/kernels/xpu/nem_prow.cl index 9337e1f02a..3a88dcecfa 100644 --- a/zig/kernels/xpu/nem_prow.cl +++ b/zig/kernels/xpu/nem_prow.cl @@ -1,10 +1,10 @@ // Nemotron-H prompt rows (any count) with the decode step's arithmetic: each output sums as the one-token path does. #define SG 16 #ifndef PR -#define PR 8 // prompt rows a sub-group holds (nem_prow_t4: 4, fewer registers for Gen12LP) +#define PR 8 // prompt rows a sub-group holds #endif #ifndef PO -#define PO 8 // output rows a sub-group holds: one load and one sum of a word's inputs serve all of them +#define PO 8 // output rows a sub-group holds: one load and one sum of a word's inputs serve all (nem_prow_gen12: 4) #endif #define PSG (16 / PO) // sub-groups a work-group: 16 output rows inline float bf(ushort v) { return as_float((uint)v << 16); } diff --git a/zig/src/xpu/kernels.zig b/zig/src/xpu/kernels.zig index 04fb4d0b34..3731089997 100644 --- a/zig/src/xpu/kernels.zig +++ b/zig/src/xpu/kernels.zig @@ -24,6 +24,6 @@ pub const nem_rows: []const u8 = if (available) &Blob("spv_nem_rows").bytes else /// nem_rows built for windows of at most 1, 2, 4 and 8 rows (fewer registers a lane). pub const nem_rows_r: [4][]const u8 = if (available) .{ &Blob("spv_nem_rows_r1").bytes, &Blob("spv_nem_rows_r2").bytes, &Blob("spv_nem_rows_r4").bytes, &Blob("spv_nem_rows_r8").bytes } else .{ &.{}, &.{}, &.{}, &.{} }; pub const nem_prow: []const u8 = if (available) &Blob("spv_nem_prow").bytes else &.{}; -/// nem_prow with 4 x 4 tiles (prompt rows x outputs a sub-group): no register spills to speak of on Gen12LP. -pub const nem_prow_t4: []const u8 = if (available) &Blob("spv_nem_prow_t4").bytes else &.{}; +/// nem_prow with 4 outputs a sub-group (8 x 4 tiles): fastest of PR, PO in {2, 4, 8} on Gen12LP, where 8 x 8 spills. +pub const nem_prow_gen12: []const u8 = if (available) &Blob("spv_nem_prow_gen12").bytes else &.{}; pub const nem_attn_dec: []const u8 = if (available) &Blob("spv_nem_attn_dec").bytes else &.{}; From 8bf9d952cba2e532a823f76394ca20403616600b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 30 Oct 2026 15:48:10 +0100 Subject: [PATCH 04/14] tools/lfm3: ref64.py runs on numpy without torch, and a bf16-points pass Without torch (or with --numpy) the forward runs on numpy: MLX affine and bf16 weights are dequantized exactly; NVFP4 and FP8 stay torch-only and are refused by name. --bf16-points rounds where the Intel GPU engine stores bf16 (embedding rows and RoPE cos/sin too, attention P and logits never). dump writes the float64 hidden state after every layer and check replays an engine's greedy replies teacher-forced. The collect and miss-table CLI is unchanged. --- tools/lfm3/REF64.md | 94 ++++++++ tools/lfm3/ref64.py | 521 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 615 insertions(+) create mode 100644 tools/lfm3/REF64.md create mode 100644 tools/lfm3/ref64.py diff --git a/tools/lfm3/REF64.md b/tools/lfm3/REF64.md new file mode 100644 index 0000000000..1f6d4730ed --- /dev/null +++ b/tools/lfm3/REF64.md @@ -0,0 +1,94 @@ +# Float64 reference receipt: LFM3-1.7B NVFP4 on CUDA + +- Checkpoint: [LiquidAI/LFM3-1.7B-NVFP4](https://huggingface.co/LiquidAI/LFM3-1.7B-NVFP4) at revision + `b30025bbbbcbf62f6c110b40d51d36af6579f6a2`. +- Engine: `9265a0f5f`, which is the PR head `08b338c7c` before its rebase (range-diff: the same patches except one + test-only line in `contract_limits.zig`). Served on a GB10 (sm_121), greedy, drafts off, + `--parallel 8 --context 32768`. The `collect` step was repeated on 2026-10-28 and returned the same 957 tokens. +- Reference: `ref64.py` on the CPU in float64 (Python 3.12, torch 2.14, safetensors 0.8, tokenizers 0.23), 2026-10-28. + +Reproduce (the reference pass takes about a minute and 20 GB of host memory): + +``` +zig-out/native/bin/tensorfold-native serve MODEL_DIR --port 8080 --parallel 8 --context 32768 +python -I tools/lfm3/ref64.py collect http://127.0.0.1:8080 MODEL_DIR engine.json +python -I tools/lfm3/ref64.py MODEL_DIR engine.json --bf16 +``` + +Rows. A prompt row is the engine's one-token greedy pick after the prefix [0, t] of a prompt, one request per prefix. +A reply row is a token of the engine's 32-token greedy reply, teacher-forced on that reply. `t` is the logits row after +token t, so prompt row P - 1 and reply token 0 are the same position. The prompts (`prompts()` in the script) are +capitals (20 tokens), code (40), a chat turn (19) and items (750, a list of 60 items and a question; it crosses the +256-key sliding window). + +The float64 pass dequantizes the NVFP4 weights exactly. The `--bf16` pass adds the engine's bf16 storage points (see +the script's docstring). It shows which side of a miss rounding alone lands on; it is not bit-exact. + +## Result + +- f64: prompt rows 825/829, reply rows 127/128 +- bf16: prompt rows 827/829, reply rows 126/128 + +Float64 top-1 misses (margin = f64 logit of its top minus f64 logit of the engine's token; bf16 margin = the same difference in the bf16-points pass, < 0 means it prefers the engine's token): + +| case | row | t | engine | float64 | f64 margin | bf16 margin | bf16 picks | +|---|---|---|---|---|---|---|---| +| capitals | reply | 24 | 11097 " Spain" | 3888 " Japan" | 0.1176 | +0.1250 (pre-round +0.1802) | f64's | +| chat | prompt | 12 | 593 " per" | 22 "." | 0.1530 | -0.0625 (pre-round -0.0671) | engine's | +| items | prompt | 5 | 1008 " first" | 1579 " number" | 0.0996 | +0.0000 (pre-round +0.1009) | engine's | +| items | prompt | 9 | 29 "5" | 654 "10" | 0.0033 | +0.0000 (pre-round +0.0093) | engine's | +| items | prompt | 180 | 28 "4" | 30 "6" | 0.0674 | +0.0000 (pre-round +0.0076) | engine's | + +bf16-points misses where float64 agrees with the engine (margin = bf16 logit of its pick minus the engine's token; f64 margin = how far float64 prefers the engine's token): + +| case | row | t | engine | bf16 pick | bf16 margin | f64 margin | +|---|---|---|---|---|---|---| +| items | prompt | 4 | 440 " The" | 229 " " | 0.0000 (pre-round -0.0985) | 0.0881 | +| items | prompt | 631 | 32 "8" | 25 "1" | 0.0000 (pre-round -0.0467) | 0.0636 | +| items | reply | 772 | 41501 " Item" | 1672 " So" | 0.0000 (pre-round -0.0814) | 0.1797 | + +## FP8 and MLX-4bit (2026-10-29) + +- Checkpoints: [LiquidAI/LFM3-1.7B-FP8](https://huggingface.co/LiquidAI/LFM3-1.7B-FP8) at + `243955da95b1cd68026f293b29de4b4c13f3638f` (block FP8, dequantized exactly: e4m3 times its block's bf16 scale) and + [LiquidAI/LFM3-1.7B-MLX-4bit](https://huggingface.co/LiquidAI/LFM3-1.7B-MLX-4bit) at + `d0ca73dcfa70ce18733754a492cdc1fb0123262d` (MLX affine, dequantized exactly). +- Served the same way with `--no-drafts`; the same prompts and rows. Only the float64 pass was run: the `--bf16` + pass models the NVFP4 engine's rounding points. + +FP8: prompt rows 824/829, reply rows 128/128. + +| case | row | t | engine | float64 | f64 margin | +|---|---|---|---|---|---| +| items | prompt | 131 | 30 "6" | 1325 "24" | 0.0348 | +| items | prompt | 582 | 654 "10" | 1035 "14" | 0.0756 | +| items | prompt | 619 | 25 "1" | 29 "5" | 0.0066 | +| items | prompt | 631 | 27 "3" | 32 "8" | 0.0290 | +| items | prompt | 693 | 9402 " stone" | 17937 " stones" | 0.0773 | + +MLX-4bit: prompt rows 825/829, reply rows 128/128. + +| case | row | t | engine | float64 | f64 margin | +|---|---|---|---|---|---| +| items | prompt | 10 | 6802 " items" | 93266 " marbles" | 0.0466 | +| items | prompt | 61 | 41501 " Item" | 29155 " Question" | 0.0042 | +| items | prompt | 595 | 27 "3" | 1035 "14" | 0.0300 | +| items | prompt | 748 | 34 ":" | 10494 " holds" | 0.0848 | + +## Without torch: the numpy path and the XPU's bf16 points + +`ref64.py` runs on numpy when torch is not installed (or with `--numpy`): MLX affine and bf16 weights are dequantized +exactly as on the torch path; NVFP4 and FP8 exports are refused by name and need torch. `--bf16-points` replaces the +`--bf16` pass with the Intel GPU (XPU) engine's storage points: the embedding rows and RoPE cos/sin rounded to bf16 too, +attention P and the logits never. `dump` writes the float64 hidden state after every layer for the XPU forward's +layer-by-layer check, and `check` replays an engine's greedy replies teacher-forced: + +``` +python -I tools/lfm3/ref64.py dump MODEL_DIR OUT_DIR +python -I tools/lfm3/ref64.py check MODEL_DIR tokens.json # {"name": {"prompt": [...], "reply": [...]}} +python -I tools/lfm3/ref64.py MODEL_DIR engine.json --bf16-points +``` + +On LiquidAI/LFM3-1.7B-MLX-4bit (Python 3.12, numpy 2.5, no torch; 3 min 14 s and 4.1 GB for `dump`): the three +prompts' bf16-points gap reaches 0.041 (capitals), 0.025 (chat) and 0.111 (items, worst row of any layer), and the +XPU engine's 64-token greedy replies agree with the float64 top-1 at 192 of 192 positions. diff --git a/tools/lfm3/ref64.py b/tools/lfm3/ref64.py new file mode 100644 index 0000000000..9f78aa46b2 --- /dev/null +++ b/tools/lfm3/ref64.py @@ -0,0 +1,521 @@ +#!/usr/bin/env python3 +"""Float64 CPU reference of the LFM3 text decoder (LFM3-1.7B), checked against the CUDA engine's greedy tokens. + +The stored weights are dequantized exactly and every op runs in float64: + nvfp4 ModelOpt NVFP4 (U8 two E2M1 a byte, low nibble first; e4m3 scale per 16 values; f32 weight_scale_2) + mlx MLX affine (bits and group size from config.json "quantization", per-module overrides) + fp8 compressed-tensors block FP8 (e4m3; a bf16 scale per 128 x 128 block) + bf16 plain bf16 weights +The math follows modeling_lfm3.py: sandwich norms (h += post_norm(mixer(norm(h)))), RMSNorm w * x * rsqrt(ms + eps); +conv y = C * causal_depthwise_conv(B * x) with the last tap on the current row; attention with per-head q/k RMSNorm, +RoPE (rotate_half; full layers partial 0.5 in the zero-tail form: inv_freq[i] = theta^(-2i/128) for i < 32, 0 after; +sliding layers all 64 frequencies), GQA, a 256-key sliding window, a per-head output gate sigmoid(gate_proj(x)); the +head is tied to embed_tokens. + +The engine's greedy tokens come from a running server (`collect`): + prompt rows the one-token greedy pick after each prefix [0, t] of a prompt (a fresh request per prefix) + reply rows a 32-token greedy reply, checked teacher-forced on the engine's own reply +Row t is the logits row after token t, so prompt row P - 1 and reply token 0 are the same position. + +--bf16 adds a second pass with the NVFP4 weight-only CUDA engine's bf16 storage points (float64 inside every op, +round-to-nearest-even to bf16 where the kernels store bf16): every RMSNorm (bf16(w * bf16(h * inv))), residual, +projection outputs, conv B*x / acc / C*y, q/k norm and RoPE, the gate, attention output, SwiGLU and the logits. Weights +are exact on both paths. Attention P is rounded to bf16 before PV on prompt rows and sliding-window decode rows and kept +unrounded on full-attention decode rows, taken against the row's final max (the kernels' running max is not +replayed). Greedy over bf16 logits breaks ties to the lower id, as the engine does. This pass shows which side of a +miss rounding alone lands on; it is not bit-exact. + +Without torch (or with --numpy) the forward runs on numpy: MLX affine and bf16 weights are dequantized exactly; NVFP4 +and FP8 stay torch-only and are refused by name. --bf16-points rounds where the XPU kernels store bf16 instead of the +CUDA engine: the embedding rows and RoPE cos and sin as well, attention P and the logits never (numpy path). + +usage: + ref64.py collect URL MODEL_DIR ENGINE_JSON the greedy rows of prompts() from a server at URL + ref64.py MODEL_DIR ENGINE_JSON [--bf16 | --bf16-points] [--numpy] + prints the miss table (markdown) and the agreement totals + ref64.py dump MODEL_DIR OUT_DIR capitals, chat, items: ids, the hidden state after every layer and + the final norm (f32), float64 top-1 per row, and the bf16-points + pass's per-layer gap (meta.json); numpy path + ref64.py check MODEL_DIR TOKENS_JSON teacher-forced float64 top-1 of {name: {prompt, reply}}; numpy path +""" +import json +import math +import os +import re +import sys +import urllib.request + +import struct + +import numpy as np + +try: + import torch + from safetensors import safe_open +except ImportError: # the numpy path: MLX affine and bf16 weights + torch = None + +if torch is not None: + torch.set_default_dtype(torch.float64) + F64 = torch.float64 + E2M1 = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6, -0.0, -0.5, -1, -1.5, -2, -3, -4, -6], dtype=F64) + + +def complete(url, ids, n): + body = {"prompt": ids, "max_tokens": n, "temperature": 0, "ignore_eos": True, "draft": False, + "return_token_ids": True} + req = urllib.request.Request(url + "/v1/completions", data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=900) as r: + out = json.loads(r.read()) + return out.get("tensorfold", out)["token_ids"] + + +def prompts(n_items): + """The receipt's prompt set: texts, encoded with the checkpoint's tokenizer.json after BOS.""" + colors = "red green blue amber violet".split() + items = " ".join(f"Item {i}: the {colors[i % 5]} box holds {7 * i % 23} stones." for i in range(n_items)) + return [("capitals", "The capital of France is Paris. The capital of Germany is Berlin. " + "The capital of Italy is"), + ("code", "def fibonacci(n):\n if n < 2:\n return n\n" + " return fibonacci(n - 1) + fibonacci(n - 2)\n\n# Rewrite iteratively:\ndef fibonacci(n):"), + ("chat", "<|im_start|>user\nExplain why the sky is blue in two sentences.<|im_end|>\n" + "<|im_start|>assistant\n"), + ("items", items + " Question: how many stones does item 3 hold? Answer: item 3 holds")] + + +def collect(url, model_dir, out_path, reply_len=32): + from tokenizers import Tokenizer + tok = Tokenizer.from_file(f"{model_dir}/tokenizer.json") + bos = json.load(open(f"{model_dir}/config.json"))["bos_token_id"] + cases = [] + for name, text in prompts(60): + ids = [bos] + tok.encode(text).ids + picks = [complete(url, ids[: t + 1], 1)[0] for t in range(len(ids))] + cases.append({"name": name, "prompt": ids, "prompt_picks": picks, "reply": complete(url, ids, reply_len)}) + print(f"{name}: {len(ids)} prompt rows, {len(cases[-1]['reply'])} reply tokens", file=sys.stderr) + json.dump({"cases": cases}, open(out_path, "w")) + + +def bf16(x): + """Nearest bf16 (8 significant bits, ties to even), kept as float64: exact, no float32 double rounding.""" + m, e = torch.frexp(x) + return torch.ldexp(torch.round(m * 256.0), (e - 8).to(torch.int32)).to(F64) + + +def same(x): + return x + + +def rot(t, cos, sin): + hd = t.shape[-1] + return t * cos + torch.cat([-t[..., hd // 2:], t[..., : hd // 2]], -1) * sin + + +def strip(k): + return re.sub(r"^(model\.language_model\.|language_model\.|model\.)+", "", k) + + +class Model: + def __init__(self, d): + self.d = d + root = json.load(open(f"{d}/config.json")) + self.cfg = root.get("text_config", root) + idx = f"{d}/model.safetensors.index.json" + if os.path.exists(idx): + self.map = json.load(open(idx))["weight_map"] + else: + with safe_open(f"{d}/model.safetensors", "pt") as f: + self.map = {k: "model.safetensors" for k in f.keys()} + self.files, self.cache = {}, {} + self.prefix = "language_model.model." if "language_model.model.embed_tokens.weight" in self.map else "model." + self.mlx = root.get("quantization") if root.get("quantization", {}).get("mode", "affine") == "affine" else None + + def raw(self, k): + f = self.map[k] + if f not in self.files: + self.files[f] = safe_open(f"{self.d}/{f}", "pt") + return self.files[f].get_tensor(k) + + def vec(self, name): + return self.raw(self.prefix + name).to(F64) + + def lin(self, name, rows=None): + """A Linear weight [out, in] in float64, dequantized exactly by its stored format.""" + p = self.prefix + name + w = self.raw(p + ".weight") + if rows is not None: + w = w[rows] + if self.mlx and p + ".scales" in self.map: + s, b = self.raw(p + ".scales"), self.raw(p + ".biases") + if rows is not None: + s, b = s[rows], b[rows] + bits = next((v["bits"] for k, v in self.mlx.items() if isinstance(v, dict) and strip(k) == strip(p)), + self.mlx["bits"]) + by = w.contiguous().view(torch.uint8).to(torch.int64) + nb = by.shape[-1] * 8 + bit = ((by.unsqueeze(-1) >> torch.arange(8)) & 1).reshape(*by.shape[:-1], nb) + q = (bit[..., : nb // bits * bits].reshape(*by.shape[:-1], -1, bits) << torch.arange(bits)).sum(-1) + q = q.to(F64).reshape(*q.shape[:-1], s.shape[-1], -1) + return (q * s.to(F64)[..., None] + b.to(F64)[..., None]).flatten(-2) + if p + ".weight_scale_2" in self.map: + sc = self.raw(p + ".weight_scale") + if rows is not None: + sc = sc[rows] + n = w.shape[0] + codes = torch.stack([w & 0xF, w >> 4], dim=-1).reshape(n, -1).long() + v = E2M1[codes].view(n, -1, 16) * sc.to(F64)[..., None] + return (v * self.raw(p + ".weight_scale_2").to(F64)).reshape(n, -1) + if w.dtype == torch.float8_e4m3fn: + sc = self.raw(p + ".weight_scale").to(F64).repeat_interleave(128, 0).repeat_interleave(128, 1) + return w.to(F64) * (sc if rows is None else sc[rows]) + return w.to(F64) + + def W(self, name): + if name not in self.cache: + self.cache[name] = self.lin(name) + return self.cache[name] + + def forward(self, ids, P, bf): + """Logits [T, V] (float64, not rounded). bf = same: plain float64; bf = bf16: the engine's storage points with + rows [0, P) on the prompt path and the rest on the decode path.""" + cfg, T = self.cfg, len(ids) + D, H, HK, hd = cfg["hidden_size"], cfg["num_attention_heads"], cfg["num_key_value_heads"], cfg["head_dim"] + taps, vec, W = cfg["conv_L_cache"], self.vec, self.W + + def rms(h, w): # the engine's form: bf16(w * bf16(h * inv)), inv from the bf16 row + return bf(w * bf(h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + cfg["norm_eps"]))) + + tabs = {} + for kind, rp in cfg["rope_parameters"].items(): + angles = int(rp.get("partial_rotary_factor", 1.0) * hd) // 2 + inv = 1.0 / rp["rope_theta"] ** (torch.arange(0, 2 * angles, 2, dtype=F64) / hd) # the full head dim's + inv = torch.cat([inv, torch.zeros(hd // 2 - angles)]) + ang = torch.arange(T, dtype=F64)[:, None] * inv[None] + tabs[kind] = torch.cat([ang.cos()] * 2, -1), torch.cat([ang.sin()] * 2, -1) + i, j = torch.arange(T)[:, None], torch.arange(T)[None] + masks = {"full_attention": j > i, "sliding_attention": (j > i) | (i - j >= cfg["sliding_window"])} + h = self.lin("embed_tokens", rows=torch.tensor(ids)) + for L, kind in enumerate(cfg["layer_types"]): + p = f"layers.{L}." + x = rms(h, vec(p + "operator_norm.weight")) + if kind == "conv": + B, C, xx = bf(x @ W(p + "conv.in_proj").T).chunk(3, -1) + bx = bf(B * xx) + w = vec(p + "conv.conv.weight").reshape(D, taps) + pad = torch.cat([torch.zeros(taps - 1, D), bx]) + acc = torch.zeros(T, D) + for k in range(taps): # oldest tap first + acc = acc + pad[k:k + T] * w[:, k] + o = bf(bf(C * bf(acc)) @ W(p + "conv.out_proj").T) + else: + q, k, v = (bf(x @ W(p + f"self_attn.{n}_proj").T) for n in "qkv") + g = bf(x @ W(p + "self_attn.gate_proj").T) # [T, H] + q = rms(q.reshape(T, H, hd), vec(p + "self_attn.q_norm.weight")).transpose(0, 1) # [H, T, hd] + k = rms(k.reshape(T, HK, hd), vec(p + "self_attn.k_norm.weight")).transpose(0, 1) + q, k = bf(rot(q, *tabs[kind])), bf(rot(k, *tabs[kind])) + k = k.repeat_interleave(H // HK, 0) + v = v.reshape(T, HK, hd).transpose(0, 1).repeat_interleave(H // HK, 0) + s = ((q @ k.transpose(1, 2)) / math.sqrt(hd)).masked_fill(masks[kind], -math.inf) + e = torch.exp(s - s.max(-1, keepdim=True).values) + pv = bf(e) + if kind == "full_attention": + pv = torch.cat([pv[:, :P], e[:, P:]], 1) # full-attention decode keeps P unrounded + a = bf((pv @ v) / e.sum(-1, keepdim=True)).transpose(0, 1) # [T, H, hd] + a = bf(a * bf(torch.sigmoid(g))[..., None]) + o = bf(a.reshape(T, H * hd) @ W(p + "self_attn.o_proj").T) + h = bf(h + rms(o, vec(p + "post_operator_norm.weight"))) + x = rms(h, vec(p + "ffn_norm.weight")) + f = p + "feed_forward." + act = bf(bf(torch.nn.functional.silu(bf(x @ W(f + "w1").T))) * bf(x @ W(f + "w3").T)) + h = bf(h + rms(bf(act @ W(f + "w2").T), vec(p + "post_ffn_norm.weight"))) + x = rms(h, vec("embedding_norm.weight")) + V = cfg["vocab_size"] + return torch.cat([x @ self.lin("embed_tokens", rows=torch.arange(a, min(a + 16000, V))).T + for a in range(0, V, 16000)], -1) + + +def torch_logits(self, ids, P, points=None): + """Logits [T, V] as a numpy array: float64, or the CUDA engine's bf16 points (points == "engine").""" + return self.forward(ids, P, bf16 if points else same).numpy() + + +Model.logits = torch_logits + + +def greedy(lg): + """argmax in (value desc, id asc) order, as the engine picks over its bf16 logits.""" + return int(np.flatnonzero(lg == lg.max())[0]) + + +def npbf16(x): + """Nearest bf16 (ties to even) of float64 values, kept as float64.""" + m, e = np.frexp(x) + return np.ldexp(np.round(m * 256.0), e - 8) + + +class NpModel: + """The float64 forward on numpy: MLX affine and bf16 weights read from the safetensors bytes, exactly.""" + + def __init__(self, d): + root = json.load(open(f"{d}/config.json")) + self.cfg = root.get("text_config", root) + idx = f"{d}/model.safetensors.index.json" + shards = sorted(set(json.load(open(idx))["weight_map"].values())) if os.path.exists(idx) else ["model.safetensors"] + self.hdr = {} + for s in shards: + with open(f"{d}/{s}", "rb") as f: + n = struct.unpack("> (np.arange(8, dtype=np.uint32) * 4)) & 15 + else: + bit = np.unpackbits(np.ascontiguousarray(w).view(np.uint8), axis=-1, bitorder="little") + nb = bit.shape[-1] // bits * bits + q = bit[..., :nb].reshape(n, -1, bits).astype(np.int64) @ (1 << np.arange(bits)) + q = q.reshape(n, s.shape[-1], -1).astype(np.float64) + return (q * s[..., None] + b[..., None]).reshape(n, -1) + + def logits(self, ids, P, points=None, dump=None): + """Logits [T, V] in float64; points None, "engine" (--bf16) or "xpu" (--bf16-points); dump(L, h) per layer.""" + cfg, T = self.cfg, len(ids) + D, H, HK, hd = cfg["hidden_size"], cfg["num_attention_heads"], cfg["num_key_value_heads"], cfg["head_dim"] + bf = npbf16 if points else same + xpu = points == "xpu" + + def rms(h, w): + return bf(w * bf(h / np.sqrt((h * h).mean(-1, keepdims=True) + cfg["norm_eps"]))) + + def mm(x, name): + return x @ self.lin(name).T + + tabs = {} + for kind, rp in cfg["rope_parameters"].items(): + angles = int(rp.get("partial_rotary_factor", 1.0) * hd) // 2 + inv = 1.0 / rp["rope_theta"] ** (np.arange(0, 2 * angles, 2, dtype=np.float64) / hd) + inv = np.concatenate([inv, np.zeros(hd // 2 - angles)]) + ang = np.arange(T, dtype=np.float64)[:, None] * inv[None] + cos, sin = np.concatenate([np.cos(ang)] * 2, -1), np.concatenate([np.sin(ang)] * 2, -1) + tabs[kind] = (bf(cos), bf(sin)) if xpu else (cos, sin) + i, j = np.arange(T)[:, None], np.arange(T)[None] + masks = {"full_attention": j > i, "sliding_attention": (j > i) | (i - j >= cfg["sliding_window"])} + h = self.lin("embed_tokens", rows=np.array(ids)) + if xpu: + h = bf(h) + for L, kind in enumerate(cfg["layer_types"]): + p = f"layers.{L}." + x = rms(h, self.vec(p + "operator_norm.weight")) + if kind == "conv": + B, C, xx = np.split(bf(mm(x, p + "conv.in_proj")), 3, -1) + bx = bf(B * xx) + w = self.vec(p + "conv.conv.weight").reshape(D, -1) + taps = w.shape[1] + pad = np.concatenate([np.zeros((taps - 1, D)), bx]) + acc = np.zeros((T, D)) + for k in range(taps): + acc = acc + pad[k:k + T] * w[:, k] + o = bf(mm(bf(C * bf(acc)), p + "conv.out_proj")) + else: + q, k, v = (bf(mm(x, p + f"self_attn.{n}_proj")) for n in "qkv") + g = bf(mm(x, p + "self_attn.gate_proj")) + q = rms(q.reshape(T, H, hd), self.vec(p + "self_attn.q_norm.weight")).transpose(1, 0, 2) + k = rms(k.reshape(T, HK, hd), self.vec(p + "self_attn.k_norm.weight")).transpose(1, 0, 2) + cos, sin = tabs[kind] + q = bf(q * cos + np.concatenate([-q[..., hd // 2:], q[..., : hd // 2]], -1) * sin) + k = bf(k * cos + np.concatenate([-k[..., hd // 2:], k[..., : hd // 2]], -1) * sin) + k = np.repeat(k, H // HK, 0) + v = np.repeat(v.reshape(T, HK, hd).transpose(1, 0, 2), H // HK, 0) + s = (q @ k.transpose(0, 2, 1)) / math.sqrt(hd) + s[:, masks[kind]] = -np.inf + e = np.exp(s - s.max(-1, keepdims=True)) + pv = e + if points == "engine": + pv = bf(e) + if kind == "full_attention": + pv[:, P:] = e[:, P:] # full-attention decode keeps P unrounded + a = bf((pv @ v) / e.sum(-1, keepdims=True)).transpose(1, 0, 2) + a = bf(a * bf(1.0 / (1.0 + np.exp(-g)))[..., None]) + o = bf(mm(a.reshape(T, H * hd), p + "self_attn.o_proj")) + h = bf(h + rms(o, self.vec(p + "post_operator_norm.weight"))) + x = rms(h, self.vec(p + "ffn_norm.weight")) + f = p + "feed_forward." + a1 = bf(mm(x, f + "w1")) + act = bf(bf(a1 / (1.0 + np.exp(-a1))) * bf(mm(x, f + "w3"))) + h = bf(h + rms(bf(mm(act, f + "w2")), self.vec(p + "post_ffn_norm.weight"))) + if dump: + dump(L, h) + x = rms(h, self.vec("embedding_norm.weight")) + if dump: + dump(-1, x) + V = cfg["vocab_size"] + return np.concatenate([x @ self.lin("embed_tokens", rows=np.arange(a, min(a + 8000, V))).T + for a in range(0, V, 8000)], -1) + + +def rel_rows(a, b): + """Per-row ||a - b|| / ||b||: mean and max.""" + r = np.linalg.norm(a - b, axis=-1) / np.linalg.norm(b, axis=-1) + return float(r.mean()), float(r.max()) + + +def dump(model_dir, out): + """capitals, chat, items (BOS, no added specials): float64 hidden states per layer, top-1, bf16-points gaps.""" + from tokenizers import Tokenizer + tok = Tokenizer.from_file(f"{model_dir}/tokenizer.json") + m = NpModel(model_dir) + bos = m.cfg["bos_token_id"] + summary = {} + for name, text in [p for p in prompts(60) if p[0] != "code"]: + ids = [bos] + tok.encode(text, add_special_tokens=False).ids + d = f"{out}/{name}" + os.makedirs(d, exist_ok=True) + f64h, xph = {}, {} + + def keep(L, h): + f64h[L] = h.copy() + h.astype(np.float32).tofile(f"{d}/h{L:02d}.f32" if L >= 0 else f"{d}/final.f32") + + lf = m.logits(ids, len(ids), None, keep) + lb = m.logits(ids, len(ids), "xpu", lambda L, h: xph.__setitem__(L, h.copy())) + gap = {L: rel_rows(xph[L], f64h[L]) for L in f64h} + top = lf.argmax(-1) + np.asarray(top, dtype=np.int32).tofile(f"{d}/top1.i32") + lf[-1].astype(np.float32).tofile(f"{d}/logits_last.f32") + agree = int((lb.argmax(-1) == top).sum()) + json.dump({"ids": ids, "T": len(ids), "top1": top.tolist(), "bf16_gap": {str(L): v for L, v in gap.items()}, + "bf16_top1_agree": agree}, open(f"{d}/meta.json", "w")) + summary[name] = {"T": len(ids), "bf16_gap_max": max(v[1] for v in gap.values()), "next": int(top[-1]), + "bf16_top1_agree": f"{agree}/{len(ids)}"} + print(name, json.dumps(summary[name]), file=sys.stderr, flush=True) + json.dump(summary, open(f"{out}/summary.json", "w"), indent=1) + + +def check(model_dir, path): + """Teacher-forced: the float64 top-1 after each prefix of prompt + reply against the reply's own next token.""" + m = NpModel(model_dir) + res = {} + for name, c in json.load(open(path)).items(): + P, reply = len(c["prompt"]), c["reply"] + lf = m.logits(c["prompt"] + reply[:-1], P)[P - 1:] + picks = lf.argmax(-1) + miss = [{"j": j, "engine": int(reply[j]), "f64": int(picks[j]), + "margin": float(lf[j].max() - lf[j, reply[j]])} for j in range(len(reply)) if picks[j] != reply[j]] + res[name] = {"agree": f"{len(reply) - len(miss)}/{len(reply)}", "misses": miss} + print(name, json.dumps(res[name]), file=sys.stderr, flush=True) + return res + + +def main(model_dir, engine_path, points, numpy_path): + mdl = NpModel(model_dir) if numpy_path or torch is None or points == "xpu" else Model(model_dir) + with_bf16 = points is not None + try: + from tokenizers import Tokenizer + tok = Tokenizer.from_file(f"{model_dir}/tokenizer.json") + + def show(t): + return f"{t} {json.dumps(tok.decode([t], skip_special_tokens=False))}".replace("|", "\\|") + except ImportError: + show = str + cases = json.load(open(engine_path))["cases"] + tot = {(m, k): [0, 0] for m in ("f64", "bf16") for k in ("prompt", "reply")} + f64_rows, bf16_rows = [], [] + for c in cases: + P = len(c["prompt"]) + ids = c["prompt"] + c["reply"][:-1] + print(f"{c['name']}: {len(ids)} rows", file=sys.stderr, flush=True) + lf = mdl.logits(ids, P) + lb_pre = mdl.logits(ids, P, points) if with_bf16 else None + lb = (npbf16(lb_pre) if points == "engine" else lb_pre) if with_bf16 else None + rows = [("prompt", t, tk) for t, tk in enumerate(c["prompt_picks"])] + rows += [("reply", P - 1 + j, tk) for j, tk in enumerate(c["reply"])] + for kind, t, eng in rows: + top = int(lf[t].argmax()) + tot["f64", kind][0] += top == eng + tot["f64", kind][1] += 1 + bcol = "" + if with_bf16: + bp = greedy(lb[t]) + tot["bf16", kind][0] += bp == eng + tot["bf16", kind][1] += 1 + if top != eng: + if with_bf16: + d, dp = lb[t, top] - lb[t, eng], lb_pre[t, top] - lb_pre[t, eng] + side = "engine's" if bp == eng else "f64's" if bp == top else f"other {show(bp)}" + bcol = f" {d:+.4f} (pre-round {dp:+.4f}) | {side} |" + f64_rows.append(f"| {c['name']} | {kind} | {t} | {show(eng)} | {show(top)} | " + f"{lf[t, top] - lf[t, eng]:.4f} |{bcol}") + elif with_bf16 and bp != eng: + bf16_rows.append(f"| {c['name']} | {kind} | {t} | {show(eng)} | {show(bp)} | " + f"{lb[t, bp] - lb[t, eng]:.4f} (pre-round {lb_pre[t, bp] - lb_pre[t, eng]:+.4f}) | " + f"{lf[t, eng] - lf[t, bp]:.4f} |") + mode = [m for m in ("f64", "bf16") if m == "f64" or with_bf16] + for m in mode: + print(f"{m}: prompt rows {tot[m, 'prompt'][0]}/{tot[m, 'prompt'][1]}, " + f"reply rows {tot[m, 'reply'][0]}/{tot[m, 'reply'][1]}") + print("\nFloat64 top-1 misses (margin = f64 logit of its top minus f64 logit of the engine's token" + + ("; bf16 margin = the same difference in the bf16-points pass, < 0 means it prefers the engine's token" + if with_bf16 else "") + "):\n") + print("| case | row | t | engine | float64 | f64 margin |" + (" bf16 margin | bf16 picks |" if with_bf16 else "")) + print("|---|---|---|---|---|---|" + ("---|---|" if with_bf16 else "")) + print("\n".join(f64_rows) or "(none)") + if with_bf16: + print("\nbf16-points misses where float64 agrees with the engine (margin = bf16 logit of its pick minus the " + "engine's token; f64 margin = how far float64 prefers the engine's token):\n") + print("| case | row | t | engine | bf16 pick | bf16 margin | f64 margin |\n|---|---|---|---|---|---|---|") + print("\n".join(bf16_rows) or "(none)") + + +if __name__ == "__main__": + a = sys.argv[1:] + flags = a[2:] + if len(a) == 4 and a[0] == "collect": + collect(*a[1:]) + elif len(a) == 3 and a[0] == "dump": + dump(a[1], a[2]) + elif len(a) == 3 and a[0] == "check": + json.dump(check(a[1], a[2]), sys.stdout, indent=1) + elif len(a) >= 2 and set(flags) <= {"--bf16", "--bf16-points", "--numpy"} and len(set(flags)) == len(flags) \ + and not {"--bf16", "--bf16-points"} <= set(flags): + pts = "engine" if "--bf16" in flags else "xpu" if "--bf16-points" in flags else None + main(a[0], a[1], pts, "--numpy" in flags) + else: + sys.exit(__doc__) From 9940ac79880a5e960dbfb0ad192519596f982962 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 30 Oct 2026 13:47:26 +0100 Subject: [PATCH 05/14] xpu: LFM3 kernels without DPAS (short conv, QK norm with partial RoPE, gated GQA attention over a ring, sandwich norm, SwiGLU) and their FP64 op tests --- zig/build/xpu.zig | 3 +- zig/kernels/xpu/lfm3.cl | 206 ++++++++++++++++++++++++++++++++++++ zig/src/xpu/kernels.zig | 1 + zig/tests/xpu/lfm3_attn.zig | 124 ++++++++++++++++++++++ zig/tests/xpu/lfm3_conv.zig | 92 ++++++++++++++++ zig/tests/xpu/lfm3_norm.zig | 112 ++++++++++++++++++++ zig/tests/xpu/lfm3_ops.zig | 95 +++++++++++++++++ zig/tests/xpu/lfm3_rope.zig | 108 +++++++++++++++++++ 8 files changed, 740 insertions(+), 1 deletion(-) create mode 100644 zig/kernels/xpu/lfm3.cl create mode 100644 zig/tests/xpu/lfm3_attn.zig create mode 100644 zig/tests/xpu/lfm3_conv.zig create mode 100644 zig/tests/xpu/lfm3_norm.zig create mode 100644 zig/tests/xpu/lfm3_ops.zig create mode 100644 zig/tests/xpu/lfm3_rope.zig diff --git a/zig/build/xpu.zig b/zig/build/xpu.zig index 1444ade410..c0d16b2e49 100644 --- a/zig/build/xpu.zig +++ b/zig/build/xpu.zig @@ -192,6 +192,7 @@ const kernels = [_]Kernel{ .{ .name = "nem_prow", .options = cl3 }, .{ .name = "nem_prow_gen12", .src = "nem_prow", .options = cl3 ++ " -DPO=4" }, .{ .name = "nem_attn_dec", .options = cl3 }, + .{ .name = "lfm3", .options = cl3 }, }; /// The runtime module; `with_kernels` false builds it host-only (empty images). @@ -257,7 +258,7 @@ fn hasDpas(device: []const u8) bool { /// The standalone op and model tests, each its own program `tf-xpu--test` (step `xpu-tests`). const test_programs = [_][]const u8{ - "panic", "nem_rows", "nem_prefill", "nem_gen", "nem_mtp", "nem_kl", "nem_norm_bench", "nem_mv_bench", "nem_slot", "nem_format", "nem_seam", "nem_attn_dec", "nem_fp64", "memory", + "panic", "nem_rows", "nem_prefill", "nem_gen", "nem_mtp", "nem_kl", "nem_norm_bench", "nem_mv_bench", "nem_slot", "nem_format", "nem_seam", "nem_attn_dec", "nem_fp64", "memory", "lfm3_conv", "lfm3_rope", "lfm3_attn", "lfm3_norm", }; /// `-Dxpu` (Linux): Nemotron's SPIR-V images, `tensorfold-xpu` and `tf-xpu-test`; needs ocloc, not nvcc or CUDA. diff --git a/zig/kernels/xpu/lfm3.cl b/zig/kernels/xpu/lfm3.cl new file mode 100644 index 0000000000..18492a2eca --- /dev/null +++ b/zig/kernels/xpu/lfm3.cl @@ -0,0 +1,206 @@ +// LFM3 ops in plain OpenCL C (no DPAS): bf16 rows, fp32 math, rounded to bf16 where modeling_lfm3.py rounds. +#pragma OPENCL FP_CONTRACT OFF +#define HD 128 // head dim +#ifndef GQ +#define GQ 4 // query heads per KV head (16 / 4) +#endif +#define NPI 16 // values an item of 256 holds in the row ops: n <= 4096 + +inline float bf(ushort v) { return as_float((uint)v << 16); } +inline ushort to_bf(float f) { + uint u = as_uint(f); + u += 0x7fffu + ((u >> 16) & 1u); + return (ushort)(u >> 16); +} +inline float bfr(float f) { return bf(to_bf(f)); } + +// Sum over a work-group of 256 items (16 sub-groups of 16) in a fixed order. +inline float wg_sum256(float v, __local float *part) { + v = sub_group_reduce_add(v); + if (get_sub_group_local_id() == 0) part[get_sub_group_id()] = v; + barrier(CLK_LOCAL_MEM_FENCE); + float t = 0.0f; + for (uint i = 0; i < 16; i++) t += part[i]; + barrier(CLK_LOCAL_MEM_FENCE); + return t; +} + +// RMSNorm of row get_group_id(0): y = bf16(w * bf16(x * rsqrt(mean(x^2) + eps))); 256 items, n % 256 == 0, n <= 4096. +__attribute__((intel_reqd_sub_group_size(16))) +__kernel void lfm3_rmsnorm(__global const ushort *x, __global const ushort *w, __global ushort *y, uint n, float eps) { + __local float part[16]; + const uint row = get_group_id(0); + const uint lid = get_local_id(0); + const uint per = n / 256; + float v[NPI]; + float ss = 0.0f; + for (uint k = 0; k < per; k++) { + v[k] = bf(x[row * n + lid + 256 * k]); + ss += v[k] * v[k]; + } + const float rs = rsqrt(wg_sum256(ss, part) / (float)n + eps); + for (uint k = 0; k < per; k++) y[row * n + lid + 256 * k] = to_bf(bf(w[lid + 256 * k]) * bfr(v[k] * rs)); +} + +// Sandwich norm step: x = bf16(x + rmsnorm(d, wp)) in place, then y = rmsnorm(x, wn); 256 items a row. +__attribute__((intel_reqd_sub_group_size(16))) +__kernel void lfm3_post_add_norm(__global const ushort *d, __global ushort *x, __global const ushort *wp, __global const ushort *wn, __global ushort *y, + uint n, float eps) { + __local float part[16]; + const uint row = get_group_id(0); + const uint lid = get_local_id(0); + const uint per = n / 256; + float v[NPI]; + float ss = 0.0f; + for (uint k = 0; k < per; k++) { + v[k] = bf(d[row * n + lid + 256 * k]); + ss += v[k] * v[k]; + } + const float rd = rsqrt(wg_sum256(ss, part) / (float)n + eps); + ss = 0.0f; + for (uint k = 0; k < per; k++) { + const uint i = row * n + lid + 256 * k; + const float dn = bfr(bf(wp[lid + 256 * k]) * bfr(v[k] * rd)); + v[k] = bfr(bf(x[i]) + dn); + x[i] = to_bf(v[k]); + ss += v[k] * v[k]; + } + const float rx = rsqrt(wg_sum256(ss, part) / (float)n + eps); + for (uint k = 0; k < per; k++) y[row * n + lid + 256 * k] = to_bf(bf(wn[lid + 256 * k]) * bfr(v[k] * rx)); +} + +// Gated short conv: rows [B | C | x], y = bf16(C * bf16(conv7(bf16(B * x)))); state: last six B * x, oldest first. +__kernel void lfm3_conv_rows(__global const ushort *proj, uint pd, __global const ushort *sin, __global ushort *sout, __global const ushort *cw, + __global ushort *y, uint n, uint D) { + const uint c = get_global_id(0); + if (c >= D) return; + float h[6], w[7]; + for (uint k = 0; k < 6; k++) h[k] = bf(sin[k * D + c]); + for (uint k = 0; k < 7; k++) w[k] = bf(cw[c * 7 + k]); + for (uint r = 0; r < n; r++) { + __global const ushort *p = proj + r * pd; + const float bx = bfr(bf(p[c]) * bf(p[2 * D + c])); + float acc = 0.0f; + for (uint k = 0; k < 6; k++) acc += h[k] * w[k]; + acc += bx * w[6]; + y[r * D + c] = to_bf(bf(p[D + c]) * bfr(acc)); + for (uint k = 0; k < 5; k++) h[k] = h[k + 1]; + h[5] = bx; + } + for (uint k = 0; k < 6; k++) sout[k * D + c] = to_bf(h[k]); +} + +// Per-head RMSNorm, RoPE on pairs (d, d + 64), bf16 cos/sin; grid (heads, rows); to row r or slot (pos0 + r) % ring. +__kernel void lfm3_qk_norm_rope(__global const ushort *src, uint src_stride, uint src_off, __global ushort *dst, uint dst_stride, + __global const ushort *w, __global const float *inv_freq, uint pos0, uint ring, float eps) { + __local float part[4]; + const uint head = get_group_id(0); + const uint r = get_group_id(1); + const uint i = get_local_id(0); + __global const ushort *s = src + r * src_stride + src_off + head * HD; + const float a = bf(s[i]), b = bf(s[i + 64]); + float ss = sub_group_reduce_add(a * a + b * b); + if (get_sub_group_local_id() == 0) part[get_sub_group_id()] = ss; + barrier(CLK_LOCAL_MEM_FENCE); + const uint nsg = get_num_sub_groups(); + float t = 0.0f; + for (uint k = 0; k < nsg; k++) t += part[k]; + const float rs = rsqrt(t / (float)HD + eps); + const float na = bfr(bf(w[i]) * bfr(a * rs)); + const float nb = bfr(bf(w[i + 64]) * bfr(b * rs)); + const float f = (float)(pos0 + r) * inv_freq[i]; + const float co = bfr(cos(f)), si = bfr(sin(f)); + const uint row = ring == 0 ? r : (pos0 + r) % ring; + __global ushort *o = dst + row * dst_stride + head * HD; + o[i] = to_bf(na * co - nb * si); + o[i + 64] = to_bf(nb * co + na * si); +} + +// Copies n_elem bf16 of row r to KV slot (pos0 + r) % ring; grid (ceil(n_elem / 64), rows). +__kernel void lfm3_v_store(__global const ushort *src, uint src_stride, uint src_off, __global ushort *dst, uint n_elem, uint pos0, uint ring) { + const uint e = get_global_id(0); + const uint r = get_group_id(1); + if (e < n_elem) dst[((pos0 + r) % ring) * n_elem + e] = src[r * src_stride + src_off + e]; +} + +// First key of row r's attention: window == 0 sees every earlier position, else the last `window` positions. +inline uint key_lo(uint p, uint window) { return (window == 0 || p + 1 <= window) ? 0 : p + 1 - window; } + +// Split-K GQA pass 1: group (kv head, chunk, row), a sub-group per q head, keys at slot k % ring; unnormalised o, m, l. +__attribute__((intel_reqd_sub_group_size(16))) +__kernel void lfm3_attn_partial(__global const ushort *q, __global const ushort *kc, __global const ushort *vc, __global float *po, __global float *pm, + __global float *pl, uint pos0, uint chunk, uint ring, uint window, uint nkv, uint maxc, float scale) { + const uint kvh = get_group_id(0); + const uint c = get_group_id(1); + const uint z = get_group_id(2); + const uint g = get_sub_group_id(); + const uint lane = get_sub_group_local_id(); + const uint nq = nkv * GQ; + const uint h = kvh * GQ + g; + const uint p = pos0 + z; + const uint lo = key_lo(p, window) + c * chunk; + const uint hi = min(p + 1, lo + chunk); + const float8 qv = as_float8(convert_uint8(vload8(0, q + (z * nq + h) * HD + lane * 8)) << 16); + float m = -INFINITY, l = 0.0f; + float8 o = (float8)(0.0f); + for (uint k = lo; k < hi; k++) { + const uint base = ((k % ring) * nkv + kvh) * HD + lane * 8; + const float8 kv = as_float8(convert_uint8(vload8(0, kc + base)) << 16); + const float8 pr = qv * kv; + const float s = sub_group_reduce_add(pr.s0 + pr.s1 + pr.s2 + pr.s3 + pr.s4 + pr.s5 + pr.s6 + pr.s7) * scale; + const float mn = fmax(m, s); + const float corr = exp(m - mn); + const float e = exp(s - mn); + const float8 vv = as_float8(convert_uint8(vload8(0, vc + base)) << 16); + l = l * corr + e; + o = o * corr + e * vv; + m = mn; + } + const uint slot = (z * nq + h) * maxc + c; + vstore8(o, 0, po + slot * HD + lane * 8); + if (lane == 0) { + pm[slot] = m; + pl[slot] = l; + } +} + +// Pass 2: row z's chunks of head h merged in order, bf16(O / L) times bf16(sigmoid(gate)); grid (nq, rows), 16 items. +__attribute__((intel_reqd_sub_group_size(16))) +__kernel void lfm3_attn_merge_gate(__global const float *po, __global const float *pm, __global const float *pl, __global const ushort *g, + __global ushort *out, uint pos0, uint chunk, uint window, uint nq, uint maxc) { + const uint h = get_group_id(0); + const uint z = get_group_id(1); + const uint lane = get_sub_group_local_id(); + const uint p = pos0 + z; + const uint len = p + 1 - key_lo(p, window); + const uint nch = (len + chunk - 1) / chunk; + const uint base = (z * nq + h) * maxc; + float mx = -INFINITY; + for (uint c = 0; c < nch; c++) mx = fmax(mx, pm[base + c]); + float L = 0.0f; + float8 O = (float8)(0.0f); + for (uint c = 0; c < nch; c++) { + const float e = exp(pm[base + c] - mx); + L += pl[base + c] * e; + O += vload8(0, po + (base + c) * HD + lane * 8) * e; + } + const float gate = bfr(1.0f / (1.0f + exp(-bf(g[z * nq + h])))); + const float8 r = O / L; + __global ushort *dst = out + (z * nq + h) * HD + lane * 8; + dst[0] = to_bf(bfr(r.s0) * gate); + dst[1] = to_bf(bfr(r.s1) * gate); + dst[2] = to_bf(bfr(r.s2) * gate); + dst[3] = to_bf(bfr(r.s3) * gate); + dst[4] = to_bf(bfr(r.s4) * gate); + dst[5] = to_bf(bfr(r.s5) * gate); + dst[6] = to_bf(bfr(r.s6) * gate); + dst[7] = to_bf(bfr(r.s7) * gate); +} + +// SwiGLU over n bf16 values: out = bf16(bf16(silu(a)) * b), a = w1 x (gate), b = w3 x (up). +__kernel void lfm3_swiglu(__global const ushort *a, __global const ushort *b, __global ushort *out, uint n) { + const uint i = get_global_id(0); + if (i >= n) return; + const float x = bf(a[i]); + out[i] = to_bf(bfr(x / (1.0f + exp(-x))) * bf(b[i])); +} diff --git a/zig/src/xpu/kernels.zig b/zig/src/xpu/kernels.zig index 3731089997..696d1a8cd4 100644 --- a/zig/src/xpu/kernels.zig +++ b/zig/src/xpu/kernels.zig @@ -27,3 +27,4 @@ pub const nem_prow: []const u8 = if (available) &Blob("spv_nem_prow").bytes else /// nem_prow with 4 outputs a sub-group (8 x 4 tiles): fastest of PR, PO in {2, 4, 8} on Gen12LP, where 8 x 8 spills. pub const nem_prow_gen12: []const u8 = if (available) &Blob("spv_nem_prow_gen12").bytes else &.{}; pub const nem_attn_dec: []const u8 = if (available) &Blob("spv_nem_attn_dec").bytes else &.{}; +pub const lfm3: []const u8 = if (available) &Blob("spv_lfm3").bytes else &.{}; diff --git a/zig/tests/xpu/lfm3_attn.zig b/zig/tests/xpu/lfm3_attn.zig new file mode 100644 index 0000000000..1bfad3d6ee --- /dev/null +++ b/zig/tests/xpu/lfm3_attn.zig @@ -0,0 +1,124 @@ +//! LFM3 GQA attention (partial + merge with gate) over a KV ring, full and sliding, against FP64 on random data. + +const std = @import("std"); +const o = @import("lfm3_ops.zig"); +const stop = @import("xpu").stop; + +const HD: usize = 128; +const NQ: u32 = 16; +const NKV: u32 = 4; +const CHUNK: u32 = 64; + +/// FP64 attention of row z (position pos0 + z) over keys by absolute position, window 0 = all, times sigmoid(gate). +fn refRows(q: []const f64, k: []const f64, v: []const f64, g: []const f64, n: usize, pos0: usize, window: usize, out: []f64) !void { + const scale = 1.0 / @sqrt(@as(f64, HD)); + const s = try o.gpa.alloc(f64, pos0 + n); + defer o.gpa.free(s); + for (0..n) |z| { + const p = pos0 + z; + const lo = if (window == 0 or p + 1 <= window) 0 else p + 1 - window; + for (0..NQ) |h| { + const kvh = h / (NQ / NKV); + const qh = q[(z * NQ + h) * HD ..][0..HD]; + var mx: f64 = -std.math.inf(f64); + for (lo..p + 1) |key| { + var a: f64 = 0; + for (0..HD) |d| a += qh[d] * k[(key * NKV + kvh) * HD + d]; + s[key] = a * scale; + mx = @max(mx, s[key]); + } + var den: f64 = 0; + const oh = out[(z * NQ + h) * HD ..][0..HD]; + @memset(oh, 0); + for (lo..p + 1) |key| { + const e = @exp(s[key] - mx); + den += e; + for (0..HD) |d| oh[d] += e * v[(key * NKV + kvh) * HD + d]; + } + const gate = 1.0 / (1.0 + @exp(-g[z * NQ + h])); + for (oh) |*x| x.* = x.* / den * gate; + } + } +} + +const Case = struct { what: []const u8, pos0: u32, n: u32, window: u32, ring: u32 }; + +pub const panic = std.debug.FullPanic(@import("rig.zig").panicDrain); +pub fn main() !void { + stop.install(); + var r = try o.rt.open(); + defer r.deinit(); + var m = try r.module(@import("xpu").kernels.lfm3); + var kp = try m.kernel("lfm3_attn_partial", .{ 64, 1, 1 }); + var km = try m.kernel("lfm3_attn_merge_gate", .{ 16, 1, 1 }); + var t: o.Tally = .{ .name = "lfm3_attn" }; + var prng = std.Random.DefaultPrng.init(13); + const rnd = prng.random(); + const maxc: u32 = 64; + const max_rows: u32 = 8; + const po = try r.alloc(@as(usize, max_rows) * NQ * maxc * HD * 4); + const pm = try r.alloc(@as(usize, max_rows) * NQ * maxc * 4); + const pl = try r.alloc(@as(usize, max_rows) * NQ * maxc * 4); + std.debug.print("GQA attention, 16 q / 4 kv heads, chunks of {d} keys, sigmoid gate per head, against FP64\n", .{CHUNK}); + + const cases = [_]Case{ + .{ .what = "full layer, 6 rows at 700 (12 chunks)", .pos0 = 700, .n = 6, .window = 0, .ring = 4096 }, + .{ .what = "full layer, decode row at 0 (one key)", .pos0 = 0, .n = 1, .window = 0, .ring = 4096 }, + .{ .what = "sliding 256, 8 rows at 900 in a 512-slot ring (wraps)", .pos0 = 900, .n = 8, .window = 256, .ring = 512 }, + .{ .what = "sliding 256, 3 rows at 100 (shorter than the window)", .pos0 = 100, .n = 3, .window = 256, .ring = 512 }, + }; + for (cases, 0..) |c, ci| { + const npos: usize = c.pos0 + c.n; + const kh = try o.gpa.alloc(u16, npos * NKV * HD); + const vh = try o.gpa.alloc(u16, kh.len); + o.randBf(kh, rnd, 1.0); + o.randBf(vh, rnd, 1.0); + const qh = try o.gpa.alloc(u16, c.n * NQ * HD); + o.randBf(qh, rnd, 1.0); + const gh = try o.gpa.alloc(u16, c.n * NQ); + o.randBf(gh, rnd, 2.0); + const planted = c.window != 0 and c.pos0 >= c.window; + if (planted) { // a key just outside each row's window, aligned with its query: any row that reads it is dominated by it + for (0..c.n) |z| { + const key = c.pos0 + z - c.window; + for (0..NKV) |kvh| for (0..HD) |d| { + kh[(key * NKV + kvh) * HD + d] = o.toBf(3.0 * o.bf(qh[(z * NQ + kvh * (NQ / NKV)) * HD + d])); + vh[(key * NKV + kvh) * HD + d] = o.toBf(8.0); + }; + } + } + // the device ring holds position p at slot p % ring; only the positions a row may read must survive + const ring_k = try o.gpa.alloc(u16, @as(usize, c.ring) * NKV * HD); + const ring_v = try o.gpa.alloc(u16, ring_k.len); + @memset(ring_k, 0); + @memset(ring_v, 0); + const first = if (npos > c.ring) npos - c.ring else 0; + for (first..npos) |p| { + const slot = p % c.ring; + @memcpy(ring_k[slot * NKV * HD ..][0 .. NKV * HD], kh[p * NKV * HD ..][0 .. NKV * HD]); + @memcpy(ring_v[slot * NKV * HD ..][0 .. NKV * HD], vh[p * NKV * HD ..][0 .. NKV * HD]); + } + const kc = try o.upload(&r, ring_k); + const vc = try o.upload(&r, ring_v); + const q = try o.upload(&r, qh); + const g = try o.upload(&r, gh); + const out = try r.alloc(c.n * NQ * HD * 2); + const span = if (c.window == 0) c.pos0 + c.n else @min(c.pos0 + c.n, c.window); + const nch = (span + CHUNK - 1) / CHUNK; + try o.run(&kp, .{ NKV, nch, c.n }, .{ q, kc, vc, po, pm, pl, c.pos0, CHUNK, c.ring, c.window, NKV, maxc, @as(f32, 1.0 / @sqrt(128.0)) }); + try o.run(&km, .{ NQ, c.n, 1 }, .{ po, pm, pl, g, out, c.pos0, CHUNK, c.window, NQ, maxc }); + const oh = try o.fetch(&r, out, c.n * NQ * HD); + const qf = try o.widen(o.gpa, qh); + const kf = try o.widen(o.gpa, kh); + const vf = try o.widen(o.gpa, vh); + const gf = try o.widen(o.gpa, gh); + const want = try o.gpa.alloc(f64, oh.len); + try refRows(qf, kf, vf, gf, c.n, c.pos0, c.window, want); + t.check(c.what, o.relErr(oh, want)); + if (ci == 2) { // planted control: a window of 257 reads the planted key and must not match + try refRows(qf, kf, vf, gf, c.n, c.pos0, c.window + 1, want); + t.control("sliding layer replayed with a window of 257", o.relErr(oh, want)); + } + } + try t.finish(); +} diff --git a/zig/tests/xpu/lfm3_conv.zig b/zig/tests/xpu/lfm3_conv.zig new file mode 100644 index 0000000000..bbfe076e9f --- /dev/null +++ b/zig/tests/xpu/lfm3_conv.zig @@ -0,0 +1,92 @@ +//! LFM3 gated short conv (lfm3_conv_rows) against an FP64 replay on random data: a prompt window, then decode steps. + +const std = @import("std"); +const o = @import("lfm3_ops.zig"); +const stop = @import("xpu").stop; + +const D: u32 = 2048; +const PD: u32 = 3 * D; + +/// FP64 replay of rows of [B | C | x]: y = C * sum_k hist[k] * w[k] over the last 7 B * x; taps reversed when `flip`. +fn refConv(proj: []const f64, w: []const f64, hist: []f64, n: usize, flip: bool, y: []f64) void { + for (0..D) |c| { + for (0..n) |r| { + const p = proj[r * PD ..]; + const bx = p[c] * p[2 * D + c]; + var acc: f64 = 0; + for (0..6) |k| acc += hist[k * D + c] * w[c * 7 + (if (flip) 6 - k else k)]; + acc += bx * w[c * 7 + @as(usize, if (flip) 0 else 6)]; + y[r * D + c] = p[D + c] * acc; + for (0..5) |k| hist[k * D + c] = hist[(k + 1) * D + c]; + hist[5 * D + c] = bx; + } + } +} + +pub const panic = std.debug.FullPanic(@import("rig.zig").panicDrain); +pub fn main() !void { + stop.install(); + var r = try o.rt.open(); + defer r.deinit(); + var m = try r.module(@import("xpu").kernels.lfm3); + var k = try m.kernel("lfm3_conv_rows", .{ 64, 1, 1 }); + var t: o.Tally = .{ .name = "lfm3_conv" }; + var prng = std.Random.DefaultPrng.init(11); + const rnd = prng.random(); + + const n_pre: u32 = 37; + const n_dec: u32 = 5; + const rows = n_pre + n_dec; + const ph = try o.gpa.alloc(u16, rows * PD); + o.randBf(ph, rnd, 1.0); + const wh = try o.gpa.alloc(u16, D * 7); + o.randBf(wh, rnd, 0.3); + const proj = try o.upload(&r, ph); + const cw = try o.upload(&r, wh); + const zero = try o.gpa.alloc(u16, 6 * D); + @memset(zero, 0); + const st = try o.upload(&r, zero); + const y = try r.alloc(rows * D * 2); + + // the prompt window from an empty state, then single decode steps on the state it leaves (state in place) + try o.run(&k, .{ D / 64, 1, 1 }, .{ proj, PD, st, st, cw, y, n_pre, D }); + for (0..n_dec) |i| { + const off = (n_pre + i) * PD * 2; + const pr: o.Buf = .{ .rt = proj.rt, .ptr = @ptrFromInt(@intFromPtr(proj.ptr.?) + off), .len = proj.len - off }; + const yo = (n_pre + i) * D * 2; + const yr: o.Buf = .{ .rt = y.rt, .ptr = @ptrFromInt(@intFromPtr(y.ptr.?) + yo), .len = y.len - yo }; + try o.run(&k, .{ D / 64, 1, 1 }, .{ pr, PD, st, st, cw, yr, @as(u32, 1), D }); + } + const yh = try o.fetch(&r, y, rows * D); + const sh = try o.fetch(&r, st, 6 * D); + + const pf = try o.widen(o.gpa, ph); + const wf = try o.widen(o.gpa, wh); + const hist = try o.gpa.alloc(f64, 6 * D); + const yf = try o.gpa.alloc(f64, rows * D); + @memset(hist, 0); + refConv(pf, wf, hist, rows, false, yf); + std.debug.print("gated short conv, D {d}: {d} prompt rows then {d} decode steps, against FP64\n", .{ D, n_pre, n_dec }); + t.check("prompt rows y", o.relErr(yh[0 .. n_pre * D], yf[0 .. n_pre * D])); + t.check("decode steps y", o.relErr(yh[n_pre * D ..], yf[n_pre * D ..])); + t.check("state after the last step (last six B*x)", o.relErr(sh, hist)); + + // the same rows one at a time: the state handoff gives the window's bits + const st1 = try o.upload(&r, zero); + const y1 = try r.alloc(rows * D * 2); + for (0..rows) |i| { + const off = i * PD * 2; + const pr: o.Buf = .{ .rt = proj.rt, .ptr = @ptrFromInt(@intFromPtr(proj.ptr.?) + off), .len = proj.len - off }; + const yo = i * D * 2; + const yr: o.Buf = .{ .rt = y1.rt, .ptr = @ptrFromInt(@intFromPtr(y1.ptr.?) + yo), .len = y1.len - yo }; + try o.run(&k, .{ D / 64, 1, 1 }, .{ pr, PD, st1, st1, cw, yr, @as(u32, 1), D }); + } + const y1h = try o.fetch(&r, y1, rows * D); + t.same("all rows stepped one at a time against the window", std.mem.eql(u16, y1h, yh)); + + // planted control: the taps in reverse order (oldest input on the newest tap) must not match + @memset(hist, 0); + refConv(pf, wf, hist, rows, true, yf); + t.control("reversed taps", o.relErr(yh, yf)); + try t.finish(); +} diff --git a/zig/tests/xpu/lfm3_norm.zig b/zig/tests/xpu/lfm3_norm.zig new file mode 100644 index 0000000000..460412ab92 --- /dev/null +++ b/zig/tests/xpu/lfm3_norm.zig @@ -0,0 +1,112 @@ +//! LFM3 row ops (lfm3_rmsnorm, the sandwich step lfm3_post_add_norm, lfm3_swiglu) against FP64 on random data. + +const std = @import("std"); +const o = @import("lfm3_ops.zig"); +const stop = @import("xpu").stop; + +const N: u32 = 2048; +const ROWS: u32 = 5; +const FF: u32 = 7680; +const eps: f64 = 1e-5; + +fn rms(x: []const f64, w: []const f64, out: []f64) void { + var ss: f64 = 0; + for (x) |v| ss += v * v; + const rs = 1.0 / @sqrt(ss / @as(f64, @floatFromInt(x.len)) + eps); + for (x, w, out) |v, wv, *y| y.* = wv * v * rs; +} + +pub const panic = std.debug.FullPanic(@import("rig.zig").panicDrain); +pub fn main() !void { + stop.install(); + var r = try o.rt.open(); + defer r.deinit(); + var m = try r.module(@import("xpu").kernels.lfm3); + var k_rms = try m.kernel("lfm3_rmsnorm", .{ 256, 1, 1 }); + var k_pan = try m.kernel("lfm3_post_add_norm", .{ 256, 1, 1 }); + var k_swi = try m.kernel("lfm3_swiglu", .{ 64, 1, 1 }); + var t: o.Tally = .{ .name = "lfm3_norm" }; + var prng = std.Random.DefaultPrng.init(14); + const rnd = prng.random(); + std.debug.print("row ops over {d} rows of {d}, against FP64\n", .{ ROWS, N }); + + const xh = try o.gpa.alloc(u16, ROWS * N); + const dh = try o.gpa.alloc(u16, ROWS * N); + const wph = try o.gpa.alloc(u16, N); + const wnh = try o.gpa.alloc(u16, N); + o.randBf(xh, rnd, 3.0); + o.randBf(dh, rnd, 0.5); + for (wph) |*v| v.* = o.toBf(0.2 + 0.1 * rnd.floatNorm(f64)); + for (wnh) |*v| v.* = o.toBf(1.0 + 0.3 * rnd.floatNorm(f64)); + const xf = try o.widen(o.gpa, xh); + const df = try o.widen(o.gpa, dh); + const wpf = try o.widen(o.gpa, wph); + const wnf = try o.widen(o.gpa, wnh); + + // RMSNorm (the final embedding_norm and every pre-norm) + { + const x = try o.upload(&r, xh); + const w = try o.upload(&r, wnh); + const y = try r.alloc(ROWS * N * 2); + try o.run(&k_rms, .{ ROWS, 1, 1 }, .{ x, w, y, N, @as(f32, 1e-5) }); + const yh = try o.fetch(&r, y, ROWS * N); + const want = try o.gpa.alloc(f64, ROWS * N); + for (0..ROWS) |z| rms(xf[z * N ..][0..N], wnf, want[z * N ..][0..N]); + t.check("rmsnorm", o.relErr(yh, want)); + } + + // sandwich step: x += post_norm(d); y = next_norm(x) + { + const x = try o.upload(&r, xh); + const d = try o.upload(&r, dh); + const wp = try o.upload(&r, wph); + const wn = try o.upload(&r, wnh); + const y = try r.alloc(ROWS * N * 2); + try o.run(&k_pan, .{ ROWS, 1, 1 }, .{ d, x, wp, wn, y, N, @as(f32, 1e-5) }); + const xo = try o.fetch(&r, x, ROWS * N); + const yh = try o.fetch(&r, y, ROWS * N); + const hw = try o.gpa.alloc(f64, ROWS * N); + const yw = try o.gpa.alloc(f64, ROWS * N); + const tmp = try o.gpa.alloc(f64, N); + for (0..ROWS) |z| { + rms(df[z * N ..][0..N], wpf, tmp); + for (0..N) |i| hw[z * N + i] = xf[z * N + i] + tmp[i]; + rms(hw[z * N ..][0..N], wnf, yw[z * N ..][0..N]); + } + t.check("post_add_norm residual h = x + post_norm(d)", o.relErr(xo, hw)); + t.check("post_add_norm output next_norm(h)", o.relErr(yh, yw)); + // planted control: the residual added before the post norm (h = post_norm(x + d)) must not match + for (0..ROWS) |z| { + for (0..N) |i| tmp[i] = xf[z * N + i] + df[z * N + i]; + rms(tmp, wpf, hw[z * N ..][0..N]); + } + t.control("residual added before the post norm", o.relErr(xo, hw)); + } + + // SwiGLU: silu(w1 x) * (w3 x) + { + const n = 3 * FF; + const ah = try o.gpa.alloc(u16, n); + const bh = try o.gpa.alloc(u16, n); + o.randBf(ah, rnd, 2.0); + o.randBf(bh, rnd, 1.0); + const a = try o.upload(&r, ah); + const b = try o.upload(&r, bh); + const y = try r.alloc(n * 2); + try o.run(&k_swi, .{ (n + 63) / 64, 1, 1 }, .{ a, b, y, n }); + const yh = try o.fetch(&r, y, n); + const want = try o.gpa.alloc(f64, n); + for (0..n) |i| { + const x = o.bf(ah[i]); + want[i] = x / (1.0 + @exp(-x)) * o.bf(bh[i]); + } + t.check("swiglu", o.relErr(yh, want)); + // planted control: gate and up swapped (silu on w3 x) must not match + for (0..n) |i| { + const x = o.bf(bh[i]); + want[i] = x / (1.0 + @exp(-x)) * o.bf(ah[i]); + } + t.control("gate and up swapped", o.relErr(yh, want)); + } + try t.finish(); +} diff --git a/zig/tests/xpu/lfm3_ops.zig b/zig/tests/xpu/lfm3_ops.zig new file mode 100644 index 0000000000..993557ff33 --- /dev/null +++ b/zig/tests/xpu/lfm3_ops.zig @@ -0,0 +1,95 @@ +//! Shared pieces of the LFM3 op tests: bf16 data, launches, the FP64 comparison and its control. + +const std = @import("std"); +pub const rt = @import("rig.zig"); + +pub const gpa = std.heap.page_allocator; +pub const Buf = rt.Buffer; +/// Largest |device - FP64| over the largest |FP64| a bf16 output may show; a control must land above it. +pub const tol: f64 = 1e-2; + +pub fn bf(v: u16) f64 { + return @as(f32, @bitCast(@as(u32, v) << 16)); +} + +/// Round to nearest even, as the kernels' to_bf. +pub fn toBf(f: f64) u16 { + var u: u32 = @bitCast(@as(f32, @floatCast(f))); + u +%= 0x7fff + ((u >> 16) & 1); + return @intCast(u >> 16); +} + +pub fn randBf(buf: []u16, rnd: std.Random, sd: f64) void { + for (buf) |*v| v.* = toBf(rnd.floatNorm(f64) * sd); +} + +pub fn widen(gpa_: std.mem.Allocator, h: []const u16) ![]f64 { + const out = try gpa_.alloc(f64, h.len); + for (h, out) |v, *o| o.* = bf(v); + return out; +} + +pub fn run(k: *rt.Kernel, groups: [3]u32, a: anytype) !void { + inline for (a, 0..) |v, i| { + const T = @TypeOf(v); + if (T == Buf) { + try k.setBuffer(i, v); + } else if (T == f32 or T == comptime_float) { + try k.setF32(i, v); + } else try k.setU32(i, @as(u32, v)); + } + try k.launch(groups); +} + +pub fn upload(r: *rt.Runtime, h: []const u16) !Buf { + const b = try r.alloc(h.len * 2); + try r.upload(b, std.mem.sliceAsBytes(h)); + return b; +} + +pub fn fetch(r: *rt.Runtime, b: Buf, n: usize) ![]u16 { + const h = try gpa.alloc(u16, n); + try r.sync(); + try r.download(std.mem.sliceAsBytes(h), b); + try r.sync(); + return h; +} + +/// max|got - want| / max|want|. +pub fn relErr(got: []const u16, want: []const f64) f64 { + var e: f64 = 0; + var m: f64 = 0; + for (got, want) |g, w| { + e = @max(e, @abs(bf(g) - w)); + m = @max(m, @abs(w)); + } + return if (m == 0) e else e / m; +} + +/// Tallies the checks of one test program: real checks must stay within tol, planted controls must exceed it. +pub const Tally = struct { + name: []const u8, + bad: bool = false, + + pub fn check(t: *Tally, what: []const u8, err: f64) void { + const ok = err <= tol; + if (!ok) t.bad = true; + std.debug.print(" {s}: rel err {e:.2} (tol {e:.0}) {s}\n", .{ what, err, tol, if (ok) "pass" else "FAIL" }); + } + + pub fn control(t: *Tally, what: []const u8, err: f64) void { + const fails = err > tol; + if (!fails) t.bad = true; + std.debug.print(" control {s}: rel err {e:.2} {s}\n", .{ what, err, if (fails) "fails as planted" else "PASSES (control broken)" }); + } + + pub fn same(t: *Tally, what: []const u8, eq: bool) void { + if (!eq) t.bad = true; + std.debug.print(" {s}: {s}\n", .{ what, if (eq) "bitwise identical" else "DIFFERENT" }); + } + + pub fn finish(t: *Tally) !void { + std.debug.print("{s}: {s}\n", .{ t.name, if (t.bad) "FAILED" else "ok" }); + if (t.bad) return error.TestFailed; + } +}; diff --git a/zig/tests/xpu/lfm3_rope.zig b/zig/tests/xpu/lfm3_rope.zig new file mode 100644 index 0000000000..d5417b321b --- /dev/null +++ b/zig/tests/xpu/lfm3_rope.zig @@ -0,0 +1,108 @@ +//! LFM3 per-head QK RMSNorm with zero-tail partial RoPE, and the V store, against FP64 on random data. + +const std = @import("std"); +const o = @import("lfm3_ops.zig"); +const stop = @import("xpu").stop; + +const HD = 128; +const eps: f64 = 1e-5; + +/// Inverse frequencies as upstream builds them in fp32: base^(-2i/128) for i < angles, then zeros. +fn invFreq(base: f32, angles: usize, out: []f32) void { + for (out, 0..) |*f, i| f.* = if (i < angles) 1.0 / std.math.pow(f32, base, @as(f32, @floatFromInt(2 * i)) / 128.0) else 0; +} + +/// FP64 replay of one row's heads: RMSNorm over 128 dims, then rotate_half RoPE at `pos` with `angles` rotated pairs. +fn refRow(src: []const f64, w: []const f64, heads: usize, base: f64, angles: usize, pos: f64, out: []f64) void { + for (0..heads) |h| { + const s = src[h * HD ..][0..HD]; + var ss: f64 = 0; + for (s) |v| ss += v * v; + const rs = 1.0 / @sqrt(ss / HD + eps); + for (0..64) |i| { + const a = w[i] * s[i] * rs; + const b = w[i + 64] * s[i + 64] * rs; + const f = if (i < angles) pos / std.math.pow(f64, base, @as(f64, @floatFromInt(2 * i)) / 128.0) else 0; + out[h * HD + i] = a * @cos(f) - b * @sin(f); + out[h * HD + i + 64] = b * @cos(f) + a * @sin(f); + } + } +} + +const Case = struct { what: []const u8, base: f32, angles: usize, heads: u32, pos0: u32, ring: u32 }; + +pub const panic = std.debug.FullPanic(@import("rig.zig").panicDrain); +pub fn main() !void { + stop.install(); + var r = try o.rt.open(); + defer r.deinit(); + var m = try r.module(@import("xpu").kernels.lfm3); + var kr = try m.kernel("lfm3_qk_norm_rope", .{ 64, 1, 1 }); + var kv = try m.kernel("lfm3_v_store", .{ 64, 1, 1 }); + var t: o.Tally = .{ .name = "lfm3_rope" }; + var prng = std.Random.DefaultPrng.init(12); + const rnd = prng.random(); + + const n: u32 = 9; + const stride: u32 = 2048 + 2 * 512; // a fused q | k | v projection row + const sh = try o.gpa.alloc(u16, n * stride); + o.randBf(sh, rnd, 1.5); + const src = try o.upload(&r, sh); + const sf = try o.widen(o.gpa, sh); + const wh = try o.gpa.alloc(u16, HD); + for (wh) |*v| v.* = o.toBf(1.0 + 0.3 * rnd.floatNorm(f64)); + const w = try o.upload(&r, wh); + const wf = try o.widen(o.gpa, wh); + const inv = try r.alloc(64 * 4); + var invh: [64]f32 = undefined; + std.debug.print("QK RMSNorm + RoPE over 128-dim heads, {d} rows, against FP64\n", .{n}); + + const cases = [_]Case{ + .{ .what = "q, full layer (theta 1e6, 32 of 64 pairs)", .base = 1e6, .angles = 32, .heads = 16, .pos0 = 7000, .ring = 0 }, + .{ .what = "q, sliding layer (theta 1e4, all 64 pairs)", .base = 1e4, .angles = 64, .heads = 16, .pos0 = 5, .ring = 0 }, + .{ .what = "k into a 256-slot ring across its wrap", .base = 1e4, .angles = 64, .heads = 4, .pos0 = 250, .ring = 256 }, + }; + var full_out: []u16 = &.{}; + for (cases) |c| { + invFreq(c.base, c.angles, &invh); + try r.upload(inv, std.mem.sliceAsBytes(&invh)); + const slots: u32 = if (c.ring == 0) n else c.ring; + const dst_stride: u32 = c.heads * HD; + const dst = try r.alloc(slots * dst_stride * 2); + const src_off: u32 = if (c.heads == 16) 0 else 2048; + try o.run(&kr, .{ c.heads, n, 1 }, .{ src, stride, src_off, dst, dst_stride, w, inv, c.pos0, c.ring, @as(f32, 1e-5) }); + const dh = try o.fetch(&r, dst, slots * dst_stride); + const got = try o.gpa.alloc(u16, n * dst_stride); + const want = try o.gpa.alloc(f64, n * dst_stride); + for (0..n) |z| { + const slot = if (c.ring == 0) z else (c.pos0 + z) % c.ring; + @memcpy(got[z * dst_stride ..][0..dst_stride], dh[slot * dst_stride ..][0..dst_stride]); + refRow(sf[z * stride + src_off ..], wf, c.heads, c.base, c.angles, @floatFromInt(c.pos0 + z), want[z * dst_stride ..]); + } + t.check(c.what, o.relErr(got, want)); + if (c.angles == 32) full_out = got; + } + + // V rows into the ring: a plain copy + { + const ring: u32 = 256; + const pos0: u32 = 250; + const dst = try r.alloc(ring * 512 * 2); + try o.run(&kv, .{ 512 / 64, n, 1 }, .{ src, stride, @as(u32, 2048 + 512), dst, @as(u32, 512), pos0, ring }); + const dh = try o.fetch(&r, dst, ring * 512); + var eq = true; + for (0..n) |z| { + const slot = (pos0 + z) % ring; + if (!std.mem.eql(u16, dh[slot * 512 ..][0..512], sh[z * stride + 2048 + 512 ..][0..512])) eq = false; + } + t.same("v rows stored at slot (pos0 + r) % 256", eq); + } + + // planted control: the full layer's q replayed with every pair rotated (partial rotary ignored) must not match + { + const want = try o.gpa.alloc(f64, n * 16 * HD); + for (0..n) |z| refRow(sf[z * stride ..], wf, 16, 1e6, 64, @floatFromInt(7000 + z), want[z * 16 * HD ..]); + t.control("full layer with all 64 pairs rotated", o.relErr(full_out, want)); + } + try t.finish(); +} From 4de7cf2842d8d8ca9b622b483a2e1fe5a14fbad0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 30 Oct 2026 15:57:17 +0100 Subject: [PATCH 06/14] xpu: LFM3 forward on the Intel GPU (prefill in windows, decode, conv state and KV rings) loading by names.zig, its checkpoint test against ref64.py's float64 dumps --- zig/build/xpu.zig | 5 +- zig/src/families/lfm3/xpu.zig | 3 + zig/src/families/lfm3/xpu_model.zig | 323 ++++++++++++++++++++++++++++ zig/tests/xpu/lfm3_fwd.zig | 163 ++++++++++++++ 4 files changed, 493 insertions(+), 1 deletion(-) create mode 100644 zig/src/families/lfm3/xpu.zig create mode 100644 zig/src/families/lfm3/xpu_model.zig create mode 100644 zig/tests/xpu/lfm3_fwd.zig diff --git a/zig/build/xpu.zig b/zig/build/xpu.zig index c0d16b2e49..885eff13c5 100644 --- a/zig/build/xpu.zig +++ b/zig/build/xpu.zig @@ -258,7 +258,7 @@ fn hasDpas(device: []const u8) bool { /// The standalone op and model tests, each its own program `tf-xpu--test` (step `xpu-tests`). const test_programs = [_][]const u8{ - "panic", "nem_rows", "nem_prefill", "nem_gen", "nem_mtp", "nem_kl", "nem_norm_bench", "nem_mv_bench", "nem_slot", "nem_format", "nem_seam", "nem_attn_dec", "nem_fp64", "memory", "lfm3_conv", "lfm3_rope", "lfm3_attn", "lfm3_norm", + "panic", "nem_rows", "nem_prefill", "nem_gen", "nem_mtp", "nem_kl", "nem_norm_bench", "nem_mv_bench", "nem_slot", "nem_format", "nem_seam", "nem_attn_dec", "nem_fp64", "memory", "lfm3_conv", "lfm3_rope", "lfm3_attn", "lfm3_norm", "lfm3_fwd", }; /// `-Dxpu` (Linux): Nemotron's SPIR-V images, `tensorfold-xpu` and `tf-xpu-test`; needs ocloc, not nvcc or CUDA. @@ -275,11 +275,14 @@ fn nemotronSteps(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std. runner.addImport("nemotron_xpu", mods.nemotron); b.installArtifact(b.addExecutable(.{ .name = "tf-xpu-test", .root_module = runner })); const tests_step = b.step("xpu-tests", "Build the standalone Intel GPU op and model tests (tf-xpu--test)"); + const lfm3 = b.createModule(.{ .root_source_file = b.path("zig/src/families/lfm3/xpu.zig"), .target = target, .optimize = optimize, .link_libc = true }); + lfm3.addImport("xpu", xpu); for (test_programs) |name| { const mod = b.createModule(.{ .root_source_file = b.path(b.fmt("zig/tests/xpu/{s}.zig", .{name})), .target = target, .optimize = optimize, .link_libc = true }); mod.addImport("xpu", xpu); mod.addImport("core", mods.core); mod.addImport("nemotron_xpu", mods.nemotron); + mod.addImport("lfm3_xpu", lfm3); tests_step.dependOn(&b.addInstallArtifact(b.addExecutable(.{ .name = b.fmt("tf-xpu-{s}-test", .{name}), .root_module = mod }), .{}).step); } } diff --git a/zig/src/families/lfm3/xpu.zig b/zig/src/families/lfm3/xpu.zig new file mode 100644 index 0000000000..c10d895352 --- /dev/null +++ b/zig/src/families/lfm3/xpu.zig @@ -0,0 +1,3 @@ +//! LFM3 on the Intel GPU (XPU backend): the forward and the config it reads. +pub const model = @import("xpu_model.zig"); +pub const config = @import("config.zig"); diff --git a/zig/src/families/lfm3/xpu_model.zig b/zig/src/families/lfm3/xpu_model.zig new file mode 100644 index 0000000000..be204ac057 --- /dev/null +++ b/zig/src/families/lfm3/xpu_model.zig @@ -0,0 +1,323 @@ +//! LFM3 forward on the Intel GPU: MLX 4-bit weights on the device, a conv state and KV rings, windows of rows, no DPAS. + +const std = @import("std"); +const xpu = @import("xpu"); +const cfg = @import("config.zig"); +const names = @import("names.zig"); +const rt = xpu.rt; +const ld = xpu.loader; + +const Buf = rt.Buffer; +const H: u32 = cfg.hidden; +const FF: u32 = cfg.intermediate; +const V: u32 = cfg.vocab; +const NQ: u32 = cfg.query_heads; +const NKV: u32 = cfg.kv_heads; +const HD: u32 = cfg.head_dim; +const QKV: u32 = (NQ + 2 * NKV) * HD; +const chunk: u32 = 128; // keys a split-K work-group scans +const argmax_parts: u32 = 64; + +pub const Table = struct { w: Buf, s: Buf, b: Buf }; +const Conv = struct { in: Table, out: Table, cw: Buf, state: Buf }; +const Attn = struct { q: Table, k: Table, v: Table, g: Table, o: Table, qn: Buf, kn: Buf, kc: Buf, vc: Buf, ring: u32, window: u32 }; +const Mixer = union(enum) { conv: Conv, attn: Attn }; +const Layer = struct { op_norm: Buf, post_op: Buf, ffn_norm: Buf, post_ffn: Buf, w1: Table, w3: Table, w2: Table, mixer: Mixer }; + +const Kernels = struct { + embed: rt.Kernel, + qmv: rt.Kernel, + head: rt.Kernel, + am_part: rt.Kernel, + am_fin: rt.Kernel, + pdense: rt.Kernel, + rms: rt.Kernel, + pan: rt.Kernel, + conv: rt.Kernel, + rope: rt.Kernel, + vst: rt.Kernel, + part: rt.Kernel, + merge: rt.Kernel, + swiglu: rt.Kernel, +}; + +/// Device scratch for a window of up to `max_rows` rows. +const Scratch = struct { ids: Buf, x: Buf, xn: Buf, d: Buf, proj: Buf, qkv: Buf, g: Buf, q: Buf, att: Buf, a: Buf, b: Buf, act: Buf, po: Buf, pm: Buf, pl: Buf, logits: Buf, am_scratch: Buf, am_out: Buf }; + +/// Called after each layer (index) and after the final norm (-1) with the window's rows, bf16 [n][2048]. +pub const Hook = struct { ctx: *anyopaque, f: *const fn (ctx: *anyopaque, layer: i32, rows: []const u16) void }; + +fn run(k: *rt.Kernel, groups: [3]u32, a: anytype) !void { + inline for (a, 0..) |v, i| { + const T = @TypeOf(v); + if (T == Buf) { + try k.setBuffer(i, v); + } else if (T == f32) { + try k.setF32(i, v); + } else try k.setU32(i, @as(u32, v)); + } + try k.launch(groups); +} + +/// Inverse RoPE frequencies as upstream builds them (fp32): theta^(-2i/128) for i < angles, then zeros (zero-tail). +pub fn invFreq(theta: f32, partial: f32, out: *[HD / 2]f32) void { + const angles: usize = @intFromFloat(partial * @as(f32, HD) / 2); + for (out, 0..) |*f, i| f.* = if (i < angles) 1.0 / std.math.pow(f32, theta, @as(f32, @floatFromInt(2 * i)) / @as(f32, HD)) else 0; +} + +pub const Model = struct { + r: *rt.Runtime, + layers: []Layer, + emb: Table, + norm_f: Buf, + inv_full: Buf, + inv_slide: Buf, + k: Kernels, + s: Scratch, + host_rows: []u16, + cap: u32, + max_rows: u32, + maxc: u32, + pos: u32 = 0, + device_bytes: u64 = 0, + + /// The affine tensors of module `m` (in layer `layer`, or a model tensor) by names.zig in namespace `ns`. + fn table(l: *ld.Loader, ns: names.Namespace, layer: ?usize, m: names.Module) !Table { + var a: [160]u8 = undefined; + var b: [176]u8 = undefined; + const base = try names.moduleName(&a, ns, layer, m); + return .{ + .w = try l.load(try std.fmt.bufPrint(&b, "{s}.weight", .{base})), + .s = try l.load(try std.fmt.bufPrint(&b, "{s}.scales", .{base})), + .b = try l.load(try std.fmt.bufPrint(&b, "{s}.biases", .{base})), + }; + } + + fn vec(l: *ld.Loader, ns: names.Namespace, layer: ?usize, m: names.Module) !Buf { + var a: [160]u8 = undefined; + return l.load(try names.moduleName(&a, ns, layer, m)); + } + + /// The checkpoint's namespace (bare `model.` or wrapped `language_model.model.`) from the embedding it holds. + fn namespace(l: *const ld.Loader) !names.Namespace { + var a: [64]u8 = undefined; + var b: [64]u8 = undefined; + return names.detect(l.map.contains(try names.embedName(&a, .bare)), l.map.contains(try names.embedName(&b, .wrapped))); + } + + /// Uploads the checkpoint: `cap` positions on full-attention layers, windows of at most `max_rows` rows. + pub fn load(gpa: std.mem.Allocator, r: *rt.Runtime, l: *ld.Loader, c: *const cfg.Config, cap: u32, max_rows: u32) !Model { + if (max_rows == 0 or max_rows > 256 or cap < max_rows) return error.Invalid; + var bm = try r.module(xpu.kernels.basic); + var am = try r.module(xpu.kernels.attn); + var pm = try r.module(xpu.kernels.nem_prow); + var fm = try r.module(xpu.kernels.lfm3); + const k: Kernels = .{ + .embed = try bm.kernel("embed4", .{ 64, 1, 1 }), + .qmv = try am.kernel("qmv4_bf", .{ 64, 1, 1 }), + .head = try am.kernel("qmv4_f32", .{ 64, 1, 1 }), + .am_part = try am.kernel("argmax_partial", .{ 256, 1, 1 }), + .am_fin = try am.kernel("argmax_final", .{ 256, 1, 1 }), + .pdense = try pm.kernel("pdense", .{ 32, 1, 1 }), + .rms = try fm.kernel("lfm3_rmsnorm", .{ 256, 1, 1 }), + .pan = try fm.kernel("lfm3_post_add_norm", .{ 256, 1, 1 }), + .conv = try fm.kernel("lfm3_conv_rows", .{ 64, 1, 1 }), + .rope = try fm.kernel("lfm3_qk_norm_rope", .{ 64, 1, 1 }), + .vst = try fm.kernel("lfm3_v_store", .{ 64, 1, 1 }), + .part = try fm.kernel("lfm3_attn_partial", .{ 64, 1, 1 }), + .merge = try fm.kernel("lfm3_attn_merge_gate", .{ 16, 1, 1 }), + .swiglu = try fm.kernel("lfm3_swiglu", .{ 64, 1, 1 }), + }; + const maxc = (cap + chunk - 1) / chunk; + const w: usize = max_rows; + const s: Scratch = .{ + .ids = try l.empty(w * 4), + .x = try l.empty(w * H * 2), + .xn = try l.empty(w * H * 2), + .d = try l.empty(w * H * 2), + .proj = try l.empty(w * 3 * H * 2), + .qkv = try l.empty(w * QKV * 2), + .g = try l.empty(w * NQ * 2), + .q = try l.empty(w * NQ * HD * 2), + .att = try l.empty(w * NQ * HD * 2), + .a = try l.empty(w * FF * 2), + .b = try l.empty(w * FF * 2), + .act = try l.empty(w * FF * 2), + .po = try l.empty(w * NQ * maxc * HD * 4), + .pm = try l.empty(w * NQ * maxc * 4), + .pl = try l.empty(w * NQ * maxc * 4), + .logits = try l.empty(@as(usize, V) * 4), + .am_scratch = try l.empty(argmax_parts * 8), + .am_out = try l.empty(4), + }; + const ns = try namespace(l); + const layers = try gpa.alloc(Layer, cfg.layers); + for (layers, 0..) |*ly, i| { + const kind = c.layer_types[i]; + ly.* = .{ + .op_norm = try vec(l, ns, i, .operator_norm), + .post_op = try vec(l, ns, i, .post_operator_norm), + .ffn_norm = try vec(l, ns, i, .ffn_norm), + .post_ffn = try vec(l, ns, i, .post_ffn_norm), + .w1 = try table(l, ns, i, .w1), + .w3 = try table(l, ns, i, .w3), + .w2 = try table(l, ns, i, .w2), + .mixer = if (kind == .conv) .{ .conv = .{ + .in = try table(l, ns, i, .conv_in), + .out = try table(l, ns, i, .conv_out), + .cw = try vec(l, ns, i, .conv), + .state = try l.zeros(6 * H * 2), + } } else blk: { + const ring: u32 = if (kind == .sliding_attention) cfg.sliding_window + max_rows else cap; + break :blk .{ .attn = .{ + .q = try table(l, ns, i, .q_proj), + .k = try table(l, ns, i, .k_proj), + .v = try table(l, ns, i, .v_proj), + .g = try table(l, ns, i, .gate_proj), + .o = try table(l, ns, i, .o_proj), + .qn = try vec(l, ns, i, .q_norm), + .kn = try vec(l, ns, i, .k_norm), + .kc = try l.empty(@as(usize, ring) * NKV * HD * 2), + .vc = try l.empty(@as(usize, ring) * NKV * HD * 2), + .ring = ring, + .window = if (kind == .sliding_attention) cfg.sliding_window else 0, + } }; + }, + }; + } + var inv: [HD / 2]f32 = undefined; + const inv_full = try l.empty(HD / 2 * 4); + invFreq(c.full_rotary.theta, c.full_rotary.partial, &inv); + try r.upload(inv_full, std.mem.sliceAsBytes(&inv)); + const inv_slide = try l.empty(HD / 2 * 4); + invFreq(c.sliding_rotary.theta, c.sliding_rotary.partial, &inv); + try r.upload(inv_slide, std.mem.sliceAsBytes(&inv)); + try r.sync(); + return .{ + .r = r, + .layers = layers, + .emb = try table(l, ns, null, .embed_tokens), + .norm_f = try vec(l, ns, null, .embedding_norm), + .inv_full = inv_full, + .inv_slide = inv_slide, + .k = k, + .s = s, + .host_rows = try gpa.alloc(u16, w * H), + .cap = cap, + .max_rows = max_rows, + .maxc = maxc, + .device_bytes = l.total, + }; + } + + /// Back to position 0: the conv states are zeroed; KV slots are only read once written. + pub fn reset(m: *Model) !void { + const zero: [6 * 256]u16 = @splat(0); + for (m.layers) |ly| switch (ly.mixer) { + .conv => |c| { + var off: usize = 0; + while (off < 6 * H * 2) : (off += zero.len * 2) { + try m.r.upload(.{ .rt = c.state.rt, .ptr = @ptrFromInt(@intFromPtr(c.state.ptr.?) + off), .len = zero.len * 2 }, std.mem.sliceAsBytes(&zero)); + } + }, + .attn => {}, + }; + try m.r.sync(); + m.pos = 0; + } + + /// y[r * y_stride + y_off + o] = x[r] . W[o] for n rows: the decode matvec for one row, pdense for a window. + fn proj(m: *Model, t: Table, x: Buf, y: Buf, in_dim: u32, rows: u32, n: u32, y_stride: u32, y_off: u32) !void { + if (n == 1) return run(&m.k.qmv, .{ (rows + 3) / 4, 1, 1 }, .{ t.w, t.s, t.b, x, y, in_dim, @as(u32, 0), y_off, rows }); + try run(&m.k.pdense, .{ (rows + 15) / 16, (n + 7) / 8, 1 }, .{ t.w, t.s, t.b, x, y, in_dim, rows, n, y_stride, y_off, @as(u32, 0) }); + } + + fn mixer(m: *Model, ly: *const Layer, n: u32) !void { + const s = &m.s; + switch (ly.mixer) { + .conv => |c| { + try m.proj(c.in, s.xn, s.proj, H, 3 * H, n, 3 * H, 0); + try run(&m.k.conv, .{ H / 64, 1, 1 }, .{ s.proj, 3 * H, c.state, c.state, c.cw, s.att, n, H }); + try m.proj(c.out, s.att, s.d, H, H, n, H, 0); + }, + .attn => |a| { + try m.proj(a.q, s.xn, s.qkv, H, NQ * HD, n, QKV, 0); + try m.proj(a.k, s.xn, s.qkv, H, NKV * HD, n, QKV, NQ * HD); + try m.proj(a.v, s.xn, s.qkv, H, NKV * HD, n, QKV, (NQ + NKV) * HD); + try m.proj(a.g, s.xn, s.g, H, NQ, n, NQ, 0); + const inv = if (a.window == 0) m.inv_full else m.inv_slide; + try run(&m.k.rope, .{ NQ, n, 1 }, .{ s.qkv, QKV, @as(u32, 0), s.q, NQ * HD, a.qn, inv, m.pos, @as(u32, 0), cfg.eps }); + try run(&m.k.rope, .{ NKV, n, 1 }, .{ s.qkv, QKV, NQ * HD, a.kc, NKV * HD, a.kn, inv, m.pos, a.ring, cfg.eps }); + try run(&m.k.vst, .{ NKV * HD / 64, n, 1 }, .{ s.qkv, QKV, (NQ + NKV) * HD, a.vc, NKV * HD, m.pos, a.ring }); + const span = if (a.window == 0) m.pos + n else @min(m.pos + n, a.window); + const scale: f32 = 1.0 / @sqrt(@as(f32, HD)); + try run(&m.k.part, .{ NKV, (span + chunk - 1) / chunk, n }, .{ s.q, a.kc, a.vc, s.po, s.pm, s.pl, m.pos, chunk, a.ring, a.window, NKV, m.maxc, scale }); + try run(&m.k.merge, .{ NQ, n, 1 }, .{ s.po, s.pm, s.pl, s.g, s.att, m.pos, chunk, a.window, NQ, m.maxc }); + try m.proj(a.o, s.att, s.d, NQ * HD, H, n, H, 0); + }, + } + } + + fn hook(m: *Model, h: ?Hook, layer: i32, b: Buf, n: u32) !void { + const k = h orelse return; + const rows = m.host_rows[0 .. n * H]; + try m.r.sync(); + try m.r.download(std.mem.sliceAsBytes(rows), b); + try m.r.sync(); + k.f(k.ctx, layer, rows); + } + + /// Runs the tokens as one window at m.pos (1 .. max_rows rows); `logits`: the last row's logits and argmax too. + pub fn forward(m: *Model, ids: []const u32, logits: bool, h: ?Hook) !void { + const n: u32 = @intCast(ids.len); + if (n == 0 or n > m.max_rows) return error.Invalid; + if (m.pos + n > m.cap) return error.ContextFull; + const s = &m.s; + try m.r.upload(s.ids, std.mem.sliceAsBytes(ids)); + try run(&m.k.embed, .{ n, 1, 1 }, .{ m.emb.w, m.emb.s, m.emb.b, s.ids, s.x, H }); + try run(&m.k.rms, .{ n, 1, 1 }, .{ s.x, m.layers[0].op_norm, s.xn, H, cfg.eps }); + for (m.layers, 0..) |*ly, i| { + try m.mixer(ly, n); + try run(&m.k.pan, .{ n, 1, 1 }, .{ s.d, s.x, ly.post_op, ly.ffn_norm, s.xn, H, cfg.eps }); + try m.proj(ly.w1, s.xn, s.a, H, FF, n, FF, 0); + try m.proj(ly.w3, s.xn, s.b, H, FF, n, FF, 0); + try run(&m.k.swiglu, .{ (n * FF + 63) / 64, 1, 1 }, .{ s.a, s.b, s.act, n * FF }); + try m.proj(ly.w2, s.act, s.d, FF, H, n, H, 0); + const next = if (i + 1 < m.layers.len) m.layers[i + 1].op_norm else m.norm_f; + try run(&m.k.pan, .{ n, 1, 1 }, .{ s.d, s.x, ly.post_ffn, next, s.xn, H, cfg.eps }); + try m.hook(h, @intCast(i), s.x, n); + } + try m.hook(h, -1, s.xn, n); + m.pos += n; + if (!logits) return; + try run(&m.k.head, .{ V / 4, 1, 1 }, .{ m.emb.w, m.emb.s, m.emb.b, s.xn, s.logits, H, (n - 1) * H, @as(u32, 0), V }); + const per = (V + argmax_parts - 1) / argmax_parts; + try run(&m.k.am_part, .{ argmax_parts, 1, 1 }, .{ s.logits, s.am_scratch, V, per }); + try run(&m.k.am_fin, .{ 1, 1, 1 }, .{ s.am_scratch, s.am_out, argmax_parts }); + } + + /// Prefill of a prompt in windows of max_rows; the last window leaves its logits and argmax. + pub fn prefill(m: *Model, ids: []const u32, h: ?Hook) !void { + var i: usize = 0; + while (i < ids.len) { + const n = @min(ids.len - i, m.max_rows); + try m.forward(ids[i .. i + n], i + n == ids.len, h); + i += n; + } + } + + /// The greedy token of the last forward with logits. + pub fn argmax(m: *Model) !u32 { + var out: [1]i32 = undefined; + try m.r.sync(); + try m.r.download(std.mem.sliceAsBytes(&out), m.s.am_out); + try m.r.sync(); + return @intCast(out[0]); + } + + pub fn fetchLogits(m: *Model, out: []f32) !void { + try m.r.sync(); + try m.r.download(std.mem.sliceAsBytes(out[0..V]), m.s.logits); + try m.r.sync(); + } +}; diff --git a/zig/tests/xpu/lfm3_fwd.zig b/zig/tests/xpu/lfm3_fwd.zig new file mode 100644 index 0000000000..3048349c35 --- /dev/null +++ b/zig/tests/xpu/lfm3_fwd.zig @@ -0,0 +1,163 @@ +//! LFM3 forward on the checkpoint: every layer against ref64.py's float64 dumps, then 64 greedy tokens, twice. + +const std = @import("std"); +const rt = @import("rig.zig"); +const lfm3 = @import("lfm3_xpu"); +const stop = @import("xpu").stop; + +const H = 2048; +const prompts = [_][]const u8{ "capitals", "chat", "items" }; +const reply_len = 64; + +fn nowNs() u64 { + var ts: std.c.timespec = undefined; + _ = std.c.clock_gettime(.MONOTONIC, &ts); + return @as(u64, @intCast(ts.sec)) * 1_000_000_000 + @as(u64, @intCast(ts.nsec)); +} + +fn bf(v: u16) f64 { + return @as(f32, @bitCast(@as(u32, v) << 16)); +} + +/// The hidden rows of every layer (index 24: the final norm) as the windows pass, bf16. +const Rows = struct { + layers: [25][]u16, + row0: usize = 0, + + fn keep(ctx: *anyopaque, layer: i32, rows: []const u16) void { + const s: *Rows = @ptrCast(@alignCast(ctx)); + const li: usize = if (layer < 0) 24 else @intCast(layer); + @memcpy(s.layers[li][s.row0 * H ..][0..rows.len], rows); + } +}; + +/// Per-row ||xpu - ref|| / ||ref||: mean and max over the rows. +fn relRows(got: []const u16, want: []const f32) [2]f64 { + const n = want.len / H; + var sum: f64 = 0; + var mx: f64 = 0; + for (0..n) |r| { + var d: f64 = 0; + var w: f64 = 0; + for (0..H) |i| { + const a = bf(got[r * H + i]); + const b: f64 = want[r * H + i]; + d += (a - b) * (a - b); + w += b * b; + } + const e = @sqrt(d / w); + sum += e; + mx = @max(mx, e); + } + return .{ sum / @as(f64, @floatFromInt(n)), mx }; +} + +fn readF32(gpa: std.mem.Allocator, io: std.Io, path: []const u8) ![]f32 { + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .unlimited); + const out = try gpa.alloc(f32, bytes.len / 4); + @memcpy(std.mem.sliceAsBytes(out), bytes); + gpa.free(bytes); + return out; +} + +pub const panic = std.debug.FullPanic(@import("rig.zig").panicDrain); +pub fn main(init: std.process.Init) !void { + stop.install(); + const gpa = init.gpa; + const io = init.io; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 3) { + std.debug.print("skipped: usage tf-xpu-lfm3_fwd-test MODEL_DIR REF_DIR [WINDOW_ROWS] [OUT_JSON]\n", .{}); + return; + } + const dir = args[1]; + const ref = args[2]; + const rows: u32 = if (args.len > 3) try std.fmt.parseInt(u32, args[3], 10) else 64; + var r = try rt.open(); + defer r.deinit(); + var c = try lfm3.config.Config.read(gpa, io, dir); + defer c.deinit(); + var l = try @import("xpu").loader.Loader.init(gpa, &r, dir); + var t0 = nowNs(); + var m = try lfm3.model.Model.load(gpa, &r, &l, &c, 4096, rows); + std.debug.print("loaded in {d:.1} s, {d:.0} MiB on the device, windows of {d} rows\n", .{ @as(f64, @floatFromInt(nowNs() - t0)) / 1e9, @as(f64, @floatFromInt(m.device_bytes)) / (1 << 20), rows }); + + var out: std.Io.Writer.Allocating = .init(gpa); + defer out.deinit(); + try out.writer.writeAll("{"); + var bad = false; + for (prompts, 0..) |name, pi| { + const meta_path = try std.fmt.allocPrint(gpa, "{s}/{s}/meta.json", .{ ref, name }); + const meta_bytes = try std.Io.Dir.cwd().readFileAlloc(io, meta_path, gpa, .unlimited); + const meta = try std.json.parseFromSlice(std.json.Value, gpa, meta_bytes, .{}); + const id_vals = meta.value.object.get("ids").?.array.items; + const ids = try gpa.alloc(u32, id_vals.len); + for (id_vals, ids) |v, *o| o.* = @intCast(v.integer); + const T = ids.len; + var keep: Rows = .{ .layers = undefined }; + for (&keep.layers) |*ly| ly.* = try gpa.alloc(u16, T * H); + const hk: lfm3.model.Hook = .{ .ctx = &keep, .f = Rows.keep }; + try m.reset(); + var i: usize = 0; + while (i < T) { + const n = @min(T - i, rows); + keep.row0 = i; + try m.forward(ids[i .. i + n], i + n == T, hk); + i += n; + } + const next = try m.argmax(); + const gaps = meta.value.object.get("bf16_gap").?.object; + std.debug.print("{s}: {d} prompt tokens; per-row relative L2 error against float64 (mean / max), and the bf16-points pass's own gap\n", .{ name, T }); + var worst: f64 = 0; + for (0..25) |li| { + const layer: i32 = if (li == 24) -1 else @intCast(li); + const path = if (li == 24) try std.fmt.allocPrint(gpa, "{s}/{s}/final.f32", .{ ref, name }) else try std.fmt.allocPrint(gpa, "{s}/{s}/h{d:0>2}.f32", .{ ref, name, li }); + const want = try readF32(gpa, io, path); + defer gpa.free(want); + const e = relRows(keep.layers[li], want); + var key: [8]u8 = undefined; + const g = gaps.get(try std.fmt.bufPrint(&key, "{d}", .{layer})).?.array.items; + const gm: f64 = switch (g[0]) { + .float => |f| f, + else => 0, + }; + const gx: f64 = switch (g[1]) { + .float => |f| f, + else => 0, + }; + const ratio = e[0] / gm; + worst = @max(worst, ratio); + std.debug.print(" {s} {d:>2}: xpu {e:.2} / {e:.2} bf16 gap {e:.2} / {e:.2} mean ratio {d:.2}\n", .{ if (li == 24) "final" else "layer", if (li == 24) 0 else li, e[0], e[1], gm, gx, ratio }); + } + const top = meta.value.object.get("top1").?.array.items; + const want_next: u32 = @intCast(top[T - 1].integer); + std.debug.print(" next token after the prompt: xpu {d}, float64 {d}; worst layer mean ratio {d:.2}\n", .{ next, want_next, worst }); + if (worst > 4.0) bad = true; + // greedy reply, twice from a fresh prefill: the same tokens + var reply: [2][reply_len]u32 = undefined; + var dec_ns: u64 = 0; + for (0..2) |run_i| { + try m.reset(); + try m.prefill(ids, null); + reply[run_i][0] = try m.argmax(); + t0 = nowNs(); + for (1..reply_len) |j| { + try m.forward(reply[run_i][j - 1 .. j], true, null); + reply[run_i][j] = try m.argmax(); + } + dec_ns = nowNs() - t0; + } + const same = std.mem.eql(u32, &reply[0], &reply[1]); + if (!same) bad = true; + std.debug.print(" {d} greedy tokens, second run {s}; decode {d:.1} tok/s (synchronous argmax per token)\n", .{ reply_len, if (same) "identical" else "DIFFERENT", @as(f64, reply_len - 1) / (@as(f64, @floatFromInt(dec_ns)) / 1e9) }); + try out.writer.print("{s}\"{s}\": {{\"prompt\": [", .{ if (pi == 0) "" else ", ", name }); + for (ids, 0..) |v, k| try out.writer.print("{s}{d}", .{ if (k == 0) "" else ", ", v }); + try out.writer.writeAll("], \"reply\": ["); + for (reply[0], 0..) |v, k| try out.writer.print("{s}{d}", .{ if (k == 0) "" else ", ", v }); + try out.writer.writeAll("]}"); + } + try out.writer.writeAll("}\n"); + if (args.len > 4) try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = args[4], .data = out.written() }); + std.debug.print("{s}\n", .{if (bad) "lfm3_fwd: FAILED" else "lfm3_fwd: ok"}); + if (bad) return error.TestFailed; +} From 0ae6b1e8e83f93973a04f695dfe1eba6d093aa0a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 30 Oct 2026 15:08:22 +0100 Subject: [PATCH 07/14] xpu: serve LFM3 (mlx-q4g64) in xpu-native through the lane host, prompt projections on nem_prow_gen12 without DPAS, Iris Xe measurements --- docs/recipes/xpu-arc.md | 11 ++++ zig/build/xpu.zig | 16 +++-- zig/src/families/lfm3/xpu.zig | 1 + zig/src/families/lfm3/xpu_engine.zig | 89 ++++++++++++++++++++++++++++ zig/src/families/lfm3/xpu_model.zig | 10 +++- zig/src/native/xpu.zig | 6 +- zig/src/native/xpu_lfm3.zig | 34 +++++++++++ zig/src/native/xpu_nem.zig | 4 +- zig/tests/xpu/lfm3_bench.zig | 46 ++++++++++++++ 9 files changed, 204 insertions(+), 13 deletions(-) create mode 100644 zig/src/families/lfm3/xpu_engine.zig create mode 100644 zig/src/native/xpu_lfm3.zig create mode 100644 zig/tests/xpu/lfm3_bench.zig diff --git a/docs/recipes/xpu-arc.md b/docs/recipes/xpu-arc.md index f73fcec648..5474455792 100644 --- a/docs/recipes/xpu-arc.md +++ b/docs/recipes/xpu-arc.md @@ -50,6 +50,17 @@ Intel Arc Pro B70 Graphics (32.5 GB maximum single allocation), Linux 7.0.0 (`xe Nemotron MLX 4-bit checkpoint (`TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit`, revision `d9d758fb83953437f7263256b0d96157e2a348b8`, 729 tensors, 17.27 GiB of weights). The commit under test is the head of the pull request. +LFM3-1.7B (`LiquidAI/LFM3-1.7B-MLX-4bit`, MLX affine 4-bit, group 64) on an Intel Iris Xe laptop GPU (Gen12LP, `tgllp`, 96 EUs, no +DPAS) under WSL2 (Ubuntu 24.04, 8 GB given to WSL), Ubuntu's compute runtime 23.43.27642 and Level Zero loader 1.16, Zig 0.17.0, +`zig build xpu-native -Dxpu-device=tgllp` (DPAS off, so the server is built without the Qwen3.8-27B runner): + +| Setup | Decode, c=1 | TTFT, 1,024 tokens | TTFT, 8,192 tokens | Memory (`serve --context 16384`) | +| --- | ---: | ---: | ---: | --- | +| Iris Xe, LFM3-1.7B MLX 4-bit | 14.8 to 15.8 tok/s | 13.1 to 14.2 s | 130.8 to 131.3 s | 1.25 GB device allocation, 441 MB host resident | + +Decode: 256 greedy tokens after a 64-token prompt; TTFT: prompts of numbered item sentences, one token generated; three runs each, +through `/v1/completions`. A `/v1/chat/completions` request with the checkpoint's chat template answers (thinking on, `finish_reason` stop). + ## Measurements Decode, 128 greedy tokens from a 5-token prompt: 5.72 ms/token, 175 tok/s (`tensorfold-xpu run`, five runs on different diff --git a/zig/build/xpu.zig b/zig/build/xpu.zig index 885eff13c5..a4564edeff 100644 --- a/zig/build/xpu.zig +++ b/zig/build/xpu.zig @@ -139,8 +139,8 @@ pub fn steps(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.buil const lanes = b.createModule(.{ .root_source_file = b.path("zig/src/core/lanes/lanes.zig"), .target = target, .optimize = optimize, .link_libc = true }); const nem = nemotronModules(b, target, optimize, ocloc, device, dpas, ze_mod, lanes); const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} }); - const engines = b.createModule(.{ .root_source_file = b.path("zig/src/native/xpu.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "engine_api", .module = api }, .{ .name = "lanes", .module = lanes }, .{ .name = "xpu", .module = nem.xpu }, .{ .name = "nemotron_xpu", .module = nem.nemotron } } }); - // the Qwen3.8-27B runner and Gemma 4 E4B (DPAS, nGEN) only with DPAS: a Gen12LP server serves Nemotron + const engines = b.createModule(.{ .root_source_file = b.path("zig/src/native/xpu.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "engine_api", .module = api }, .{ .name = "lanes", .module = lanes }, .{ .name = "xpu", .module = nem.xpu }, .{ .name = "nemotron_xpu", .module = nem.nemotron }, .{ .name = "lfm3_xpu", .module = lfm3Module(b, target, optimize, nem.xpu) } } }); + // the Qwen3.8-27B runner and Gemma 4 E4B (DPAS, nGEN) only with DPAS: a Gen12LP server serves Nemotron and LFM3 const engine_options = b.addOptions(); engine_options.addOption(bool, "qwen27", dpas); engine_options.addOption(bool, "gemma4", dpas); @@ -256,9 +256,16 @@ fn hasDpas(device: []const u8) bool { return true; } +/// LFM3's forward and engine on the runtime module `xpu`. +fn lfm3Module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, xpu: *std.Build.Module) *std.Build.Module { + const m = b.createModule(.{ .root_source_file = b.path("zig/src/families/lfm3/xpu.zig"), .target = target, .optimize = optimize, .link_libc = true }); + m.addImport("xpu", xpu); + return m; +} + /// The standalone op and model tests, each its own program `tf-xpu--test` (step `xpu-tests`). const test_programs = [_][]const u8{ - "panic", "nem_rows", "nem_prefill", "nem_gen", "nem_mtp", "nem_kl", "nem_norm_bench", "nem_mv_bench", "nem_slot", "nem_format", "nem_seam", "nem_attn_dec", "nem_fp64", "memory", "lfm3_conv", "lfm3_rope", "lfm3_attn", "lfm3_norm", "lfm3_fwd", + "panic", "nem_rows", "nem_prefill", "nem_gen", "nem_mtp", "nem_kl", "nem_norm_bench", "nem_mv_bench", "nem_slot", "nem_format", "nem_seam", "nem_attn_dec", "nem_fp64", "memory", "lfm3_conv", "lfm3_rope", "lfm3_attn", "lfm3_norm", "lfm3_fwd", "lfm3_bench", }; /// `-Dxpu` (Linux): Nemotron's SPIR-V images, `tensorfold-xpu` and `tf-xpu-test`; needs ocloc, not nvcc or CUDA. @@ -275,8 +282,7 @@ fn nemotronSteps(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std. runner.addImport("nemotron_xpu", mods.nemotron); b.installArtifact(b.addExecutable(.{ .name = "tf-xpu-test", .root_module = runner })); const tests_step = b.step("xpu-tests", "Build the standalone Intel GPU op and model tests (tf-xpu--test)"); - const lfm3 = b.createModule(.{ .root_source_file = b.path("zig/src/families/lfm3/xpu.zig"), .target = target, .optimize = optimize, .link_libc = true }); - lfm3.addImport("xpu", xpu); + const lfm3 = lfm3Module(b, target, optimize, xpu); for (test_programs) |name| { const mod = b.createModule(.{ .root_source_file = b.path(b.fmt("zig/tests/xpu/{s}.zig", .{name})), .target = target, .optimize = optimize, .link_libc = true }); mod.addImport("xpu", xpu); diff --git a/zig/src/families/lfm3/xpu.zig b/zig/src/families/lfm3/xpu.zig index c10d895352..3cefba81da 100644 --- a/zig/src/families/lfm3/xpu.zig +++ b/zig/src/families/lfm3/xpu.zig @@ -1,3 +1,4 @@ //! LFM3 on the Intel GPU (XPU backend): the forward and the config it reads. pub const model = @import("xpu_model.zig"); +pub const engine = @import("xpu_engine.zig"); pub const config = @import("config.zig"); diff --git a/zig/src/families/lfm3/xpu_engine.zig b/zig/src/families/lfm3/xpu_engine.zig new file mode 100644 index 0000000000..0d47070b42 --- /dev/null +++ b/zig/src/families/lfm3/xpu_engine.zig @@ -0,0 +1,89 @@ +//! LFM3 behind the Intel GPU lane host: one sequence on the device, the prompt in windows, greedy argmax on the device. +const std = @import("std"); +const xpu = @import("xpu"); +const cfg = @import("config.zig"); +const model = @import("xpu_model.zig"); + +pub const Options = struct { context: usize = 4096, window_rows: u32 = 64 }; + +pub const Engine = struct { + gpa: std.mem.Allocator, + stream: xpu.Stream, + r: xpu.rt.Runtime, + loader: xpu.loader.Loader, + c: cfg.Config, + m: model.Model, + max_len: usize, + window_rows: u32, + load_seconds: f64, + + pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const xpu.Context, dir: []const u8, o: Options) !*Engine { + const t0 = std.Io.Clock.awake.now(io); + const e = try gpa.create(Engine); + errdefer gpa.destroy(e); + e.gpa = gpa; + e.c = try cfg.Config.read(gpa, io, dir); + errdefer e.c.deinit(); + if (e.c.kind != .text) return error.UnsupportedFormat; + const q = e.c.quant.global; + if (q.bits != 4 or q.group != 64 or e.c.quant.modules.count() != 0) return error.UnsupportedFormat; + e.stream = try xpu.Stream.init(ctx); + errdefer e.stream.deinit(); + e.r = try xpu.rt.Runtime.init(ctx, e.stream); + errdefer e.r.deinit(); + e.loader = try xpu.loader.Loader.init(gpa, &e.r, dir); + e.window_rows = o.window_rows; + e.m = try model.Model.load(gpa, &e.r, &e.loader, &e.c, @intCast(o.context), o.window_rows); + e.max_len = o.context; + e.load_seconds = xpu.decode.seconds(io, t0); + return e; + } + + /// The process exits right after: device memory goes with it once the queue is drained. + pub fn deinit(e: *Engine) void { + e.r.deinit(); + e.stream.deinit(); + e.c.deinit(); + e.gpa.destroy(e); + } + + pub fn deviceBytes(e: *const Engine) u64 { + return e.loader.total; + } + + pub fn vocab(_: *const Engine) usize { + return cfg.vocab; + } + + pub fn reset(e: *Engine) !void { + try e.m.reset(); + } + + pub fn feed(e: *Engine, token: u32, logits: bool) !void { + try e.m.forward(&.{token}, logits, null); + } + + pub fn sync(e: *Engine) !void { + try e.r.sync(); + } + + pub fn argmax(e: *Engine) !u32 { + return e.m.argmax(); + } + + pub fn fetchLogits(e: *Engine, out: []f32) !void { + try e.m.fetchLogits(out); + } + + /// The prompt head in windows of at most `rows` tokens, without logits. + pub fn prefillWindows(e: *Engine, head: []const u32, rows: u32) !void { + var t: usize = 0; + const w = @min(rows, e.window_rows); + while (t < head.len) { + const n = @min(w, head.len - t); + try e.m.forward(head[t .. t + n], false, null); + t += n; + } + try e.r.sync(); + } +}; diff --git a/zig/src/families/lfm3/xpu_model.zig b/zig/src/families/lfm3/xpu_model.zig index be204ac057..6f9d51c468 100644 --- a/zig/src/families/lfm3/xpu_model.zig +++ b/zig/src/families/lfm3/xpu_model.zig @@ -17,6 +17,10 @@ const HD: u32 = cfg.head_dim; const QKV: u32 = (NQ + 2 * NKV) * HD; const chunk: u32 = 128; // keys a split-K work-group scans const argmax_parts: u32 = 64; +/// Without DPAS (Gen12LP) prompt projections take nem_prow_gen12: 4 outputs a sub-group, as 8 x 8 spills on Gen12LP. +const gen12 = !xpu.kernels.dpas; +const tile_rows: u32 = 8; +const tile_outs: u32 = if (gen12) 4 else 8; pub const Table = struct { w: Buf, s: Buf, b: Buf }; const Conv = struct { in: Table, out: Table, cw: Buf, state: Buf }; @@ -110,7 +114,7 @@ pub const Model = struct { if (max_rows == 0 or max_rows > 256 or cap < max_rows) return error.Invalid; var bm = try r.module(xpu.kernels.basic); var am = try r.module(xpu.kernels.attn); - var pm = try r.module(xpu.kernels.nem_prow); + var pm = try r.module(if (gen12) xpu.kernels.nem_prow_gen12 else xpu.kernels.nem_prow); var fm = try r.module(xpu.kernels.lfm3); const k: Kernels = .{ .embed = try bm.kernel("embed4", .{ 64, 1, 1 }), @@ -118,7 +122,7 @@ pub const Model = struct { .head = try am.kernel("qmv4_f32", .{ 64, 1, 1 }), .am_part = try am.kernel("argmax_partial", .{ 256, 1, 1 }), .am_fin = try am.kernel("argmax_final", .{ 256, 1, 1 }), - .pdense = try pm.kernel("pdense", .{ 32, 1, 1 }), + .pdense = try pm.kernel("pdense", .{ 16 * (16 / tile_outs), 1, 1 }), .rms = try fm.kernel("lfm3_rmsnorm", .{ 256, 1, 1 }), .pan = try fm.kernel("lfm3_post_add_norm", .{ 256, 1, 1 }), .conv = try fm.kernel("lfm3_conv_rows", .{ 64, 1, 1 }), @@ -229,7 +233,7 @@ pub const Model = struct { /// y[r * y_stride + y_off + o] = x[r] . W[o] for n rows: the decode matvec for one row, pdense for a window. fn proj(m: *Model, t: Table, x: Buf, y: Buf, in_dim: u32, rows: u32, n: u32, y_stride: u32, y_off: u32) !void { if (n == 1) return run(&m.k.qmv, .{ (rows + 3) / 4, 1, 1 }, .{ t.w, t.s, t.b, x, y, in_dim, @as(u32, 0), y_off, rows }); - try run(&m.k.pdense, .{ (rows + 15) / 16, (n + 7) / 8, 1 }, .{ t.w, t.s, t.b, x, y, in_dim, rows, n, y_stride, y_off, @as(u32, 0) }); + try run(&m.k.pdense, .{ (rows + 15) / 16, (n + tile_rows - 1) / tile_rows, 1 }, .{ t.w, t.s, t.b, x, y, in_dim, rows, n, y_stride, y_off, @as(u32, 0) }); } fn mixer(m: *Model, ly: *const Layer, n: u32) !void { diff --git a/zig/src/native/xpu.zig b/zig/src/native/xpu.zig index 6ec6e1ebf4..cce64198d1 100644 --- a/zig/src/native/xpu.zig +++ b/zig/src/native/xpu.zig @@ -15,7 +15,7 @@ const Allocator = std.mem.Allocator; pub const backends: []const []const u8 = &.{"xpu"}; const qwen27_family: []const api.Family = if (has_qwen27) &.{.{ .model_type = "qwen3_5", .formats = &.{ "mlx-q4g64", "gptq-int4-g128" } }} else &.{}; const gemma4_family: []const api.Family = if (nem.has_gemma4) &.{.{ .model_type = "gemma4", .formats = &.{"mlx-q4g64"} }} else &.{}; -pub const families: []const api.Family = qwen27_family ++ &[_]api.Family{.{ .model_type = "nemotron_h", .formats = &.{"mlx-q4g64"} }} ++ gemma4_family; +pub const families: []const api.Family = qwen27_family ++ &[_]api.Family{ .{ .model_type = "nemotron_h", .formats = &.{"mlx-q4g64"} }, .{ .model_type = "lfm3", .formats = &.{"mlx-q4g64"} } } ++ gemma4_family; pub fn chip(_: Allocator) ?[]const u8 { return "Intel Xe2 (Level Zero)"; @@ -181,8 +181,8 @@ pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]c if (nem.serves(o.model_type)) return nem.open(a, gpa, io, o, problem); if (comptime has_qwen27) { if (std.mem.eql(u8, o.model_type, "qwen3_5")) return openQwen27(a, gpa, io, o, problem); - problem.* = try std.fmt.allocPrint(a, "the XPU engine serves Qwen3.8-27B (qwen3_5, MLX affine 4/64 or symmetric GPTQ 4/128) and Nemotron-H (nemotron_h, MLX 4/64), not {s}", .{o.model_type}); - } else problem.* = try std.fmt.allocPrint(a, "this XPU build (no DPAS) serves Nemotron-H (nemotron_h, MLX 4/64), not {s}: Qwen3.8-27B and Gemma 4 E4B need DPAS (-Dxpu-dpas=true)", .{o.model_type}); + problem.* = try std.fmt.allocPrint(a, "the XPU engine serves Qwen3.8-27B (qwen3_5, MLX affine 4/64 or symmetric GPTQ 4/128), Nemotron-H (nemotron_h, MLX 4/64) and LFM3 (lfm3, MLX 4/64), not {s}", .{o.model_type}); + } else problem.* = try std.fmt.allocPrint(a, "this XPU build (no DPAS) serves Nemotron-H (nemotron_h, MLX 4/64) and LFM3 (lfm3, MLX 4/64), not {s}: Qwen3.8-27B and Gemma 4 E4B need DPAS (-Dxpu-dpas=true)", .{o.model_type}); return null; } diff --git a/zig/src/native/xpu_lfm3.zig b/zig/src/native/xpu_lfm3.zig new file mode 100644 index 0000000000..007aa005f9 --- /dev/null +++ b/zig/src/native/xpu_lfm3.zig @@ -0,0 +1,34 @@ +//! The LFM3 entry of the Intel GPU lane host: its family table row, loader and prompt window. +const std = @import("std"); +const xpu = @import("xpu"); +const api = @import("engine_api"); +const lfm3 = @import("lfm3_xpu"); +const Allocator = std.mem.Allocator; + +pub const model_type = "lfm3"; +pub const formats: []const []const u8 = &.{"mlx-q4g64"}; +pub const name = "xpu-lfm3"; +pub const Engine = lfm3.engine.Engine; + +/// Rows of one prompt window (pdense above one row; the scratch is sized for it). +const prompt_rows: u32 = 64; + +pub fn promptRows(_: *const Engine) u32 { + return prompt_rows; +} + +/// What the startup line adds about the engine's settings. +pub fn note(a: Allocator, e: *const Engine) ![]const u8 { + return std.fmt.allocPrint(a, ", prompt windows of {d} rows, no draft head", .{e.window_rows}); +} + +/// The engine for `o.dir`, or null with `problem` set (the caller releases the device). +pub fn init(a: Allocator, io: std.Io, ctx: *const xpu.Context, o: api.Open, window: usize, problem: *[]const u8) !?*Engine { + return Engine.init(std.heap.page_allocator, io, ctx, o.dir, .{ .context = window, .window_rows = prompt_rows }) catch |e| { + problem.* = if (e == error.UnsupportedFormat) + try std.fmt.allocPrint(a, "{s}: the Intel GPU engine reads LFM3 text checkpoints in MLX affine 4-bit, group 64", .{o.dir}) + else + try std.fmt.allocPrint(a, "the Intel GPU engine cannot load {s} ({s})", .{ o.dir, @errorName(e) }); + return null; + }; +} diff --git a/zig/src/native/xpu_nem.zig b/zig/src/native/xpu_nem.zig index 61344fbb54..2b87e1a333 100644 --- a/zig/src/native/xpu_nem.zig +++ b/zig/src/native/xpu_nem.zig @@ -1,4 +1,4 @@ -//! Nemotron on the XPU server: its lane host (one family behind the lane core, one stream, host sampling). +//! Nemotron and LFM3 on the XPU server: their lane host (a family behind the lane core, one stream, host sampling). const std = @import("std"); const xpu = @import("xpu"); const api = @import("engine_api"); @@ -10,7 +10,7 @@ const Allocator = std.mem.Allocator; pub const has_gemma4 = @import("xpu_engine_options").gemma4; /// The families this host serves, each a file with its table row, loader and prompt window. -const specs = if (has_gemma4) .{ @import("xpu_nemotron.zig"), @import("xpu_gemma4.zig") } else .{@import("xpu_nemotron.zig")}; +const specs = if (has_gemma4) .{ @import("xpu_nemotron.zig"), @import("xpu_gemma4.zig"), @import("xpu_lfm3.zig") } else .{ @import("xpu_nemotron.zig"), @import("xpu_lfm3.zig") }; /// Prompt plus reply tokens when --context names none (and the model's window is larger). const default_context: i64 = 32768; diff --git a/zig/tests/xpu/lfm3_bench.zig b/zig/tests/xpu/lfm3_bench.zig new file mode 100644 index 0000000000..4ced6e492f --- /dev/null +++ b/zig/tests/xpu/lfm3_bench.zig @@ -0,0 +1,46 @@ +//! LFM3 prefill and decode timing on the device: a prompt of N random ids in windows of W rows, then 64 greedy tokens. + +const std = @import("std"); +const rt = @import("rig.zig"); +const lfm3 = @import("lfm3_xpu"); +const stop = @import("xpu").stop; + +fn nowNs() u64 { + var ts: std.c.timespec = undefined; + _ = std.c.clock_gettime(.MONOTONIC, &ts); + return @as(u64, @intCast(ts.sec)) * 1_000_000_000 + @as(u64, @intCast(ts.nsec)); +} + +pub const panic = std.debug.FullPanic(@import("rig.zig").panicDrain); +pub fn main(init: std.process.Init) !void { + stop.install(); + const gpa = init.gpa; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 2) { + std.debug.print("skipped: usage tf-xpu-lfm3_bench-test MODEL_DIR [PROMPT_TOKENS] [WINDOW_ROWS]\n", .{}); + return; + } + const n: usize = if (args.len > 2) try std.fmt.parseInt(usize, args[2], 10) else 1024; + const w: u32 = if (args.len > 3) try std.fmt.parseInt(u32, args[3], 10) else 64; + var r = try rt.open(); + defer r.deinit(); + var c = try lfm3.config.Config.read(gpa, init.io, args[1]); + defer c.deinit(); + var l = try @import("xpu").loader.Loader.init(gpa, &r, args[1]); + var m = try lfm3.model.Model.load(gpa, &r, &l, &c, @intCast(n + 128), w); + var prng = std.Random.DefaultPrng.init(5); + const ids = try gpa.alloc(u32, n); + for (ids) |*v| v.* = prng.random().uintLessThan(u32, 120000); + try m.reset(); + var t0 = nowNs(); + try m.prefill(ids, null); + var tok = try m.argmax(); + const pre = @as(f64, @floatFromInt(nowNs() - t0)) / 1e9; + t0 = nowNs(); + for (0..64) |_| { + try m.forward(&.{tok}, true, null); + tok = try m.argmax(); + } + const dec = @as(f64, @floatFromInt(nowNs() - t0)) / 1e9; + std.debug.print("prompt {d} tokens in windows of {d}: prefill {d:.2} s ({d:.1} tok/s); decode after it {d:.1} tok/s; device {d:.0} MiB\n", .{ n, w, pre, @as(f64, @floatFromInt(n)) / pre, 64.0 / dec, @as(f64, @floatFromInt(m.device_bytes)) / (1 << 20) }); +} From ddfa1feecc2399e7c4c9f679a407f0241bff5dea Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 31 Oct 2026 13:53:59 +0100 Subject: [PATCH 08/14] xpu: LFM3's per-stream state in sequences, with kept prompt states copied on the device LFM3's per-stream state (the conv rows of 13 layers, the KV slots of 11, the position) moves out of the layers into a sequence the model binds; the engine makes, frees and binds sequences and saves and restores prompt states (the conv rows, the full layers' slots up to the mark, the sliding rings whole). The lane host serves LFM3 on one sequence until the engine verifies windows (a later commit). --- zig/src/families/lfm3/xpu_engine.zig | 66 ++++++++++++ zig/src/families/lfm3/xpu_model.zig | 152 +++++++++++++++++++++++---- 2 files changed, 196 insertions(+), 22 deletions(-) diff --git a/zig/src/families/lfm3/xpu_engine.zig b/zig/src/families/lfm3/xpu_engine.zig index 0d47070b42..848d235970 100644 --- a/zig/src/families/lfm3/xpu_engine.zig +++ b/zig/src/families/lfm3/xpu_engine.zig @@ -7,6 +7,8 @@ const model = @import("xpu_model.zig"); pub const Options = struct { context: usize = 4096, window_rows: u32 = 64 }; pub const Engine = struct { + pub const Sequence = model.Seq; + pub const Snapshot = model.Snap; gpa: std.mem.Allocator, stream: xpu.Stream, r: xpu.rt.Runtime, @@ -16,6 +18,8 @@ pub const Engine = struct { max_len: usize, window_rows: u32, load_seconds: f64, + /// The first stream's state, held here so the lane host can tell it from the streams it adds (newSeq). + own: Sequence, pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const xpu.Context, dir: []const u8, o: Options) !*Engine { const t0 = std.Io.Clock.awake.now(io); @@ -34,6 +38,10 @@ pub const Engine = struct { e.loader = try xpu.loader.Loader.init(gpa, &e.r, dir); e.window_rows = o.window_rows; e.m = try model.Model.load(gpa, &e.r, &e.loader, &e.c, @intCast(o.context), o.window_rows); + e.own = e.m.own.*; + gpa.destroy(e.m.own); + e.m.own = &e.own; + e.m.seq = &e.own; e.max_len = o.context; e.load_seconds = xpu.decode.seconds(io, t0); return e; @@ -75,6 +83,64 @@ pub const Engine = struct { try e.m.fetchLogits(out); } + pub fn newSeq(e: *Engine) !*Sequence { + return e.m.newSeq(); + } + + pub fn freeSeq(e: *Engine, q: *Sequence) void { + e.m.freeSeq(q); + } + + pub fn bind(e: *Engine, q: *Sequence) void { + e.m.bind(q); + } + + /// Prompt tokens [from, to) of the bound stream in windows of at most `rows`; the last window keeps logits. + pub fn prefillSpan(e: *Engine, ids: []const u32, from: usize, to: usize, rows: u32) !void { + if (e.m.pos != from) return error.PositionMismatch; + const w = @min(rows, e.window_rows); + var t = from; + while (t < to) { + const n = @min(w, to - t); + try e.m.forward(ids[t .. t + n], t + n == ids.len, null); + t += n; + } + } + + /// The bound stream's next token after its last forward with logits (one row). + pub fn rowTokens(e: *Engine, out: []i32) !void { + if (out.len != 1) return error.WindowTooWide; + out[0] = @intCast(try e.m.argmax()); + } + + pub fn rowLogits(e: *Engine, out: []f32) !void { + if (out.len != cfg.vocab) return error.WindowTooWide; + try e.m.fetchLogits(out); + } + + /// One decode step of the bound stream, with its logits. + pub fn stepRow(e: *Engine, token: u32) !void { + try e.m.forward(&.{token}, true, null); + } + + pub fn snapBytes(e: *const Engine, at: u32) u64 { + return e.m.snapBytes(at); + } + + /// The bound stream's state at prompt position `at`, which it has just reached. + pub fn snapSave(e: *Engine, at: u32) !*Snapshot { + if (at != e.m.pos) return error.PositionMismatch; + return e.m.snapSave(); + } + + pub fn snapRestore(e: *Engine, sn: *const Snapshot) !void { + try e.m.snapRestore(sn); + } + + pub fn snapDrop(e: *Engine, sn: *Snapshot) void { + e.m.snapDrop(sn); + } + /// The prompt head in windows of at most `rows` tokens, without logits. pub fn prefillWindows(e: *Engine, head: []const u32, rows: u32) !void { var t: usize = 0; diff --git a/zig/src/families/lfm3/xpu_model.zig b/zig/src/families/lfm3/xpu_model.zig index 6f9d51c468..d96c828a71 100644 --- a/zig/src/families/lfm3/xpu_model.zig +++ b/zig/src/families/lfm3/xpu_model.zig @@ -23,8 +23,8 @@ const tile_rows: u32 = 8; const tile_outs: u32 = if (gen12) 4 else 8; pub const Table = struct { w: Buf, s: Buf, b: Buf }; -const Conv = struct { in: Table, out: Table, cw: Buf, state: Buf }; -const Attn = struct { q: Table, k: Table, v: Table, g: Table, o: Table, qn: Buf, kn: Buf, kc: Buf, vc: Buf, ring: u32, window: u32 }; +const Conv = struct { in: Table, out: Table, cw: Buf }; +const Attn = struct { q: Table, k: Table, v: Table, g: Table, o: Table, qn: Buf, kn: Buf, ring: u32, window: u32 }; const Mixer = union(enum) { conv: Conv, attn: Attn }; const Layer = struct { op_norm: Buf, post_op: Buf, ffn_norm: Buf, post_ffn: Buf, w1: Table, w3: Table, w2: Table, mixer: Mixer }; @@ -45,6 +45,20 @@ const Kernels = struct { swiglu: rt.Kernel, }; +/// One stream's state: each conv layer's last six B*x rows, each attention layer's KV slots, and its position. +pub const Seq = struct { conv: [cfg.layers]Buf = undefined, kc: [cfg.layers]Buf = undefined, vc: [cfg.layers]Buf = undefined, pos: u32 = 0 }; + +/// A kept prompt state at `at`: conv rows, the full layers' first `at` KV slots, the sliding rings, in one buffer. +pub const Snap = struct { at: u32, buf: Buf }; + +const conv_bytes: usize = 6 * H * 2; +const slot_bytes: usize = NKV * HD * 2; + +/// `b` from byte `off` on. +fn from(b: Buf, off: usize) Buf { + return .{ .rt = b.rt, .ptr = @ptrFromInt(@intFromPtr(b.ptr.?) + off), .len = b.len - off }; +} + /// Device scratch for a window of up to `max_rows` rows. const Scratch = struct { ids: Buf, x: Buf, xn: Buf, d: Buf, proj: Buf, qkv: Buf, g: Buf, q: Buf, att: Buf, a: Buf, b: Buf, act: Buf, po: Buf, pm: Buf, pl: Buf, logits: Buf, am_scratch: Buf, am_out: Buf }; @@ -70,6 +84,7 @@ pub fn invFreq(theta: f32, partial: f32, out: *[HD / 2]f32) void { } pub const Model = struct { + gpa: std.mem.Allocator, r: *rt.Runtime, layers: []Layer, emb: Table, @@ -84,6 +99,9 @@ pub const Model = struct { maxc: u32, pos: u32 = 0, device_bytes: u64 = 0, + /// The bound stream's state (`own` until another is bound); `pos` is the bound stream's position. + seq: *Seq = undefined, + own: *Seq = undefined, /// The affine tensors of module `m` (in layer `layer`, or a model tensor) by names.zig in namespace `ns`. fn table(l: *ld.Loader, ns: names.Namespace, layer: ?usize, m: names.Module) !Table { @@ -170,7 +188,6 @@ pub const Model = struct { .in = try table(l, ns, i, .conv_in), .out = try table(l, ns, i, .conv_out), .cw = try vec(l, ns, i, .conv), - .state = try l.zeros(6 * H * 2), } } else blk: { const ring: u32 = if (kind == .sliding_attention) cfg.sliding_window + max_rows else cap; break :blk .{ .attn = .{ @@ -181,8 +198,6 @@ pub const Model = struct { .o = try table(l, ns, i, .o_proj), .qn = try vec(l, ns, i, .q_norm), .kn = try vec(l, ns, i, .k_norm), - .kc = try l.empty(@as(usize, ring) * NKV * HD * 2), - .vc = try l.empty(@as(usize, ring) * NKV * HD * 2), .ring = ring, .window = if (kind == .sliding_attention) cfg.sliding_window else 0, } }; @@ -197,7 +212,8 @@ pub const Model = struct { invFreq(c.sliding_rotary.theta, c.sliding_rotary.partial, &inv); try r.upload(inv_slide, std.mem.sliceAsBytes(&inv)); try r.sync(); - return .{ + var m: Model = .{ + .gpa = gpa, .r = r, .layers = layers, .emb = try table(l, ns, null, .embed_tokens), @@ -212,36 +228,128 @@ pub const Model = struct { .maxc = maxc, .device_bytes = l.total, }; + m.own = try m.newSeq(); + m.seq = m.own; + return m; } - /// Back to position 0: the conv states are zeroed; KV slots are only read once written. - pub fn reset(m: *Model) !void { - const zero: [6 * 256]u16 = @splat(0); - for (m.layers) |ly| switch (ly.mixer) { - .conv => |c| { - var off: usize = 0; - while (off < 6 * H * 2) : (off += zero.len * 2) { - try m.r.upload(.{ .rt = c.state.rt, .ptr = @ptrFromInt(@intFromPtr(c.state.ptr.?) + off), .len = zero.len * 2 }, std.mem.sliceAsBytes(&zero)); - } + /// A new stream state at position 0 (zeroed conv rows; KV slots are only read once written). + pub fn newSeq(m: *Model) !*Seq { + const q = try m.gpa.create(Seq); + errdefer m.gpa.destroy(q); + q.* = .{}; + for (m.layers, 0..) |ly, i| switch (ly.mixer) { + .conv => q.conv[i] = try m.r.alloc(conv_bytes), + .attn => |a| { + q.kc[i] = try m.r.alloc(@as(usize, a.ring) * slot_bytes); + q.vc[i] = try m.r.alloc(@as(usize, a.ring) * slot_bytes); }, - .attn => {}, + }; + try m.zeroConv(q); + return q; + } + + pub fn freeSeq(m: *Model, q: *Seq) void { + if (q == m.seq) { // `own` bound again with its own position, not the freed stream's + m.seq = m.own; + m.pos = m.own.pos; + } + for (m.layers, 0..) |ly, i| switch (ly.mixer) { + .conv => q.conv[i].free(), + .attn => { + q.kc[i].free(); + q.vc[i].free(); + }, + }; + m.gpa.destroy(q); + } + + /// Makes `q` the stream the next forwards read and write. + pub fn bind(m: *Model, q: *Seq) void { + m.seq.pos = m.pos; + m.seq = q; + m.pos = q.pos; + } + + fn zeroConv(m: *Model, q: *Seq) !void { + const zero: [6 * 256]u16 = @splat(0); + for (m.layers, 0..) |ly, i| if (ly.mixer == .conv) { + var off: usize = 0; + while (off < conv_bytes) : (off += zero.len * 2) try m.r.upload(from(q.conv[i], off), std.mem.sliceAsBytes(&zero)); }; try m.r.sync(); + } + + /// The bound stream back to position 0. + pub fn reset(m: *Model) !void { + try m.zeroConv(m.seq); m.pos = 0; } + /// Bytes of a kept state at position `at`: conv rows, `at` slots of each full layer, each sliding ring whole. + pub fn snapBytes(m: *const Model, at: u32) u64 { + var n: u64 = 0; + for (m.layers) |ly| n += switch (ly.mixer) { + .conv => conv_bytes, + .attn => |a| 2 * @as(u64, if (a.window == 0) at else a.ring) * slot_bytes, + }; + return n; + } + + /// One piece of a kept state at byte `off` of its buffer, to (save) or from it; the next piece's offset. + fn piece(m: *Model, sn: *const Snap, save: bool, b: Buf, bytes: usize, off: usize) !usize { + if (bytes > 0) { + if (save) try m.r.copy(from(sn.buf, off), b, bytes) else try m.r.copy(b, from(sn.buf, off), bytes); + } + return off + bytes; + } + + /// Copies the bound stream's state to (save) or from a kept state's buffer, on the device. + fn snapCopy(m: *Model, sn: *const Snap, save: bool) !void { + var off: usize = 0; + for (m.layers, 0..) |ly, i| switch (ly.mixer) { + .conv => off = try m.piece(sn, save, m.seq.conv[i], conv_bytes, off), + .attn => |a| { + const bytes = @as(usize, if (a.window == 0) sn.at else a.ring) * slot_bytes; + off = try m.piece(sn, save, m.seq.kc[i], bytes, off); + off = try m.piece(sn, save, m.seq.vc[i], bytes, off); + }, + }; + try m.r.sync(); + } + + /// The bound stream's state at its position, kept. + pub fn snapSave(m: *Model) !*Snap { + const sn = try m.gpa.create(Snap); + errdefer m.gpa.destroy(sn); + sn.* = .{ .at = m.pos, .buf = try m.r.alloc(@intCast(@max(m.snapBytes(m.pos), 1))) }; + try m.snapCopy(sn, true); + return sn; + } + + /// The bound stream continues from a kept state. + pub fn snapRestore(m: *Model, sn: *const Snap) !void { + m.pos = sn.at; + try m.snapCopy(sn, false); + } + + pub fn snapDrop(m: *Model, sn: *Snap) void { + sn.buf.free(); + m.gpa.destroy(sn); + } + /// y[r * y_stride + y_off + o] = x[r] . W[o] for n rows: the decode matvec for one row, pdense for a window. fn proj(m: *Model, t: Table, x: Buf, y: Buf, in_dim: u32, rows: u32, n: u32, y_stride: u32, y_off: u32) !void { if (n == 1) return run(&m.k.qmv, .{ (rows + 3) / 4, 1, 1 }, .{ t.w, t.s, t.b, x, y, in_dim, @as(u32, 0), y_off, rows }); try run(&m.k.pdense, .{ (rows + 15) / 16, (n + tile_rows - 1) / tile_rows, 1 }, .{ t.w, t.s, t.b, x, y, in_dim, rows, n, y_stride, y_off, @as(u32, 0) }); } - fn mixer(m: *Model, ly: *const Layer, n: u32) !void { + fn mixer(m: *Model, ly: *const Layer, li: usize, n: u32) !void { const s = &m.s; switch (ly.mixer) { .conv => |c| { try m.proj(c.in, s.xn, s.proj, H, 3 * H, n, 3 * H, 0); - try run(&m.k.conv, .{ H / 64, 1, 1 }, .{ s.proj, 3 * H, c.state, c.state, c.cw, s.att, n, H }); + try run(&m.k.conv, .{ H / 64, 1, 1 }, .{ s.proj, 3 * H, m.seq.conv[li], m.seq.conv[li], c.cw, s.att, n, H }); try m.proj(c.out, s.att, s.d, H, H, n, H, 0); }, .attn => |a| { @@ -251,11 +359,11 @@ pub const Model = struct { try m.proj(a.g, s.xn, s.g, H, NQ, n, NQ, 0); const inv = if (a.window == 0) m.inv_full else m.inv_slide; try run(&m.k.rope, .{ NQ, n, 1 }, .{ s.qkv, QKV, @as(u32, 0), s.q, NQ * HD, a.qn, inv, m.pos, @as(u32, 0), cfg.eps }); - try run(&m.k.rope, .{ NKV, n, 1 }, .{ s.qkv, QKV, NQ * HD, a.kc, NKV * HD, a.kn, inv, m.pos, a.ring, cfg.eps }); - try run(&m.k.vst, .{ NKV * HD / 64, n, 1 }, .{ s.qkv, QKV, (NQ + NKV) * HD, a.vc, NKV * HD, m.pos, a.ring }); + try run(&m.k.rope, .{ NKV, n, 1 }, .{ s.qkv, QKV, NQ * HD, m.seq.kc[li], NKV * HD, a.kn, inv, m.pos, a.ring, cfg.eps }); + try run(&m.k.vst, .{ NKV * HD / 64, n, 1 }, .{ s.qkv, QKV, (NQ + NKV) * HD, m.seq.vc[li], NKV * HD, m.pos, a.ring }); const span = if (a.window == 0) m.pos + n else @min(m.pos + n, a.window); const scale: f32 = 1.0 / @sqrt(@as(f32, HD)); - try run(&m.k.part, .{ NKV, (span + chunk - 1) / chunk, n }, .{ s.q, a.kc, a.vc, s.po, s.pm, s.pl, m.pos, chunk, a.ring, a.window, NKV, m.maxc, scale }); + try run(&m.k.part, .{ NKV, (span + chunk - 1) / chunk, n }, .{ s.q, m.seq.kc[li], m.seq.vc[li], s.po, s.pm, s.pl, m.pos, chunk, a.ring, a.window, NKV, m.maxc, scale }); try run(&m.k.merge, .{ NQ, n, 1 }, .{ s.po, s.pm, s.pl, s.g, s.att, m.pos, chunk, a.window, NQ, m.maxc }); try m.proj(a.o, s.att, s.d, NQ * HD, H, n, H, 0); }, @@ -281,7 +389,7 @@ pub const Model = struct { try run(&m.k.embed, .{ n, 1, 1 }, .{ m.emb.w, m.emb.s, m.emb.b, s.ids, s.x, H }); try run(&m.k.rms, .{ n, 1, 1 }, .{ s.x, m.layers[0].op_norm, s.xn, H, cfg.eps }); for (m.layers, 0..) |*ly, i| { - try m.mixer(ly, n); + try m.mixer(ly, i, n); try run(&m.k.pan, .{ n, 1, 1 }, .{ s.d, s.x, ly.post_op, ly.ffn_norm, s.xn, H, cfg.eps }); try m.proj(ly.w1, s.xn, s.a, H, FF, n, FF, 0); try m.proj(ly.w3, s.xn, s.b, H, FF, n, FF, 0); From db61ed59f720498e3e4b1e10dfe18a1f88a76e55 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 31 Oct 2026 13:58:34 +0100 Subject: [PATCH 09/14] xpu: LFM3 decode fuses w1, w3 and SwiGLU, the four attention matvecs, and RoPE with the KV store lfm3_mv.cl takes attn.cl's qmv_row verbatim (same compile options, so the same contraction) and runs w1 and w3 of a row with their SwiGLU in one launch, and q, k, v and the gate in one; lfm3_rope_store normalises and rotates the q and k heads and stores k and v in one launch, through the same per-head function lfm3_qk_norm_rope now calls. Both RoPE kernels require sub-groups of 16: their norm sums one partial a sub-group into part[4], which a compiler-chosen width of 8 would overrun and 32 would sum in another order. Every value is rounded as before: the greedy tokens of the three reference prompts are unchanged. Decode on Iris Xe: 15.2 to 17.2 tok/s (LFM3_FUSE selects the fusions for A/B timing). --- zig/build/xpu.zig | 1 + zig/kernels/xpu/lfm3.cl | 48 ++++++++++++++----- zig/kernels/xpu/lfm3_mv.cl | 71 +++++++++++++++++++++++++++++ zig/src/families/lfm3/xpu_model.zig | 42 +++++++++++++---- zig/src/xpu/kernels.zig | 1 + zig/tests/xpu/lfm3_bench.zig | 9 ++++ 6 files changed, 151 insertions(+), 21 deletions(-) create mode 100644 zig/kernels/xpu/lfm3_mv.cl diff --git a/zig/build/xpu.zig b/zig/build/xpu.zig index a4564edeff..3bc189b8f0 100644 --- a/zig/build/xpu.zig +++ b/zig/build/xpu.zig @@ -193,6 +193,7 @@ const kernels = [_]Kernel{ .{ .name = "nem_prow_gen12", .src = "nem_prow", .options = cl3 ++ " -DPO=4" }, .{ .name = "nem_attn_dec", .options = cl3 }, .{ .name = "lfm3", .options = cl3 }, + .{ .name = "lfm3_mv", .options = cl3 }, }; /// The runtime module; `with_kernels` false builds it host-only (empty images). diff --git a/zig/kernels/xpu/lfm3.cl b/zig/kernels/xpu/lfm3.cl index 18492a2eca..9a55972574 100644 --- a/zig/kernels/xpu/lfm3.cl +++ b/zig/kernels/xpu/lfm3.cl @@ -90,14 +90,10 @@ __kernel void lfm3_conv_rows(__global const ushort *proj, uint pd, __global cons for (uint k = 0; k < 6; k++) sout[k * D + c] = to_bf(h[k]); } -// Per-head RMSNorm, RoPE on pairs (d, d + 64), bf16 cos/sin; grid (heads, rows); to row r or slot (pos0 + r) % ring. -__kernel void lfm3_qk_norm_rope(__global const ushort *src, uint src_stride, uint src_off, __global ushort *dst, uint dst_stride, - __global const ushort *w, __global const float *inv_freq, uint pos0, uint ring, float eps) { - __local float part[4]; - const uint head = get_group_id(0); - const uint r = get_group_id(1); - const uint i = get_local_id(0); - __global const ushort *s = src + r * src_stride + src_off + head * HD; +// One head of row r at position pos: RMSNorm over 128 dims, RoPE on pairs (d, d + 64) with bf16 cos/sin; 64 items. +inline void norm_rope_head(__global const ushort *s, __global ushort *o, __global const ushort *w, __global const float *inv_freq, uint pos, + float eps, __local float *part) { + const uint i = get_local_id(0); // 64 items in 4 sub-groups of 16 (the kernels require 16): part[4] const float a = bf(s[i]), b = bf(s[i + 64]); float ss = sub_group_reduce_add(a * a + b * b); if (get_sub_group_local_id() == 0) part[get_sub_group_id()] = ss; @@ -108,14 +104,44 @@ __kernel void lfm3_qk_norm_rope(__global const ushort *src, uint src_stride, uin const float rs = rsqrt(t / (float)HD + eps); const float na = bfr(bf(w[i]) * bfr(a * rs)); const float nb = bfr(bf(w[i + 64]) * bfr(b * rs)); - const float f = (float)(pos0 + r) * inv_freq[i]; + const float f = (float)pos * inv_freq[i]; const float co = bfr(cos(f)), si = bfr(sin(f)); - const uint row = ring == 0 ? r : (pos0 + r) % ring; - __global ushort *o = dst + row * dst_stride + head * HD; o[i] = to_bf(na * co - nb * si); o[i + 64] = to_bf(nb * co + na * si); } +// Per-head RMSNorm, RoPE on pairs (d, d + 64), bf16 cos/sin; grid (heads, rows); to row r or slot (pos0 + r) % ring. +__attribute__((intel_reqd_sub_group_size(16))) +__kernel void lfm3_qk_norm_rope(__global const ushort *src, uint src_stride, uint src_off, __global ushort *dst, uint dst_stride, + __global const ushort *w, __global const float *inv_freq, uint pos0, uint ring, float eps) { + __local float part[4]; + const uint head = get_group_id(0); + const uint r = get_group_id(1); + const uint row = ring == 0 ? r : (pos0 + r) % ring; + norm_rope_head(src + r * src_stride + src_off + head * HD, dst + row * dst_stride + head * HD, w, inv_freq, pos0 + r, eps, part); +} + +// q heads, then k heads and v rows into slot (pos0 + r) % ring, in one launch; grid (nq + nkv + nkv * 2, rows). +__attribute__((intel_reqd_sub_group_size(16))) +__kernel void lfm3_rope_store(__global const ushort *qkv, uint stride, __global ushort *q, __global ushort *kc, __global ushort *vc, + __global const ushort *qw, __global const ushort *kw, __global const float *inv_freq, uint nq, uint nkv, + uint pos0, uint ring, float eps) { + __local float part[4]; + const uint g = get_group_id(0); + const uint r = get_group_id(1); + const uint slot = (pos0 + r) % ring; + __global const ushort *src = qkv + r * stride; + if (g < nq) { + norm_rope_head(src + g * HD, q + (r * nq + g) * HD, qw, inv_freq, pos0 + r, eps, part); + } else if (g < nq + nkv) { + const uint h = g - nq; + norm_rope_head(src + (nq + h) * HD, kc + (slot * nkv + h) * HD, kw, inv_freq, pos0 + r, eps, part); + } else { + const uint e = (g - nq - nkv) * 64 + get_local_id(0); + vc[slot * nkv * HD + e] = src[(nq + nkv) * HD + e]; + } +} + // Copies n_elem bf16 of row r to KV slot (pos0 + r) % ring; grid (ceil(n_elem / 64), rows). __kernel void lfm3_v_store(__global const ushort *src, uint src_stride, uint src_off, __global ushort *dst, uint n_elem, uint pos0, uint ring) { const uint e = get_global_id(0); diff --git a/zig/kernels/xpu/lfm3_mv.cl b/zig/kernels/xpu/lfm3_mv.cl new file mode 100644 index 0000000000..67b19e4709 --- /dev/null +++ b/zig/kernels/xpu/lfm3_mv.cl @@ -0,0 +1,71 @@ +// LFM3 decode matvecs fused by launch: attn.cl's qmv_row unchanged (same contraction), more per launch. +inline float bf(ushort v) { return as_float((uint)v << 16); } +inline ushort to_bf(float f) { + uint u = as_uint(f); + u += 0x7fffu + ((u >> 16) & 1u); + return (ushort)(u >> 16); +} + +// attn.cl's qmv_row, verbatim: one 16-lane sub-group a row, the sum valid on every lane. +inline float qmv_row(__global const uint *w, __global const ushort *scales, __global const ushort *biases, + __global const ushort *x, uint row, uint in_dim) { + const uint lane = get_sub_group_local_id(); + const uint words = in_dim / 8; + const uint groups = in_dim / 64; + float acc = 0.0f; + for (uint wi = lane; wi < words; wi += 16) { + const uint g = wi / 8; + const float sc = bf(scales[row * groups + g]); + const float bi = bf(biases[row * groups + g]); + const uint pk = w[row * words + wi]; + const float8 xv = as_float8(convert_uint8(vload8(wi, x)) << 16); + float dot = 0.0f, sx = 0.0f; + dot += xv.s0 * (float)(pk & 15u); sx += xv.s0; + dot += xv.s1 * (float)((pk >> 4) & 15u); sx += xv.s1; + dot += xv.s2 * (float)((pk >> 8) & 15u); sx += xv.s2; + dot += xv.s3 * (float)((pk >> 12) & 15u); sx += xv.s3; + dot += xv.s4 * (float)((pk >> 16) & 15u); sx += xv.s4; + dot += xv.s5 * (float)((pk >> 20) & 15u); sx += xv.s5; + dot += xv.s6 * (float)((pk >> 24) & 15u); sx += xv.s6; + dot += xv.s7 * (float)((pk >> 28) & 15u); sx += xv.s7; + acc += sc * dot + bi * sx; + } + return sub_group_reduce_add(acc); +} + +// Four tables in one launch (q, k, v, gate): rows r0 .. r3 into y0 + o0 .. y3 + o3; 4 rows a work-group. +__attribute__((intel_reqd_sub_group_size(16))) +__kernel void qmv4_bf_x4(__global const uint *w0, __global const ushort *s0, __global const ushort *b0, __global ushort *y0, uint r0, uint o0, + __global const uint *w1, __global const ushort *s1, __global const ushort *b1, __global ushort *y1, uint r1, uint o1, + __global const uint *w2, __global const ushort *s2, __global const ushort *b2, __global ushort *y2, uint r2, uint o2, + __global const uint *w3, __global const ushort *s3, __global const ushort *b3, __global ushort *y3, uint r3, uint o3, + __global const ushort *x, uint in_dim) { + uint row = get_group_id(0) * 4 + get_sub_group_id(); + float acc; + __global ushort *y; + if (row < r0) { + acc = qmv_row(w0, s0, b0, x, row, in_dim); + y = y0 + o0; + } else if ((row -= r0) < r1) { + acc = qmv_row(w1, s1, b1, x, row, in_dim); + y = y1 + o1; + } else if ((row -= r1) < r2) { + acc = qmv_row(w2, s2, b2, x, row, in_dim); + y = y2 + o2; + } else if ((row -= r2) < r3) { + acc = qmv_row(w3, s3, b3, x, row, in_dim); + y = y3 + o3; + } else return; + if (get_sub_group_local_id() == 0) y[row] = to_bf(acc); +} + +// w1 and w3 rows and their SwiGLU: act = bf16(bf16(silu(bf16(w1 x))) * bf16(w3 x)), as qmv4_bf + lfm3_swiglu. +__attribute__((intel_reqd_sub_group_size(16))) +__kernel void qmv4_swiglu(__global const uint *w1, __global const ushort *s1, __global const ushort *b1, __global const uint *w3, __global const ushort *s3, + __global const ushort *b3, __global const ushort *x, __global ushort *act, uint in_dim, uint rows) { + const uint row = get_group_id(0) * 4 + get_sub_group_id(); + if (row >= rows) return; + const float a = bf(to_bf(qmv_row(w1, s1, b1, x, row, in_dim))); + const float u = bf(to_bf(qmv_row(w3, s3, b3, x, row, in_dim))); + if (get_sub_group_local_id() == 0) act[row] = to_bf(bf(to_bf(a / (1.0f + exp(-a)))) * u); +} diff --git a/zig/src/families/lfm3/xpu_model.zig b/zig/src/families/lfm3/xpu_model.zig index d96c828a71..53e299898f 100644 --- a/zig/src/families/lfm3/xpu_model.zig +++ b/zig/src/families/lfm3/xpu_model.zig @@ -43,6 +43,9 @@ const Kernels = struct { part: rt.Kernel, merge: rt.Kernel, swiglu: rt.Kernel, + qkvg: rt.Kernel, + w13: rt.Kernel, + rope_store: rt.Kernel, }; /// One stream's state: each conv layer's last six B*x rows, each attention layer's KV slots, and its position. @@ -102,6 +105,8 @@ pub const Model = struct { /// The bound stream's state (`own` until another is bound); `pos` is the bound stream's position. seq: *Seq = undefined, own: *Seq = undefined, + /// Fused launches by bit (1 w1+w3+SwiGLU, 2 q/k/v/gate, 4 rope and KV store); LFM3_FUSE sets them for A/B. + fuse: u32 = 7, /// The affine tensors of module `m` (in layer `layer`, or a model tensor) by names.zig in namespace `ns`. fn table(l: *ld.Loader, ns: names.Namespace, layer: ?usize, m: names.Module) !Table { @@ -134,6 +139,7 @@ pub const Model = struct { var am = try r.module(xpu.kernels.attn); var pm = try r.module(if (gen12) xpu.kernels.nem_prow_gen12 else xpu.kernels.nem_prow); var fm = try r.module(xpu.kernels.lfm3); + var vm = try r.module(xpu.kernels.lfm3_mv); const k: Kernels = .{ .embed = try bm.kernel("embed4", .{ 64, 1, 1 }), .qmv = try am.kernel("qmv4_bf", .{ 64, 1, 1 }), @@ -149,6 +155,9 @@ pub const Model = struct { .part = try fm.kernel("lfm3_attn_partial", .{ 64, 1, 1 }), .merge = try fm.kernel("lfm3_attn_merge_gate", .{ 16, 1, 1 }), .swiglu = try fm.kernel("lfm3_swiglu", .{ 64, 1, 1 }), + .qkvg = try vm.kernel("qmv4_bf_x4", .{ 64, 1, 1 }), + .w13 = try vm.kernel("qmv4_swiglu", .{ 64, 1, 1 }), + .rope_store = try fm.kernel("lfm3_rope_store", .{ 64, 1, 1 }), }; const maxc = (cap + chunk - 1) / chunk; const w: usize = max_rows; @@ -228,6 +237,7 @@ pub const Model = struct { .maxc = maxc, .device_bytes = l.total, }; + if (std.c.getenv("LFM3_FUSE")) |v| m.fuse = std.fmt.parseInt(u32, std.mem.span(v), 10) catch 7; m.own = try m.newSeq(); m.seq = m.own; return m; @@ -353,14 +363,22 @@ pub const Model = struct { try m.proj(c.out, s.att, s.d, H, H, n, H, 0); }, .attn => |a| { - try m.proj(a.q, s.xn, s.qkv, H, NQ * HD, n, QKV, 0); - try m.proj(a.k, s.xn, s.qkv, H, NKV * HD, n, QKV, NQ * HD); - try m.proj(a.v, s.xn, s.qkv, H, NKV * HD, n, QKV, (NQ + NKV) * HD); - try m.proj(a.g, s.xn, s.g, H, NQ, n, NQ, 0); + if (n == 1 and m.fuse & 2 != 0) { // the four decode matvecs in one launch, each row as qmv4_bf sums it + try run(&m.k.qkvg, .{ (QKV + NQ + 3) / 4, 1, 1 }, .{ a.q.w, a.q.s, a.q.b, s.qkv, NQ * HD, @as(u32, 0), a.k.w, a.k.s, a.k.b, s.qkv, NKV * HD, NQ * HD, a.v.w, a.v.s, a.v.b, s.qkv, NKV * HD, (NQ + NKV) * HD, a.g.w, a.g.s, a.g.b, s.g, NQ, @as(u32, 0), s.xn, H }); + } else { + try m.proj(a.q, s.xn, s.qkv, H, NQ * HD, n, QKV, 0); + try m.proj(a.k, s.xn, s.qkv, H, NKV * HD, n, QKV, NQ * HD); + try m.proj(a.v, s.xn, s.qkv, H, NKV * HD, n, QKV, (NQ + NKV) * HD); + try m.proj(a.g, s.xn, s.g, H, NQ, n, NQ, 0); + } const inv = if (a.window == 0) m.inv_full else m.inv_slide; - try run(&m.k.rope, .{ NQ, n, 1 }, .{ s.qkv, QKV, @as(u32, 0), s.q, NQ * HD, a.qn, inv, m.pos, @as(u32, 0), cfg.eps }); - try run(&m.k.rope, .{ NKV, n, 1 }, .{ s.qkv, QKV, NQ * HD, m.seq.kc[li], NKV * HD, a.kn, inv, m.pos, a.ring, cfg.eps }); - try run(&m.k.vst, .{ NKV * HD / 64, n, 1 }, .{ s.qkv, QKV, (NQ + NKV) * HD, m.seq.vc[li], NKV * HD, m.pos, a.ring }); + if (m.fuse & 4 != 0) { // q, k and v of every row in one launch + try run(&m.k.rope_store, .{ NQ + NKV + NKV * HD / 64, n, 1 }, .{ s.qkv, QKV, s.q, m.seq.kc[li], m.seq.vc[li], a.qn, a.kn, inv, NQ, NKV, m.pos, a.ring, cfg.eps }); + } else { + try run(&m.k.rope, .{ NQ, n, 1 }, .{ s.qkv, QKV, @as(u32, 0), s.q, NQ * HD, a.qn, inv, m.pos, @as(u32, 0), cfg.eps }); + try run(&m.k.rope, .{ NKV, n, 1 }, .{ s.qkv, QKV, NQ * HD, m.seq.kc[li], NKV * HD, a.kn, inv, m.pos, a.ring, cfg.eps }); + try run(&m.k.vst, .{ NKV * HD / 64, n, 1 }, .{ s.qkv, QKV, (NQ + NKV) * HD, m.seq.vc[li], NKV * HD, m.pos, a.ring }); + } const span = if (a.window == 0) m.pos + n else @min(m.pos + n, a.window); const scale: f32 = 1.0 / @sqrt(@as(f32, HD)); try run(&m.k.part, .{ NKV, (span + chunk - 1) / chunk, n }, .{ s.q, m.seq.kc[li], m.seq.vc[li], s.po, s.pm, s.pl, m.pos, chunk, a.ring, a.window, NKV, m.maxc, scale }); @@ -391,9 +409,13 @@ pub const Model = struct { for (m.layers, 0..) |*ly, i| { try m.mixer(ly, i, n); try run(&m.k.pan, .{ n, 1, 1 }, .{ s.d, s.x, ly.post_op, ly.ffn_norm, s.xn, H, cfg.eps }); - try m.proj(ly.w1, s.xn, s.a, H, FF, n, FF, 0); - try m.proj(ly.w3, s.xn, s.b, H, FF, n, FF, 0); - try run(&m.k.swiglu, .{ (n * FF + 63) / 64, 1, 1 }, .{ s.a, s.b, s.act, n * FF }); + if (n == 1 and m.fuse & 1 != 0) { // w1, w3 and their SwiGLU in one launch, each value as the three launches round it + try run(&m.k.w13, .{ FF / 4, 1, 1 }, .{ ly.w1.w, ly.w1.s, ly.w1.b, ly.w3.w, ly.w3.s, ly.w3.b, s.xn, s.act, H, FF }); + } else { + try m.proj(ly.w1, s.xn, s.a, H, FF, n, FF, 0); + try m.proj(ly.w3, s.xn, s.b, H, FF, n, FF, 0); + try run(&m.k.swiglu, .{ (n * FF + 63) / 64, 1, 1 }, .{ s.a, s.b, s.act, n * FF }); + } try m.proj(ly.w2, s.act, s.d, FF, H, n, H, 0); const next = if (i + 1 < m.layers.len) m.layers[i + 1].op_norm else m.norm_f; try run(&m.k.pan, .{ n, 1, 1 }, .{ s.d, s.x, ly.post_ffn, next, s.xn, H, cfg.eps }); diff --git a/zig/src/xpu/kernels.zig b/zig/src/xpu/kernels.zig index 696d1a8cd4..93178cba5e 100644 --- a/zig/src/xpu/kernels.zig +++ b/zig/src/xpu/kernels.zig @@ -28,3 +28,4 @@ pub const nem_prow: []const u8 = if (available) &Blob("spv_nem_prow").bytes else pub const nem_prow_gen12: []const u8 = if (available) &Blob("spv_nem_prow_gen12").bytes else &.{}; pub const nem_attn_dec: []const u8 = if (available) &Blob("spv_nem_attn_dec").bytes else &.{}; pub const lfm3: []const u8 = if (available) &Blob("spv_lfm3").bytes else &.{}; +pub const lfm3_mv: []const u8 = if (available) &Blob("spv_lfm3_mv").bytes else &.{}; diff --git a/zig/tests/xpu/lfm3_bench.zig b/zig/tests/xpu/lfm3_bench.zig index 4ced6e492f..8f40ebbcea 100644 --- a/zig/tests/xpu/lfm3_bench.zig +++ b/zig/tests/xpu/lfm3_bench.zig @@ -42,5 +42,14 @@ pub fn main(init: std.process.Init) !void { tok = try m.argmax(); } const dec = @as(f64, @floatFromInt(nowNs() - t0)) / 1e9; + if (std.c.getenv("LFM3_PROF") != null) { // 16 more steps with a sync around every launch: device time by kernel + @import("xpu").rt.prof_on = true; + for (0..16) |_| { + try m.forward(&.{tok}, true, null); + tok = try m.argmax(); + } + @import("xpu").rt.prof_on = false; + @import("xpu").rt.profDump("16 decode steps"); + } std.debug.print("prompt {d} tokens in windows of {d}: prefill {d:.2} s ({d:.1} tok/s); decode after it {d:.1} tok/s; device {d:.0} MiB\n", .{ n, w, pre, @as(f64, @floatFromInt(n)) / pre, 64.0 / dec, @as(f64, @floatFromInt(m.device_bytes)) / (1 << 20) }); } From 0c25bf7709a2c49c3608fd01eea7be6cd067ae04 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 31 Oct 2026 11:29:45 +0100 Subject: [PATCH 10/14] xpu: LFM3 prompts in windows of 128 rows On Iris Xe a prompt window of 128 rows beats 32 and 64 with the 8 x 4 prompt tiles (1,024-token prefill 84 against 78 tok/s at 64). The greedy tokens of the three reference prompts are unchanged at windows of 128 and of 64. --- zig/src/native/xpu_lfm3.zig | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zig/src/native/xpu_lfm3.zig b/zig/src/native/xpu_lfm3.zig index 007aa005f9..9366e4e54e 100644 --- a/zig/src/native/xpu_lfm3.zig +++ b/zig/src/native/xpu_lfm3.zig @@ -10,8 +10,8 @@ pub const formats: []const []const u8 = &.{"mlx-q4g64"}; pub const name = "xpu-lfm3"; pub const Engine = lfm3.engine.Engine; -/// Rows of one prompt window (pdense above one row; the scratch is sized for it). -const prompt_rows: u32 = 64; +/// Rows of one prompt window (pdense above one row; the scratch is sized for it): 128 beat 32 and 64 on Iris Xe. +const prompt_rows: u32 = 128; pub fn promptRows(_: *const Engine) u32 { return prompt_rows; From 3b5c257a028c166c189932c13bf3a6a3f6f074e1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sat, 31 Oct 2026 14:42:43 +0100 Subject: [PATCH 11/14] xpu: LFM3 as a copy-drafting engine of the lane host: a sequence a stream, verified windows, shared forwards LFM3's engine takes the contract Gemma 4 E4B's takes on the XPU lane host (copy_only): one sequence a stream in up to 8 slots, the lane core's context copies verified in windows of up to 16 rows (verifyRows, keepRows), several streams' windows in one forward (verifyMulti, slotOf), window costs timed at load (1, 2, 4, 8 and 16 rows, the widths between interpolated: 8 s to load instead of 21 s, the same round plans), and prompt states kept anywhere (cache_rules). One line in the adapter (xpu_nem.zig): the host takes a family's copy_match when its table row declares one, before XPU_COPY_MATCH. LFM3's is 8: at the core's 4, copies cost a short chat 6% on Iris Xe; at 8, prose and code decode 65% and 38% faster than with --no-drafts, and the two chats within 2% of it (session noise). A window's rows sum as the decode steps do: its projections and head run on pdense, whose bf16 and fp32 outputs equal qmv4_bf and qmv4_f32 bit for bit (tf-xpu-lfm3_rows-test). Conv and attention run on each stream's own state and position; lfm3_conv_hist also records each row's B * x, so keeping the first c rows of a window restores the conv state those rows leave. Decode ids are kernel arguments, so no host memory the caller may reuse is read later. tf-xpu-lfm3_streams-test checks streams in shared forwards, freed and joined mid-way, and windows rolled back to their kept rows, against each stream alone. --- zig/build/xpu.zig | 2 +- zig/kernels/xpu/lfm3.cl | 25 ++++ zig/src/families/lfm3/xpu_engine.zig | 93 +++++++++++- zig/src/families/lfm3/xpu_model.zig | 169 ++++++++++++++++++---- zig/src/native/xpu_lfm3.zig | 8 +- zig/src/native/xpu_nem.zig | 1 + zig/tests/xpu/lfm3_rows.zig | 65 +++++++++ zig/tests/xpu/lfm3_streams.zig | 204 +++++++++++++++++++++++++++ 8 files changed, 528 insertions(+), 39 deletions(-) create mode 100644 zig/tests/xpu/lfm3_rows.zig create mode 100644 zig/tests/xpu/lfm3_streams.zig diff --git a/zig/build/xpu.zig b/zig/build/xpu.zig index 3bc189b8f0..4981242a4d 100644 --- a/zig/build/xpu.zig +++ b/zig/build/xpu.zig @@ -266,7 +266,7 @@ fn lfm3Module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bui /// The standalone op and model tests, each its own program `tf-xpu--test` (step `xpu-tests`). const test_programs = [_][]const u8{ - "panic", "nem_rows", "nem_prefill", "nem_gen", "nem_mtp", "nem_kl", "nem_norm_bench", "nem_mv_bench", "nem_slot", "nem_format", "nem_seam", "nem_attn_dec", "nem_fp64", "memory", "lfm3_conv", "lfm3_rope", "lfm3_attn", "lfm3_norm", "lfm3_fwd", "lfm3_bench", + "panic", "nem_rows", "nem_prefill", "nem_gen", "nem_mtp", "nem_kl", "nem_norm_bench", "nem_mv_bench", "nem_slot", "nem_format", "nem_seam", "nem_attn_dec", "nem_fp64", "memory", "lfm3_conv", "lfm3_rope", "lfm3_attn", "lfm3_norm", "lfm3_fwd", "lfm3_bench", "lfm3_rows", "lfm3_streams", }; /// `-Dxpu` (Linux): Nemotron's SPIR-V images, `tensorfold-xpu` and `tf-xpu-test`; needs ocloc, not nvcc or CUDA. diff --git a/zig/kernels/xpu/lfm3.cl b/zig/kernels/xpu/lfm3.cl index 9a55972574..e5144072ea 100644 --- a/zig/kernels/xpu/lfm3.cl +++ b/zig/kernels/xpu/lfm3.cl @@ -90,6 +90,31 @@ __kernel void lfm3_conv_rows(__global const ushort *proj, uint pd, __global cons for (uint k = 0; k < 6; k++) sout[k * D + c] = to_bf(h[k]); } +// lfm3_conv_rows for a verified window, also ext rows 0..5 = the state before it, 6 + r = row r's B * x (rollback). +__kernel void lfm3_conv_hist(__global const ushort *proj, uint pd, __global const ushort *sin, __global ushort *sout, __global const ushort *cw, + __global ushort *y, uint n, uint D, __global ushort *ext) { + const uint c = get_global_id(0); + if (c >= D) return; + float h[6], w[7]; + for (uint k = 0; k < 6; k++) { + h[k] = bf(sin[k * D + c]); + ext[k * D + c] = to_bf(h[k]); + } + for (uint k = 0; k < 7; k++) w[k] = bf(cw[c * 7 + k]); + for (uint r = 0; r < n; r++) { + __global const ushort *p = proj + r * pd; + const float bx = bfr(bf(p[c]) * bf(p[2 * D + c])); + float acc = 0.0f; + for (uint k = 0; k < 6; k++) acc += h[k] * w[k]; + acc += bx * w[6]; + y[r * D + c] = to_bf(bf(p[D + c]) * bfr(acc)); + ext[(6 + r) * D + c] = to_bf(bx); + for (uint k = 0; k < 5; k++) h[k] = h[k + 1]; + h[5] = bx; + } + for (uint k = 0; k < 6; k++) sout[k * D + c] = to_bf(h[k]); +} + // One head of row r at position pos: RMSNorm over 128 dims, RoPE on pairs (d, d + 64) with bf16 cos/sin; 64 items. inline void norm_rope_head(__global const ushort *s, __global ushort *o, __global const ushort *w, __global const float *inv_freq, uint pos, float eps, __local float *part) { diff --git a/zig/src/families/lfm3/xpu_engine.zig b/zig/src/families/lfm3/xpu_engine.zig index 848d235970..6b28bb97bd 100644 --- a/zig/src/families/lfm3/xpu_engine.zig +++ b/zig/src/families/lfm3/xpu_engine.zig @@ -1,4 +1,4 @@ -//! LFM3 behind the Intel GPU lane host: one sequence on the device, the prompt in windows, greedy argmax on the device. +//! LFM3 behind the Intel GPU lane host: a sequence a stream, windows verified bit for bit the steps. const std = @import("std"); const xpu = @import("xpu"); const cfg = @import("config.zig"); @@ -9,6 +9,9 @@ pub const Options = struct { context: usize = 4096, window_rows: u32 = 64 }; pub const Engine = struct { pub const Sequence = model.Seq; pub const Snapshot = model.Snap; + /// The lane core drafts by copying the context (no head): windows of up to 16 rows, bit for bit the steps. + pub const copy_only = true; + head: ?u8 = null, gpa: std.mem.Allocator, stream: xpu.Stream, r: xpu.rt.Runtime, @@ -20,6 +23,8 @@ pub const Engine = struct { load_seconds: f64, /// The first stream's state, held here so the lane host can tell it from the streams it adds (newSeq). own: Sequence, + /// Each slot's sequence (0 is `own`); verifyMulti names streams by slot. + slots: [model.max_streams]?*Sequence, pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const xpu.Context, dir: []const u8, o: Options) !*Engine { const t0 = std.Io.Clock.awake.now(io); @@ -42,6 +47,9 @@ pub const Engine = struct { gpa.destroy(e.m.own); e.m.own = &e.own; e.m.seq = &e.own; + e.head = null; + e.slots = @splat(null); + e.slots[0] = &e.own; e.max_len = o.context; e.load_seconds = xpu.decode.seconds(io, t0); return e; @@ -84,13 +92,24 @@ pub const Engine = struct { } pub fn newSeq(e: *Engine) !*Sequence { - return e.m.newSeq(); + const slot = for (e.slots[1..], 1..) |q, i| { + if (q == null) break i; + } else return error.NoFreeSlot; + const q = try e.m.newSeq(); + q.slot = @intCast(slot); + e.slots[slot] = q; + return q; } pub fn freeSeq(e: *Engine, q: *Sequence) void { + e.slots[q.slot] = null; e.m.freeSeq(q); } + pub fn slotOf(_: *const Engine, q: *const Sequence) u32 { + return q.slot; + } + pub fn bind(e: *Engine, q: *Sequence) void { e.m.bind(q); } @@ -107,15 +126,75 @@ pub const Engine = struct { } } - /// The bound stream's next token after its last forward with logits (one row). + /// The pending token and drafts of the bound stream as one window; every row keeps its logits and argmax. + pub fn verifyRows(e: *Engine, ids: []const u32) !void { + try e.m.forwardSegs(&.{.{ .q = e.m.seq, .ids = ids }}, .rows, true, null); + } + + /// Windows of several streams (by slot) in one forward over the weights, each row at its stream's position. + pub fn verifyMulti(e: *Engine, slots: []const u32, windows: []const []const u32) !void { + var segs: [model.max_streams]model.Seg = undefined; + if (slots.len > segs.len) return error.WindowTooWide; + for (slots, windows, segs[0..slots.len]) |sl, ids, *g| g.* = .{ .q = e.slots[sl] orelse return error.UnknownStream, .ids = ids }; + try e.m.forwardSegs(segs[0..slots.len], .rows, true, null); + } + + /// Keeps the first c rows of the bound stream's last window. + pub fn keepRows(e: *Engine, c: u32) !void { + try e.m.keepRows(c); + } + + /// The last window's greedy tokens, a row each (all streams' rows in order after verifyMulti). pub fn rowTokens(e: *Engine, out: []i32) !void { - if (out.len != 1) return error.WindowTooWide; - out[0] = @intCast(try e.m.argmax()); + try e.m.rowTokens(out); } pub fn rowLogits(e: *Engine, out: []f32) !void { - if (out.len != cfg.vocab) return error.WindowTooWide; - try e.m.fetchLogits(out); + try e.m.rowLogits(out); + } + + /// No head: the lane core's copies are the drafts. + pub fn draftKept(_: *Engine, _: []const u32, _: u32, _: u32, _: []u32) !usize { + return 0; + } + + pub fn draftFirst(_: *Engine, _: u32, _: u32, _: []u32) !usize { + return 0; + } + + /// Window ms by rows on a scratch sequence: 1, 2, 4, 8, 16 timed (best of 3), the rest interpolated; no head: 0. + pub fn timeRounds(e: *Engine, ms: []f64) !f64 { + const prev = e.m.seq; + const q = try e.m.newSeq(); + e.m.bind(q); + defer { + e.m.bind(prev); + e.m.freeSeq(q); + } + const ids: [model.window_rows]u32 = @splat(1); + const timed = [_]usize{ 1, 2, 4, 8, model.window_rows }; + std.debug.assert(ms.len == model.window_rows); + for (timed) |w| { + var best: f64 = std.math.inf(f64); + for (0..3) |_| { + try e.m.reset(); + try e.r.sync(); + var ts0: std.c.timespec = undefined; + _ = std.c.clock_gettime(.MONOTONIC, &ts0); + try e.m.forwardSegs(&.{.{ .q = q, .ids = ids[0..w] }}, .rows, true, null); + try e.r.sync(); + var ts1: std.c.timespec = undefined; + _ = std.c.clock_gettime(.MONOTONIC, &ts1); + best = @min(best, @as(f64, @floatFromInt(ts1.sec - ts0.sec)) * 1e3 + @as(f64, @floatFromInt(ts1.nsec - ts0.nsec)) / 1e6); + } + ms[w - 1] = best; + } + for (timed[0 .. timed.len - 1], timed[1..]) |a, b| for (a + 1..b) |w| { + const f = @as(f64, @floatFromInt(w - a)) / @as(f64, @floatFromInt(b - a)); + ms[w - 1] = ms[a - 1] + f * (ms[b - 1] - ms[a - 1]); + }; + try e.m.reset(); + return 0; } /// One decode step of the bound stream, with its logits. diff --git a/zig/src/families/lfm3/xpu_model.zig b/zig/src/families/lfm3/xpu_model.zig index 53e299898f..92099105b7 100644 --- a/zig/src/families/lfm3/xpu_model.zig +++ b/zig/src/families/lfm3/xpu_model.zig @@ -30,6 +30,8 @@ const Layer = struct { op_norm: Buf, post_op: Buf, ffn_norm: Buf, post_ffn: Buf, const Kernels = struct { embed: rt.Kernel, + embed_tok: rt.Kernel, + conv_hist: rt.Kernel, qmv: rt.Kernel, head: rt.Kernel, am_part: rt.Kernel, @@ -49,12 +51,34 @@ const Kernels = struct { }; /// One stream's state: each conv layer's last six B*x rows, each attention layer's KV slots, and its position. -pub const Seq = struct { conv: [cfg.layers]Buf = undefined, kc: [cfg.layers]Buf = undefined, vc: [cfg.layers]Buf = undefined, pos: u32 = 0 }; +pub const Seq = struct { + conv: [cfg.layers]Buf = undefined, + kc: [cfg.layers]Buf = undefined, + vc: [cfg.layers]Buf = undefined, + pos: u32 = 0, + /// Each conv layer's state before the last verified window, then its rows' B * x: keepRows' source. + ext: [cfg.layers]Buf = undefined, + /// The last verified window still to keep: its first position and rows (0: none). + win_start: u32 = 0, + win_rows: u32 = 0, + /// The engine's slot for this stream (0 is the engine's own). + slot: u32 = 0, +}; + +/// Rows of a verified window at most (all streams' rows of a shared forward), and streams a shared forward takes. +pub const window_rows: u32 = 16; +pub const max_streams: u32 = 8; + +/// A stream's rows in one forward: `ids` at the stream's own position (its window, or one row of several streams). +pub const Seg = struct { q: *Seq, ids: []const u32 }; +/// What a forward leaves: nothing, the last row's logits and argmax (a prompt), or every row's (verified windows). +pub const Out = enum { none, last, rows }; /// A kept prompt state at `at`: conv rows, the full layers' first `at` KV slots, the sliding rings, in one buffer. pub const Snap = struct { at: u32, buf: Buf }; const conv_bytes: usize = 6 * H * 2; +const ext_bytes: usize = (6 + window_rows) * H * 2; const slot_bytes: usize = NKV * HD * 2; /// `b` from byte `off` on. @@ -142,6 +166,8 @@ pub const Model = struct { var vm = try r.module(xpu.kernels.lfm3_mv); const k: Kernels = .{ .embed = try bm.kernel("embed4", .{ 64, 1, 1 }), + .embed_tok = try bm.kernel("embed4_tok", .{ 64, 1, 1 }), + .conv_hist = try fm.kernel("lfm3_conv_hist", .{ 64, 1, 1 }), .qmv = try am.kernel("qmv4_bf", .{ 64, 1, 1 }), .head = try am.kernel("qmv4_f32", .{ 64, 1, 1 }), .am_part = try am.kernel("argmax_partial", .{ 256, 1, 1 }), @@ -177,9 +203,9 @@ pub const Model = struct { .po = try l.empty(w * NQ * maxc * HD * 4), .pm = try l.empty(w * NQ * maxc * 4), .pl = try l.empty(w * NQ * maxc * 4), - .logits = try l.empty(@as(usize, V) * 4), - .am_scratch = try l.empty(argmax_parts * 8), - .am_out = try l.empty(4), + .logits = try l.empty(@as(usize, window_rows) * V * 4), + .am_scratch = try l.empty(window_rows * argmax_parts * 8), + .am_out = try l.empty(window_rows * 4), }; const ns = try namespace(l); const layers = try gpa.alloc(Layer, cfg.layers); @@ -249,7 +275,10 @@ pub const Model = struct { errdefer m.gpa.destroy(q); q.* = .{}; for (m.layers, 0..) |ly, i| switch (ly.mixer) { - .conv => q.conv[i] = try m.r.alloc(conv_bytes), + .conv => { + q.conv[i] = try m.r.alloc(conv_bytes); + q.ext[i] = try m.r.alloc(ext_bytes); + }, .attn => |a| { q.kc[i] = try m.r.alloc(@as(usize, a.ring) * slot_bytes); q.vc[i] = try m.r.alloc(@as(usize, a.ring) * slot_bytes); @@ -265,7 +294,10 @@ pub const Model = struct { m.pos = m.own.pos; } for (m.layers, 0..) |ly, i| switch (ly.mixer) { - .conv => q.conv[i].free(), + .conv => { + q.conv[i].free(); + q.ext[i].free(); + }, .attn => { q.kc[i].free(); q.vc[i].free(); @@ -354,12 +386,22 @@ pub const Model = struct { try run(&m.k.pdense, .{ (rows + 15) / 16, (n + tile_rows - 1) / tile_rows, 1 }, .{ t.w, t.s, t.b, x, y, in_dim, rows, n, y_stride, y_off, @as(u32, 0) }); } - fn mixer(m: *Model, ly: *const Layer, li: usize, n: u32) !void { + /// The token mixer of every segment's rows: matvecs over all rows, conv and attention per stream. + fn mixer(m: *Model, ly: *const Layer, li: usize, segs: []const Seg, n: u32, hist: bool) !void { const s = &m.s; switch (ly.mixer) { .conv => |c| { try m.proj(c.in, s.xn, s.proj, H, 3 * H, n, 3 * H, 0); - try run(&m.k.conv, .{ H / 64, 1, 1 }, .{ s.proj, 3 * H, m.seq.conv[li], m.seq.conv[li], c.cw, s.att, n, H }); + var r0: usize = 0; + for (segs) |g| { + const k: u32 = @intCast(g.ids.len); + const pr = from(s.proj, r0 * 3 * H * 2); + const y = from(s.att, r0 * H * 2); + if (hist) { + try run(&m.k.conv_hist, .{ H / 64, 1, 1 }, .{ pr, 3 * H, g.q.conv[li], g.q.conv[li], c.cw, y, k, H, g.q.ext[li] }); + } else try run(&m.k.conv, .{ H / 64, 1, 1 }, .{ pr, 3 * H, g.q.conv[li], g.q.conv[li], c.cw, y, k, H }); + r0 += k; + } try m.proj(c.out, s.att, s.d, H, H, n, H, 0); }, .attn => |a| { @@ -372,17 +414,25 @@ pub const Model = struct { try m.proj(a.g, s.xn, s.g, H, NQ, n, NQ, 0); } const inv = if (a.window == 0) m.inv_full else m.inv_slide; - if (m.fuse & 4 != 0) { // q, k and v of every row in one launch - try run(&m.k.rope_store, .{ NQ + NKV + NKV * HD / 64, n, 1 }, .{ s.qkv, QKV, s.q, m.seq.kc[li], m.seq.vc[li], a.qn, a.kn, inv, NQ, NKV, m.pos, a.ring, cfg.eps }); - } else { - try run(&m.k.rope, .{ NQ, n, 1 }, .{ s.qkv, QKV, @as(u32, 0), s.q, NQ * HD, a.qn, inv, m.pos, @as(u32, 0), cfg.eps }); - try run(&m.k.rope, .{ NKV, n, 1 }, .{ s.qkv, QKV, NQ * HD, m.seq.kc[li], NKV * HD, a.kn, inv, m.pos, a.ring, cfg.eps }); - try run(&m.k.vst, .{ NKV * HD / 64, n, 1 }, .{ s.qkv, QKV, (NQ + NKV) * HD, m.seq.vc[li], NKV * HD, m.pos, a.ring }); - } - const span = if (a.window == 0) m.pos + n else @min(m.pos + n, a.window); const scale: f32 = 1.0 / @sqrt(@as(f32, HD)); - try run(&m.k.part, .{ NKV, (span + chunk - 1) / chunk, n }, .{ s.q, m.seq.kc[li], m.seq.vc[li], s.po, s.pm, s.pl, m.pos, chunk, a.ring, a.window, NKV, m.maxc, scale }); - try run(&m.k.merge, .{ NQ, n, 1 }, .{ s.po, s.pm, s.pl, s.g, s.att, m.pos, chunk, a.window, NQ, m.maxc }); + var r0: usize = 0; + for (segs) |g| { + const k: u32 = @intCast(g.ids.len); + const pos = g.q.pos; + const qkv = from(s.qkv, r0 * QKV * 2); + const q = from(s.q, r0 * NQ * HD * 2); + if (m.fuse & 4 != 0) { // q, k and v of every row in one launch + try run(&m.k.rope_store, .{ NQ + NKV + NKV * HD / 64, k, 1 }, .{ qkv, QKV, q, g.q.kc[li], g.q.vc[li], a.qn, a.kn, inv, NQ, NKV, pos, a.ring, cfg.eps }); + } else { + try run(&m.k.rope, .{ NQ, k, 1 }, .{ qkv, QKV, @as(u32, 0), q, NQ * HD, a.qn, inv, pos, @as(u32, 0), cfg.eps }); + try run(&m.k.rope, .{ NKV, k, 1 }, .{ qkv, QKV, NQ * HD, g.q.kc[li], NKV * HD, a.kn, inv, pos, a.ring, cfg.eps }); + try run(&m.k.vst, .{ NKV * HD / 64, k, 1 }, .{ qkv, QKV, (NQ + NKV) * HD, g.q.vc[li], NKV * HD, pos, a.ring }); + } + const span = if (a.window == 0) pos + k else @min(pos + k, a.window); + try run(&m.k.part, .{ NKV, (span + chunk - 1) / chunk, k }, .{ q, g.q.kc[li], g.q.vc[li], s.po, s.pm, s.pl, pos, chunk, a.ring, a.window, NKV, m.maxc, scale }); + try run(&m.k.merge, .{ NQ, k, 1 }, .{ s.po, s.pm, s.pl, from(s.g, r0 * NQ * 2), from(s.att, r0 * NQ * HD * 2), pos, chunk, a.window, NQ, m.maxc }); + r0 += k; + } try m.proj(a.o, s.att, s.d, NQ * HD, H, n, H, 0); }, } @@ -397,17 +447,35 @@ pub const Model = struct { k.f(k.ctx, layer, rows); } - /// Runs the tokens as one window at m.pos (1 .. max_rows rows); `logits`: the last row's logits and argmax too. + /// The tokens as one window of the bound stream (1 .. max_rows rows); `logits`: the last row's logits and argmax. pub fn forward(m: *Model, ids: []const u32, logits: bool, h: ?Hook) !void { - const n: u32 = @intCast(ids.len); - if (n == 0 or n > m.max_rows) return error.Invalid; - if (m.pos + n > m.cap) return error.ContextFull; + try m.forwardSegs(&.{.{ .q = m.seq, .ids = ids }}, if (logits) .last else .none, false, h); + } + + /// One forward over every segment's rows, each at its stream's position; `hist` keeps a window's rollback source. + pub fn forwardSegs(m: *Model, segs: []const Seg, out: Out, hist: bool, h: ?Hook) !void { + var n: u32 = 0; + m.seq.pos = m.pos; + for (segs) |g| { + if (g.ids.len == 0 or (hist and g.ids.len > window_rows)) return error.Invalid; + if (g.q.pos + g.ids.len > m.cap) return error.ContextFull; + n += @intCast(g.ids.len); + } + if (n > m.max_rows or (out == .rows and n > window_rows)) return error.Invalid; const s = &m.s; - try m.r.upload(s.ids, std.mem.sliceAsBytes(ids)); - try run(&m.k.embed, .{ n, 1, 1 }, .{ m.emb.w, m.emb.s, m.emb.b, s.ids, s.x, H }); + if (n == 1 or out == .rows) { // ids as kernel arguments: the queue reads no host memory the caller may reuse + var r: usize = 0; + for (segs) |g| for (g.ids) |id| { + try run(&m.k.embed_tok, .{ 1, 1, 1 }, .{ m.emb.w, m.emb.s, m.emb.b, id, from(s.x, r * H * 2), H }); + r += 1; + }; + } else { // a prompt window: its ids live as long as its stream + try m.r.upload(s.ids, std.mem.sliceAsBytes(segs[0].ids)); + try run(&m.k.embed, .{ n, 1, 1 }, .{ m.emb.w, m.emb.s, m.emb.b, s.ids, s.x, H }); + } try run(&m.k.rms, .{ n, 1, 1 }, .{ s.x, m.layers[0].op_norm, s.xn, H, cfg.eps }); for (m.layers, 0..) |*ly, i| { - try m.mixer(ly, i, n); + try m.mixer(ly, i, segs, n, hist); try run(&m.k.pan, .{ n, 1, 1 }, .{ s.d, s.x, ly.post_op, ly.ffn_norm, s.xn, H, cfg.eps }); if (n == 1 and m.fuse & 1 != 0) { // w1, w3 and their SwiGLU in one launch, each value as the three launches round it try run(&m.k.w13, .{ FF / 4, 1, 1 }, .{ ly.w1.w, ly.w1.s, ly.w1.b, ly.w3.w, ly.w3.s, ly.w3.b, s.xn, s.act, H, FF }); @@ -422,12 +490,53 @@ pub const Model = struct { try m.hook(h, @intCast(i), s.x, n); } try m.hook(h, -1, s.xn, n); - m.pos += n; - if (!logits) return; - try run(&m.k.head, .{ V / 4, 1, 1 }, .{ m.emb.w, m.emb.s, m.emb.b, s.xn, s.logits, H, (n - 1) * H, @as(u32, 0), V }); + for (segs) |g| { + if (hist) { + g.q.win_start = g.q.pos; + g.q.win_rows = @intCast(g.ids.len); + } + g.q.pos += @intCast(g.ids.len); + } + m.pos = m.seq.pos; const per = (V + argmax_parts - 1) / argmax_parts; - try run(&m.k.am_part, .{ argmax_parts, 1, 1 }, .{ s.logits, s.am_scratch, V, per }); - try run(&m.k.am_fin, .{ 1, 1, 1 }, .{ s.am_scratch, s.am_out, argmax_parts }); + switch (out) { + .none => {}, + .last => { + try run(&m.k.head, .{ V / 4, 1, 1 }, .{ m.emb.w, m.emb.s, m.emb.b, s.xn, s.logits, H, (n - 1) * H, @as(u32, 0), V }); + try run(&m.k.am_part, .{ argmax_parts, 1, 1 }, .{ s.logits, s.am_scratch, V, per }); + try run(&m.k.am_fin, .{ 1, 1, 1 }, .{ s.am_scratch, s.am_out, argmax_parts }); + }, + .rows => { + // every row's logits: pdense's fp32 epilogue sums each as qmv4_f32 does (tf-xpu-lfm3_rows-test) + try run(&m.k.pdense, .{ V / 16, (n + tile_rows - 1) / tile_rows, 1 }, .{ m.emb.w, m.emb.s, m.emb.b, s.xn, s.logits, H, V, n, V, @as(u32, 0), @as(u32, 1) }); + try run(&m.k.am_part, .{ argmax_parts, n, 1 }, .{ s.logits, s.am_scratch, V, per }); + try run(&m.k.am_fin, .{ n, 1, 1 }, .{ s.am_scratch, s.am_out, argmax_parts }); + }, + } + } + + /// Keeps the first c rows of the bound stream's last verified window: its position, and its conv rows from ext. + pub fn keepRows(m: *Model, c: u32) !void { + const q = m.seq; + if (q.win_rows == 0) return; + if (c > q.win_rows) return error.Invalid; + if (c < q.win_rows) for (m.layers, 0..) |ly, i| if (ly.mixer == .conv) try m.r.copy(q.conv[i], from(q.ext[i], c * H * 2), conv_bytes); + m.pos = q.win_start + c; + q.win_rows = 0; + } + + /// The greedy tokens of the last forward's rows (Out.rows), in row order. + pub fn rowTokens(m: *Model, out: []i32) !void { + try m.r.sync(); + try m.r.download(std.mem.sliceAsBytes(out), m.s.am_out); + try m.r.sync(); + } + + /// The logits of the last forward's rows (Out.rows), [rows][vocab]. + pub fn rowLogits(m: *Model, out: []f32) !void { + try m.r.sync(); + try m.r.download(std.mem.sliceAsBytes(out), m.s.logits); + try m.r.sync(); } /// Prefill of a prompt in windows of max_rows; the last window leaves its logits and argmax. diff --git a/zig/src/native/xpu_lfm3.zig b/zig/src/native/xpu_lfm3.zig index 9366e4e54e..f0fb4ec536 100644 --- a/zig/src/native/xpu_lfm3.zig +++ b/zig/src/native/xpu_lfm3.zig @@ -10,6 +10,12 @@ pub const formats: []const []const u8 = &.{"mlx-q4g64"}; pub const name = "xpu-lfm3"; pub const Engine = lfm3.engine.Engine; +/// Windows cut at any row give the same bits: states are kept and resumed anywhere, short prompts too. +pub const cache_rules: api.prompt_cache.Rules = .{ .min_prompt = 0, .min_gap = 64 }; + +/// Copies drafted after an 8-token match (the core's 4 lost 6% on a short chat on Iris Xe; XPU_COPY_MATCH overrides). +pub const copy_match: i64 = 8; + /// Rows of one prompt window (pdense above one row; the scratch is sized for it): 128 beat 32 and 64 on Iris Xe. const prompt_rows: u32 = 128; @@ -19,7 +25,7 @@ pub fn promptRows(_: *const Engine) u32 { /// What the startup line adds about the engine's settings. pub fn note(a: Allocator, e: *const Engine) ![]const u8 { - return std.fmt.allocPrint(a, ", prompt windows of {d} rows, no draft head", .{e.window_rows}); + return std.fmt.allocPrint(a, ", prompt windows of {d} rows, copy drafts of the lane core (no head)", .{e.window_rows}); } /// The engine for `o.dir`, or null with `problem` set (the caller releases the device). diff --git a/zig/src/native/xpu_nem.zig b/zig/src/native/xpu_nem.zig index 2b87e1a333..3d5ae12fe8 100644 --- a/zig/src/native/xpu_nem.zig +++ b/zig/src/native/xpu_nem.zig @@ -146,6 +146,7 @@ fn openFamily(comptime S: type, a: Allocator, gpa: Allocator, io: std.Io, o: api h.backend.started = chainStopHandlers; try h.backend.measure(); if (@hasDecl(S.Engine, "copy_only")) { + if (@hasDecl(S, "copy_match")) h.backend.copy_match = S.copy_match; if (getenv("XPU_COPY_MATCH")) |v| h.backend.copy_match = std.fmt.parseInt(i64, v, 10) catch 4; h.backend.streams = @min(@max(1, o.lanes), 8); // one sequence each (the engine holds up to 8) } else if (h.backend.facts().exact_width > 1) h.backend.streams = @max(1, o.lanes); // drafted windows: up to --parallel streams a round diff --git a/zig/tests/xpu/lfm3_rows.zig b/zig/tests/xpu/lfm3_rows.zig new file mode 100644 index 0000000000..36bec81272 --- /dev/null +++ b/zig/tests/xpu/lfm3_rows.zig @@ -0,0 +1,65 @@ +//! LFM3's window matvecs (pdense, bf16 and fp32 out) against the one-row kernels, bit for bit, on random data. + +const std = @import("std"); +const o = @import("lfm3_ops.zig"); +const stop = @import("xpu").stop; +const kernels = @import("xpu").kernels; + +const ROWS: u32 = 1024; + +fn randWords(r: *o.rt.Runtime, n: usize, rnd: std.Random) !o.Buf { + const h = try o.gpa.alloc(u32, n); + for (h) |*v| v.* = rnd.int(u32); + const b = try r.alloc(n * 4); + try r.upload(b, std.mem.sliceAsBytes(h)); + return b; +} + +fn fetch32(r: *o.rt.Runtime, b: o.Buf, n: usize) ![]u32 { + const h = try o.gpa.alloc(u32, n); + try r.sync(); + try r.download(std.mem.sliceAsBytes(h), b); + try r.sync(); + return h; +} + +pub const panic = std.debug.FullPanic(@import("rig.zig").panicDrain); +pub fn main() !void { + stop.install(); + var r = try o.rt.open(); + defer r.deinit(); + var am = try r.module(kernels.attn); + var pm = try r.module(if (kernels.dpas) kernels.nem_prow else kernels.nem_prow_gen12); + const tile_outs: u32 = if (kernels.dpas) 8 else 4; + var k_bf = try am.kernel("qmv4_bf", .{ 64, 1, 1 }); + var k_f32 = try am.kernel("qmv4_f32", .{ 64, 1, 1 }); + var k_pd = try pm.kernel("pdense", .{ 16 * (16 / tile_outs), 1, 1 }); + var t: o.Tally = .{ .name = "lfm3_rows" }; + var prng = std.Random.DefaultPrng.init(21); + const rnd = prng.random(); + std.debug.print("window matvecs (pdense) against the one-row kernels, {d} output rows, bit for bit\n", .{ROWS}); + for ([_]u32{ 2048, 7680 }) |in_dim| { + const sh = try o.gpa.alloc(u16, ROWS * in_dim / 64); + const bh = try o.gpa.alloc(u16, sh.len); + o.randBf(sh, rnd, 0.02); + o.randBf(bh, rnd, 0.1); + const w = try randWords(&r, ROWS * in_dim / 8, rnd); + const s = try o.upload(&r, sh); + const b = try o.upload(&r, bh); + const xh = try o.gpa.alloc(u16, 16 * in_dim); + o.randBf(xh, rnd, 1.0); + const x = try o.upload(&r, xh); + const y1 = try r.alloc(16 * ROWS * 4); + const yn = try r.alloc(16 * ROWS * 4); + for ([_]u32{ 1, 2, 3, 4, 8, 13, 16 }) |n| { + var buf: [96]u8 = undefined; + for (0..n) |z| try o.run(&k_bf, .{ ROWS / 4, 1, 1 }, .{ w, s, b, x, y1, in_dim, @as(u32, @intCast(z)) * in_dim, @as(u32, @intCast(z)) * ROWS, ROWS }); + try o.run(&k_pd, .{ ROWS / 16, (n + 7) / 8, 1 }, .{ w, s, b, x, yn, in_dim, ROWS, n, ROWS, @as(u32, 0), @as(u32, 0) }); + t.same(try std.fmt.bufPrint(&buf, "pdense bf16 in {d} n {d:>2} vs qmv4_bf", .{ in_dim, n }), std.mem.eql(u16, try o.fetch(&r, y1, n * ROWS), try o.fetch(&r, yn, n * ROWS))); + for (0..n) |z| try o.run(&k_f32, .{ ROWS / 4, 1, 1 }, .{ w, s, b, x, y1, in_dim, @as(u32, @intCast(z)) * in_dim, @as(u32, @intCast(z)) * ROWS, ROWS }); + try o.run(&k_pd, .{ ROWS / 16, (n + 7) / 8, 1 }, .{ w, s, b, x, yn, in_dim, ROWS, n, ROWS, @as(u32, 0), @as(u32, 1) }); + t.same(try std.fmt.bufPrint(&buf, "pdense fp32 in {d} n {d:>2} vs qmv4_f32", .{ in_dim, n }), std.mem.eql(u32, try fetch32(&r, y1, n * ROWS), try fetch32(&r, yn, n * ROWS))); + } + } + try t.finish(); +} diff --git a/zig/tests/xpu/lfm3_streams.zig b/zig/tests/xpu/lfm3_streams.zig new file mode 100644 index 0000000000..ee6682d296 --- /dev/null +++ b/zig/tests/xpu/lfm3_streams.zig @@ -0,0 +1,204 @@ +//! LFM3 streams in shared forwards and verified windows with rollback, against each stream alone. + +const std = @import("std"); +const rt = @import("rig.zig"); +const lfm3 = @import("lfm3_xpu"); +const stop = @import("xpu").stop; +const Model = lfm3.model.Model; +const Seq = lfm3.model.Seq; + +const steps = 64; + +fn readIds(gpa: std.mem.Allocator, io: std.Io, ref: []const u8, name: []const u8) ![]u32 { + const path = try std.fmt.allocPrint(gpa, "{s}/{s}/meta.json", .{ ref, name }); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .unlimited); + const meta = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + const vals = meta.value.object.get("ids").?.array.items; + const ids = try gpa.alloc(u32, vals.len); + for (vals, ids) |v, *o| o.* = @intCast(v.integer); + return ids; +} + +/// The stream's prompt on its own sequence (bound), then its first token. +fn prefill(m: *Model, q: *Seq, ids: []const u32) !u32 { + m.bind(q); + try m.reset(); + try m.prefill(ids, null); + return m.argmax(); +} + +/// One row of each stream in one forward (verifyMulti's one-row windows); the rows' greedy tokens in `out`. +fn shared(m: *Model, qs: []const *Seq, toks: []const u32, out: []i32) !void { + var segs: [8]lfm3.model.Seg = undefined; + for (qs, toks, segs[0..qs.len]) |q, *t, *g| g.* = .{ .q = q, .ids = t[0..1] }; + try m.forwardSegs(segs[0..qs.len], .rows, false, null); + try m.rowTokens(out[0..qs.len]); +} + +/// Stream q alone: the prompt, then `steps` greedy tokens. +fn solo(m: *Model, q: *Seq, ids: []const u32, out: []u32) !void { + out[0] = try prefill(m, q, ids); + for (1..out.len) |j| { + try m.forward(out[j - 1 .. j], true, null); + out[j] = try m.argmax(); + } +} + +pub const panic = std.debug.FullPanic(@import("rig.zig").panicDrain); +pub fn main(init: std.process.Init) !void { + stop.install(); + const gpa = init.gpa; + const io = init.io; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 3) { + std.debug.print("skipped: usage tf-xpu-lfm3_streams-test MODEL_DIR REF_DIR\n", .{}); + return; + } + var r = try rt.open(); + defer r.deinit(); + var c = try lfm3.config.Config.read(gpa, io, args[1]); + defer c.deinit(); + var l = try @import("xpu").loader.Loader.init(gpa, &r, args[1]); + var m = try Model.load(gpa, &r, &l, &c, 2048, 128); + const names = [_][]const u8{ "chat", "capitals", "items" }; + var ids: [3][]u32 = undefined; + for (names, &ids) |n, *v| v.* = try readIds(gpa, io, args[2], n); + var want: [3][steps]u32 = undefined; + const q0 = try m.newSeq(); + for (0..3) |k| try solo(&m, q0, ids[k], &want[k]); + // shared: chat and items decode together; items is freed at step 20 and capitals starts on a new sequence + const qa = try m.newSeq(); + const qb = try m.newSeq(); + var got: [3][steps]u32 = undefined; + got[0][0] = try prefill(&m, qa, ids[0]); + got[2][0] = try prefill(&m, qb, ids[2]); + var qs: [2]*Seq = .{ qa, qb }; + var toks: [2]u32 = undefined; + var row: [2]i32 = undefined; + var ja: usize = 1; + var jb: usize = 1; + while (ja < 21) : ({ + ja += 1; + jb += 1; + }) { + toks = .{ got[0][ja - 1], got[2][jb - 1] }; + try shared(&m, &qs, &toks, &row); + got[0][ja] = @intCast(row[0]); + got[2][jb] = @intCast(row[1]); + } + m.freeSeq(qb); + const qc = try m.newSeq(); + got[1][0] = try prefill(&m, qc, ids[1]); + qs = .{ qa, qc }; + var jc: usize = 1; + while (jc < steps - 21) : ({ + ja += 1; + jc += 1; + }) { + toks = .{ got[0][ja - 1], got[1][jc - 1] }; + try shared(&m, &qs, &toks, &row); + got[0][ja] = @intCast(row[0]); + got[1][jc] = @intCast(row[1]); + } + var bad = false; + // four streams in every step: chat, capitals, items, chat again (a duplicate must equal its twin) + { + const order = [_]usize{ 0, 1, 2, 0 }; + var q4: [4]*Seq = undefined; + var g4: [4][steps]u32 = undefined; + for (&q4, order, 0..) |*q, k, z| { + q.* = try m.newSeq(); + g4[z][0] = try prefill(&m, q.*, ids[k]); + } + for (1..steps) |j| { + var t4: [4]u32 = undefined; + for (0..4) |z| t4[z] = g4[z][j - 1]; + var r4: [4]i32 = undefined; + try shared(&m, &q4, &t4, &r4); + for (0..4) |z| g4[z][j] = @intCast(r4[z]); + } + for (order, 0..) |k, z| { + const eq = std.mem.eql(u32, &g4[z], &want[k]); + if (!eq) bad = true; + std.debug.print(" four streams, row {d} ({s}): {s}\n", .{ z, names[k], if (eq) "equal to alone" else "DIFFERS from alone" }); + } + } + // the server's round: haiku (a chat prompt from CHAT_JSON), chat, capitals and items as rows 0..3, 66 steps + if (args.len > 3) { + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, args[3], gpa, .unlimited); + const doc = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + const hv = doc.value.object.get("chat_haiku").?.array.items; + const hids = try gpa.alloc(u32, hv.len); + for (hv, hids) |v, *d| d.* = @intCast(v.integer); + var hwant: [70]u32 = undefined; + try solo(&m, q0, hids, &hwant); + const prompts = [_][]const u32{ hids, ids[0], ids[1], ids[2] }; + var q4: [4]*Seq = undefined; + var g4: [4][70]u32 = undefined; + for (&q4, prompts, 0..) |*q, pr, z| { + q.* = try m.newSeq(); + g4[z][0] = try prefill(&m, q.*, pr); + } + for (1..70) |j| { + var t4: [4]u32 = undefined; + for (0..4) |z| t4[z] = g4[z][j - 1]; + var r4: [4]i32 = undefined; + try shared(&m, &q4, &t4, &r4); + for (0..4) |z| g4[z][j] = @intCast(r4[z]); + } + std.debug.print(" haiku alone, tokens 58..66: {any}\n", .{hwant[58..66]}); + const eq = std.mem.eql(u32, &g4[0], &hwant); + if (!eq) bad = true; + const first = std.mem.indexOfDiff(u32, &g4[0], &hwant); + std.debug.print(" haiku in the server's four-stream round: {s} {any}\n", .{ if (eq) "equal to alone" else "DIFFERS from alone at", first }); + for (1..4) |z| { + const k = z - 1; + const e2 = std.mem.eql(u32, g4[z][0..steps], &want[k]); + if (!e2) bad = true; + std.debug.print(" row {d} ({s}) in that round: {s}\n", .{ z, names[k], if (e2) "equal to alone" else "DIFFERS from alone" }); + } + } + // chat in verified windows (the pending token, up to 4 copied drafts, every third wrong), kept rows only + { + const q = try m.newSeq(); + var got_w: [steps]u32 = undefined; + got_w[0] = try prefill(&m, q, ids[0]); + var j: usize = 1; + var round: usize = 0; + var kept_total: usize = 0; + while (j < steps) : (round += 1) { + const room = @min(5, steps - j + 1); + var win: [5]u32 = undefined; + win[0] = got_w[j - 1]; + for (1..room) |d| win[d] = if ((round + d) % 3 == 0) 7 else want[0][j - 1 + d]; + try m.forwardSegs(&.{.{ .q = q, .ids = win[0..room] }}, .rows, true, null); + var g: [5]i32 = undefined; + try m.rowTokens(g[0..room]); + var keep: usize = 1; + while (keep < room and win[keep] == @as(u32, @intCast(g[keep - 1]))) keep += 1; + m.bind(q); + try m.keepRows(@intCast(keep)); + for (0..keep) |k| if (j + k < steps) { + got_w[j + k] = @intCast(g[k]); + }; + j += keep; + kept_total += keep; + } + const eq = std.mem.eql(u32, &got_w, &want[0]); + if (!eq) bad = true; + std.debug.print(" chat in verified windows with rollback ({d} rounds for {d} tokens): {s} {any}\n", .{ round, kept_total, if (eq) "equal to alone" else "DIFFERS from alone at", std.mem.indexOfDiff(u32, &got_w, &want[0]) }); + } + const checks = [_]struct { k: usize, n: usize, what: []const u8 }{ + .{ .k = 0, .n = ja, .what = "chat, shared all the way (a partner freed and a new one prefilled at step 20)" }, + .{ .k = 2, .n = 21, .what = "items, shared until freed at step 20" }, + .{ .k = 1, .n = jc, .what = "capitals, prefilled mid-way, then shared" }, + }; + for (checks) |ck| { + const eq = std.mem.eql(u32, got[ck.k][0..ck.n], want[ck.k][0..ck.n]); + if (!eq) bad = true; + const first = std.mem.indexOfDiff(u32, got[ck.k][0..ck.n], want[ck.k][0..ck.n]); + std.debug.print(" {s}: {d} tokens {s}{s}\n", .{ ck.what, ck.n, if (eq) "equal to alone" else "DIFFER from alone at ", if (first) |f| try std.fmt.allocPrint(gpa, "{d}", .{f}) else "" }); + } + std.debug.print("{s}\n", .{if (bad) "lfm3_streams: FAILED" else "lfm3_streams: ok"}); + if (bad) return error.TestFailed; +} From 88c079d76775f3f6d54b3789c668a64429e0e2c6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 30 Oct 2026 18:09:46 +0100 Subject: [PATCH 12/14] docs: the Iris Xe LFM3 row after the fused decode and 8 x 4 prompt tiles, with copy drafts, streams and the prompt cache --- docs/recipes/xpu-arc.md | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/docs/recipes/xpu-arc.md b/docs/recipes/xpu-arc.md index 5474455792..bc6be8c855 100644 --- a/docs/recipes/xpu-arc.md +++ b/docs/recipes/xpu-arc.md @@ -54,12 +54,18 @@ LFM3-1.7B (`LiquidAI/LFM3-1.7B-MLX-4bit`, MLX affine 4-bit, group 64) on an Inte DPAS) under WSL2 (Ubuntu 24.04, 8 GB given to WSL), Ubuntu's compute runtime 23.43.27642 and Level Zero loader 1.16, Zig 0.17.0, `zig build xpu-native -Dxpu-device=tgllp` (DPAS off, so the server is built without the Qwen3.8-27B runner): -| Setup | Decode, c=1 | TTFT, 1,024 tokens | TTFT, 8,192 tokens | Memory (`serve --context 16384`) | -| --- | ---: | ---: | ---: | --- | -| Iris Xe, LFM3-1.7B MLX 4-bit | 14.8 to 15.8 tok/s | 13.1 to 14.2 s | 130.8 to 131.3 s | 1.25 GB device allocation, 441 MB host resident | - -Decode: 256 greedy tokens after a 64-token prompt; TTFT: prompts of numbered item sentences, one token generated; three runs each, -through `/v1/completions`. A `/v1/chat/completions` request with the checkpoint's chat template answers (thinking on, `finish_reason` stop). +| Setup | Decode, c=1, `--no-drafts` | Decode, c=1, copy drafts | Decode, c=4, aggregate | TTFT, 1,024 tokens | TTFT, 8,192 tokens | Memory (`serve --context 16384`) | +| --- | ---: | ---: | ---: | ---: | ---: | --- | +| Iris Xe, LFM3-1.7B MLX 4-bit | 15.6 to 16.8 tok/s | 15.9 (chat) to 26.3 (prose) tok/s | 25 to 32 tok/s | 11.7 to 12.2 s | 116.3 to 117.7 s | 1.33 GB device allocation, 448 MB host resident | + +Decode, c=1: two short chats (160 tokens), a 1,001-token prose prompt and a 601-token code prompt (128 tokens), greedy, three runs +each, with copy drafts (the default) and with `--no-drafts`; c=4: four concurrent 128-token completions of distinct 32-token prompts, +six runs a mode, with or without drafts; TTFT: prompts of numbered item sentences, one token generated; three runs each, through +`/v1/completions`. Copy-draft gains depend on how repetitive the text is: prose and code that repeat their own context decode at 26 +and 22 tok/s, the chats, with few copies to draft, at the speed without drafts. A `/v1/chat/completions` request with the +checkpoint's chat template answers (thinking on, `finish_reason` stop). +With `--parallel N` each stream has its own sequence (about 196 MiB of KV at 16,384 tokens); a second chat turn over an +8,037-token document resumes from the prompt cache in 0.80 s instead of 122.6 s, with the same reply. ## Measurements From d9e490cec58388485b7b2599dbbb958bb1f5bb50 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sun, 1 Nov 2026 12:32:56 +0100 Subject: [PATCH 13/14] xpu: LFM3 takes the rope frequencies, ring and window from the family xpu_model.zig drops its own invFreq (theta^e in f32, 1-3 ulp from torch in 16 of the full layers' 32 frequencies and 28 of the sliding layers' 64) for Rotary.inverseFrequencies() (f64 pow, then f32: torch's values), and sizes the sliding rings and windows by cfg.ring and LayerType.window, as the CUDA backend does. --- zig/src/families/lfm3/xpu_model.zig | 18 +++++------------- 1 file changed, 5 insertions(+), 13 deletions(-) diff --git a/zig/src/families/lfm3/xpu_model.zig b/zig/src/families/lfm3/xpu_model.zig index 92099105b7..6d7f75c91c 100644 --- a/zig/src/families/lfm3/xpu_model.zig +++ b/zig/src/families/lfm3/xpu_model.zig @@ -104,12 +104,6 @@ fn run(k: *rt.Kernel, groups: [3]u32, a: anytype) !void { try k.launch(groups); } -/// Inverse RoPE frequencies as upstream builds them (fp32): theta^(-2i/128) for i < angles, then zeros (zero-tail). -pub fn invFreq(theta: f32, partial: f32, out: *[HD / 2]f32) void { - const angles: usize = @intFromFloat(partial * @as(f32, HD) / 2); - for (out, 0..) |*f, i| f.* = if (i < angles) 1.0 / std.math.pow(f32, theta, @as(f32, @floatFromInt(2 * i)) / @as(f32, HD)) else 0; -} - pub const Model = struct { gpa: std.mem.Allocator, r: *rt.Runtime, @@ -224,7 +218,7 @@ pub const Model = struct { .out = try table(l, ns, i, .conv_out), .cw = try vec(l, ns, i, .conv), } } else blk: { - const ring: u32 = if (kind == .sliding_attention) cfg.sliding_window + max_rows else cap; + const ring: u32 = if (kind == .sliding_attention) @intCast(cfg.ring(max_rows)) else cap; break :blk .{ .attn = .{ .q = try table(l, ns, i, .q_proj), .k = try table(l, ns, i, .k_proj), @@ -234,18 +228,16 @@ pub const Model = struct { .qn = try vec(l, ns, i, .q_norm), .kn = try vec(l, ns, i, .k_norm), .ring = ring, - .window = if (kind == .sliding_attention) cfg.sliding_window else 0, + .window = @intCast(kind.window()), } }; }, }; } - var inv: [HD / 2]f32 = undefined; + // the family's RoPE frequencies (torch's float32 values), as every backend's const inv_full = try l.empty(HD / 2 * 4); - invFreq(c.full_rotary.theta, c.full_rotary.partial, &inv); - try r.upload(inv_full, std.mem.sliceAsBytes(&inv)); + try r.upload(inv_full, std.mem.sliceAsBytes(&c.full_rotary.inverseFrequencies())); const inv_slide = try l.empty(HD / 2 * 4); - invFreq(c.sliding_rotary.theta, c.sliding_rotary.partial, &inv); - try r.upload(inv_slide, std.mem.sliceAsBytes(&inv)); + try r.upload(inv_slide, std.mem.sliceAsBytes(&c.sliding_rotary.inverseFrequencies())); try r.sync(); var m: Model = .{ .gpa = gpa, From bb38b814b324ec1b42423be4914eed2156583bef Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Sun, 1 Nov 2026 15:27:58 +0100 Subject: [PATCH 14/14] xpu: LFM3 prompt windows take a tiled GEMM with 4-bit dequant in the loader (2.7x TTFT at 1k), verified windows keep pdense pgemm (lfm3_mv.cl) tiles 64 x 64 in SLM as vit_gemm does and dequantizes q * scale + bias as it loads a weight tile; each output sums an fma chain a 64-wide k block and adds the blocks in order, so a row's sum does not depend on how the prompt is cut. Model.forwardPrompt routes every projection of a prompt window to it (prefill, prefillSpan, prefillWindows); verified windows, shared rounds and decode keep pdense and qmv, so drafts stay bit-equal to one-row steps. lfm3_pgemm checks it against FP64 in groups of 64 and 32 (the other group size as the planted control) and every row alone = in a window; lfm3_chunks checks the model's windows of 64, 128, 256 and a ragged cut bit for bit (rows, logits, the next step), pdense's order as the control. On Iris Xe through the server, TTFT at 1,024 tokens goes from 12.3-12.4 s to 4.6-4.7 s and at 8,192 from 116.3-116.7 s to 55.4-55.6 s. --- zig/build/xpu.zig | 2 +- zig/kernels/xpu/lfm3_mv.cl | 83 +++++++++++++++++++++++++ zig/src/families/lfm3/xpu_engine.zig | 4 +- zig/src/families/lfm3/xpu_model.zig | 24 ++++++- zig/tests/xpu/lfm3_chunks.zig | 93 ++++++++++++++++++++++++++++ zig/tests/xpu/lfm3_fwd.zig | 2 +- zig/tests/xpu/lfm3_pgemm.zig | 78 +++++++++++++++++++++++ 7 files changed, 280 insertions(+), 6 deletions(-) create mode 100644 zig/tests/xpu/lfm3_chunks.zig create mode 100644 zig/tests/xpu/lfm3_pgemm.zig diff --git a/zig/build/xpu.zig b/zig/build/xpu.zig index 4981242a4d..10bf54badc 100644 --- a/zig/build/xpu.zig +++ b/zig/build/xpu.zig @@ -266,7 +266,7 @@ fn lfm3Module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bui /// The standalone op and model tests, each its own program `tf-xpu--test` (step `xpu-tests`). const test_programs = [_][]const u8{ - "panic", "nem_rows", "nem_prefill", "nem_gen", "nem_mtp", "nem_kl", "nem_norm_bench", "nem_mv_bench", "nem_slot", "nem_format", "nem_seam", "nem_attn_dec", "nem_fp64", "memory", "lfm3_conv", "lfm3_rope", "lfm3_attn", "lfm3_norm", "lfm3_fwd", "lfm3_bench", "lfm3_rows", "lfm3_streams", + "panic", "nem_rows", "nem_prefill", "nem_gen", "nem_mtp", "nem_kl", "nem_norm_bench", "nem_mv_bench", "nem_slot", "nem_format", "nem_seam", "nem_attn_dec", "nem_fp64", "memory", "lfm3_conv", "lfm3_rope", "lfm3_attn", "lfm3_norm", "lfm3_fwd", "lfm3_bench", "lfm3_rows", "lfm3_streams", "lfm3_chunks", "lfm3_pgemm", }; /// `-Dxpu` (Linux): Nemotron's SPIR-V images, `tensorfold-xpu` and `tf-xpu-test`; needs ocloc, not nvcc or CUDA. diff --git a/zig/kernels/xpu/lfm3_mv.cl b/zig/kernels/xpu/lfm3_mv.cl index 67b19e4709..05db2de4cf 100644 --- a/zig/kernels/xpu/lfm3_mv.cl +++ b/zig/kernels/xpu/lfm3_mv.cl @@ -69,3 +69,86 @@ __kernel void qmv4_swiglu(__global const uint *w1, __global const ushort *s1, __ const float u = bf(to_bf(qmv_row(w3, s3, b3, x, row, in_dim))); if (get_sub_group_local_id() == 0) act[row] = to_bf(bf(to_bf(a / (1.0f + exp(-a)))) * u); } + +// Prompt windows: bf16(x . (q * scale + bias)), one fma chain a 64-wide k block, blocks in order: any cut, same rows. +__attribute__((reqd_work_group_size(16, 16, 1))) +__kernel void pgemm(__global const uint *w, __global const ushort *scales, __global const ushort *biases, __global const ushort *x, + __global ushort *y, uint k, uint n, uint m, uint ldy, uint y_off, uint qg) { + __local float xs[16][64]; + __local float ws[16][64]; + const uint lx = get_local_id(0), ly = get_local_id(1), t = ly * 16 + lx; + const ulong m0 = (ulong)get_group_id(1) * 64, n0 = (ulong)get_group_id(0) * 64; + const uint lr = t / 4, lk = (t % 4) * 4, words = k / 8, groups = k / qg; + float acc[4][4]; +#pragma unroll + for (uint i = 0; i < 4; i++) +#pragma unroll + for (uint j = 0; j < 4; j++) acc[i][j] = 0.0f; + float part[4][4]; + for (uint k0 = 0; k0 < k; k0 += 16) { + const ulong xr = m0 + lr, wr = n0 + lr; + const uint kk = k0 + lk; + ushort4 xv = (ushort4)(0); + if (xr < m) xv = vload4(0, x + xr * k + kk); + float4 wv = (float4)(0.0f); + if (wr < n) { + const uint p = w[wr * words + kk / 8] >> (4 * (kk % 8)); + const float sc = bf(scales[wr * groups + kk / qg]), bi = bf(biases[wr * groups + kk / qg]); + wv = (float4)((float)(p & 15u) * sc + bi, (float)((p >> 4) & 15u) * sc + bi, (float)((p >> 8) & 15u) * sc + bi, + (float)((p >> 12) & 15u) * sc + bi); + } + xs[lk][lr] = bf(xv.s0); + xs[lk + 1][lr] = bf(xv.s1); + xs[lk + 2][lr] = bf(xv.s2); + xs[lk + 3][lr] = bf(xv.s3); + ws[lk][lr] = wv.s0; + ws[lk + 1][lr] = wv.s1; + ws[lk + 2][lr] = wv.s2; + ws[lk + 3][lr] = wv.s3; + barrier(CLK_LOCAL_MEM_FENCE); + if (k0 % 64 == 0) { +#pragma unroll + for (uint i = 0; i < 4; i++) +#pragma unroll + for (uint j = 0; j < 4; j++) part[i][j] = 0.0f; + } +#pragma unroll + for (uint q = 0; q < 16; q++) { + const float4 a = vload4(ly, &xs[q][0]); + const float4 b = vload4(lx, &ws[q][0]); + part[0][0] = fma(a.s0, b.s0, part[0][0]); + part[0][1] = fma(a.s0, b.s1, part[0][1]); + part[0][2] = fma(a.s0, b.s2, part[0][2]); + part[0][3] = fma(a.s0, b.s3, part[0][3]); + part[1][0] = fma(a.s1, b.s0, part[1][0]); + part[1][1] = fma(a.s1, b.s1, part[1][1]); + part[1][2] = fma(a.s1, b.s2, part[1][2]); + part[1][3] = fma(a.s1, b.s3, part[1][3]); + part[2][0] = fma(a.s2, b.s0, part[2][0]); + part[2][1] = fma(a.s2, b.s1, part[2][1]); + part[2][2] = fma(a.s2, b.s2, part[2][2]); + part[2][3] = fma(a.s2, b.s3, part[2][3]); + part[3][0] = fma(a.s3, b.s0, part[3][0]); + part[3][1] = fma(a.s3, b.s1, part[3][1]); + part[3][2] = fma(a.s3, b.s2, part[3][2]); + part[3][3] = fma(a.s3, b.s3, part[3][3]); + } + if ((k0 + 16) % 64 == 0 || k0 + 16 >= k) { +#pragma unroll + for (uint i = 0; i < 4; i++) +#pragma unroll + for (uint j = 0; j < 4; j++) acc[i][j] += part[i][j]; + } + barrier(CLK_LOCAL_MEM_FENCE); + } +#pragma unroll + for (uint i = 0; i < 4; i++) { + const ulong r = m0 + ly * 4 + i; + if (r >= m) continue; +#pragma unroll + for (uint j = 0; j < 4; j++) { + const ulong o = n0 + lx * 4 + j; + if (o < n) y[r * ldy + y_off + o] = to_bf(acc[i][j]); + } + } +} diff --git a/zig/src/families/lfm3/xpu_engine.zig b/zig/src/families/lfm3/xpu_engine.zig index 6b28bb97bd..f71d117870 100644 --- a/zig/src/families/lfm3/xpu_engine.zig +++ b/zig/src/families/lfm3/xpu_engine.zig @@ -121,7 +121,7 @@ pub const Engine = struct { var t = from; while (t < to) { const n = @min(w, to - t); - try e.m.forward(ids[t .. t + n], t + n == ids.len, null); + try e.m.forwardPrompt(ids[t .. t + n], t + n == ids.len, null); t += n; } } @@ -226,7 +226,7 @@ pub const Engine = struct { const w = @min(rows, e.window_rows); while (t < head.len) { const n = @min(w, head.len - t); - try e.m.forward(head[t .. t + n], false, null); + try e.m.forwardPrompt(head[t .. t + n], false, null); t += n; } try e.r.sync(); diff --git a/zig/src/families/lfm3/xpu_model.zig b/zig/src/families/lfm3/xpu_model.zig index 6d7f75c91c..3fda63bdd4 100644 --- a/zig/src/families/lfm3/xpu_model.zig +++ b/zig/src/families/lfm3/xpu_model.zig @@ -48,6 +48,7 @@ const Kernels = struct { qkvg: rt.Kernel, w13: rt.Kernel, rope_store: rt.Kernel, + pgemm: rt.Kernel, }; /// One stream's state: each conv layer's last six B*x rows, each attention layer's KV slots, and its position. @@ -125,6 +126,10 @@ pub const Model = struct { own: *Seq = undefined, /// Fused launches by bit (1 w1+w3+SwiGLU, 2 q/k/v/gate, 4 rope and KV store); LFM3_FUSE sets them for A/B. fuse: u32 = 7, + /// In forwardPrompt: projections take pgemm, whose rows are the same however the prompt is cut. + prompt: bool = false, + /// The MLX group size of the projections. + qg: u32 = 64, /// The affine tensors of module `m` (in layer `layer`, or a model tensor) by names.zig in namespace `ns`. fn table(l: *ld.Loader, ns: names.Namespace, layer: ?usize, m: names.Module) !Table { @@ -178,6 +183,7 @@ pub const Model = struct { .qkvg = try vm.kernel("qmv4_bf_x4", .{ 64, 1, 1 }), .w13 = try vm.kernel("qmv4_swiglu", .{ 64, 1, 1 }), .rope_store = try fm.kernel("lfm3_rope_store", .{ 64, 1, 1 }), + .pgemm = try vm.kernel("pgemm", .{ 16, 16, 1 }), }; const maxc = (cap + chunk - 1) / chunk; const w: usize = max_rows; @@ -254,6 +260,7 @@ pub const Model = struct { .max_rows = max_rows, .maxc = maxc, .device_bytes = l.total, + .qg = c.quant.global.group, }; if (std.c.getenv("LFM3_FUSE")) |v| m.fuse = std.fmt.parseInt(u32, std.mem.span(v), 10) catch 7; m.own = try m.newSeq(); @@ -372,8 +379,9 @@ pub const Model = struct { m.gpa.destroy(sn); } - /// y[r * y_stride + y_off + o] = x[r] . W[o] for n rows: the decode matvec for one row, pdense for a window. + /// y[r * y_stride + y_off + o] = x[r] . W[o] for n rows: pgemm in a prompt window, else qmv (one row) or pdense. fn proj(m: *Model, t: Table, x: Buf, y: Buf, in_dim: u32, rows: u32, n: u32, y_stride: u32, y_off: u32) !void { + if (m.prompt) return run(&m.k.pgemm, .{ (rows + 63) / 64, (n + 63) / 64, 1 }, .{ t.w, t.s, t.b, x, y, in_dim, rows, n, y_stride, y_off, m.qg }); if (n == 1) return run(&m.k.qmv, .{ (rows + 3) / 4, 1, 1 }, .{ t.w, t.s, t.b, x, y, in_dim, @as(u32, 0), y_off, rows }); try run(&m.k.pdense, .{ (rows + 15) / 16, (n + tile_rows - 1) / tile_rows, 1 }, .{ t.w, t.s, t.b, x, y, in_dim, rows, n, y_stride, y_off, @as(u32, 0) }); } @@ -444,6 +452,18 @@ pub const Model = struct { try m.forwardSegs(&.{.{ .q = m.seq, .ids = ids }}, if (logits) .last else .none, false, h); } + /// A prompt window of the bound stream, every projection in pgemm (verified windows and decode keep theirs). + pub fn forwardPrompt(m: *Model, ids: []const u32, logits: bool, h: ?Hook) !void { + const fuse = m.fuse; + m.prompt = true; + m.fuse = fuse & 4; // a one-row window too: not the decode matvecs + defer { + m.prompt = false; + m.fuse = fuse; + } + try m.forward(ids, logits, h); + } + /// One forward over every segment's rows, each at its stream's position; `hist` keeps a window's rollback source. pub fn forwardSegs(m: *Model, segs: []const Seg, out: Out, hist: bool, h: ?Hook) !void { var n: u32 = 0; @@ -536,7 +556,7 @@ pub const Model = struct { var i: usize = 0; while (i < ids.len) { const n = @min(ids.len - i, m.max_rows); - try m.forward(ids[i .. i + n], i + n == ids.len, h); + try m.forwardPrompt(ids[i .. i + n], i + n == ids.len, h); i += n; } } diff --git a/zig/tests/xpu/lfm3_chunks.zig b/zig/tests/xpu/lfm3_chunks.zig new file mode 100644 index 0000000000..8c27c736c0 --- /dev/null +++ b/zig/tests/xpu/lfm3_chunks.zig @@ -0,0 +1,93 @@ +//! LFM3 prompt windows of 64, 128, 256 (the most) rows and a ragged cut: rows, logits and the next step, bit for bit. + +const std = @import("std"); +const rt = @import("rig.zig"); +const lfm3 = @import("lfm3_xpu"); +const stop = @import("xpu").stop; +const Model = lfm3.model.Model; +const H = lfm3.config.hidden; +const V = lfm3.config.vocab; + +const Rows = struct { + out: []u16, + at: usize = 0, + fn take(ctx: *anyopaque, layer: i32, rows: []const u16) void { + if (layer != -1) return; + const r: *Rows = @ptrCast(@alignCast(ctx)); + @memcpy(r.out[r.at..][0..rows.len], rows); + r.at += rows.len; + } +}; + +/// The prompt in windows of `cut` (cycled): each row's final norm, the last logits, then one decode step's logits. +fn run(m: *Model, ids: []const u32, cut: []const usize, prompt: bool, rows: []u16, last: []f32, step: []f32) !u32 { + var cap: Rows = .{ .out = rows }; + try m.reset(); + var i: usize = 0; + var c: usize = 0; + while (i < ids.len) : (c += 1) { + const n = @min(cut[c % cut.len], ids.len - i); + const hk: lfm3.model.Hook = .{ .ctx = &cap, .f = Rows.take }; + if (prompt) try m.forwardPrompt(ids[i .. i + n], i + n == ids.len, hk) else try m.forward(ids[i .. i + n], i + n == ids.len, hk); + i += n; + } + try m.fetchLogits(last); + const tok = try m.argmax(); + try m.forward(&.{tok}, true, null); + try m.fetchLogits(step); + return tok; +} + +pub const panic = std.debug.FullPanic(@import("rig.zig").panicDrain); +pub fn main(init: std.process.Init) !void { + stop.install(); + const gpa = init.gpa; + const io = init.io; + const args = try init.minimal.args.toSlice(init.arena.allocator()); + if (args.len < 3) { + std.debug.print("skipped: usage tf-xpu-lfm3_chunks-test MODEL_DIR REF_DIR (its items prompt, repeated to 2,100 tokens)\n", .{}); + return; + } + var r = try rt.open(); + defer r.deinit(); + var c = try lfm3.config.Config.read(gpa, io, args[1]); + defer c.deinit(); + var l = try @import("xpu").loader.Loader.init(gpa, &r, args[1]); + var m = try Model.load(gpa, &r, &l, &c, 2304, 256); + const path = try std.fmt.allocPrint(gpa, "{s}/items/meta.json", .{args[2]}); + const meta = try std.json.parseFromSlice(std.json.Value, gpa, try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .unlimited), .{}); + const vals = meta.value.object.get("ids").?.array.items; + const ids = try gpa.alloc(u32, 2100); + for (ids, 0..) |*o, i| o.* = @intCast(vals[if (i < vals.len) i else 1 + (i - 1) % (vals.len - 1)].integer); + const cuts = [_]struct { name: []const u8, cut: []const usize, prompt: bool = true }{ + .{ .name = "windows of 128", .cut = &.{128} }, + .{ .name = "windows of 64", .cut = &.{64} }, + .{ .name = "windows of 256", .cut = &.{256} }, + .{ .name = "ragged (1, 7, 100, 3, 250, 64, 1)", .cut = &.{ 1, 7, 100, 3, 250, 64, 1 } }, + .{ .name = "planted control: windows of 128 in pdense's order (must differ)", .cut = &.{128}, .prompt = false }, + }; + var ref_rows: []u16 = &.{}; + var ref_last: []f32 = &.{}; + var ref_step: []f32 = &.{}; + var ref_tok: u32 = 0; + var bad = false; + for (cuts, 0..) |ct, k| { + const rows = try gpa.alloc(u16, ids.len * H); + const last = try gpa.alloc(f32, V); + const step = try gpa.alloc(f32, V); + const tok = try run(&m, ids, ct.cut, ct.prompt, rows, last, step); + if (k == 0) { + ref_rows = rows; + ref_last = last; + ref_step = step; + ref_tok = tok; + std.debug.print(" {s}: the reference ({d} rows, next token {d})\n", .{ ct.name, ids.len, tok }); + continue; + } + const same = std.mem.eql(u16, rows, ref_rows) and std.mem.eql(f32, last, ref_last) and std.mem.eql(f32, step, ref_step) and tok == ref_tok; + if (same != ct.prompt) bad = true; + std.debug.print(" {s}: every row, the logits and the decode step after {s}\n", .{ ct.name, if (same) "bit for bit equal" else "differ" }); + } + std.debug.print("{s}\n", .{if (bad) "lfm3_chunks: FAILED" else "lfm3_chunks: ok"}); + if (bad) return error.TestFailed; +} diff --git a/zig/tests/xpu/lfm3_fwd.zig b/zig/tests/xpu/lfm3_fwd.zig index 3048349c35..4aa1fdfc11 100644 --- a/zig/tests/xpu/lfm3_fwd.zig +++ b/zig/tests/xpu/lfm3_fwd.zig @@ -102,7 +102,7 @@ pub fn main(init: std.process.Init) !void { while (i < T) { const n = @min(T - i, rows); keep.row0 = i; - try m.forward(ids[i .. i + n], i + n == T, hk); + try m.forwardPrompt(ids[i .. i + n], i + n == T, hk); i += n; } const next = try m.argmax(); diff --git a/zig/tests/xpu/lfm3_pgemm.zig b/zig/tests/xpu/lfm3_pgemm.zig new file mode 100644 index 0000000000..7b289e52b4 --- /dev/null +++ b/zig/tests/xpu/lfm3_pgemm.zig @@ -0,0 +1,78 @@ +//! pgemm, the prompt windows' GEMM: MLX 4-bit in groups of 64 and 32 against FP64; each row alone = in a window. + +const std = @import("std"); +const o = @import("lfm3_ops.zig"); +const stop = @import("xpu").stop; +const kernels = @import("xpu").kernels; + +/// W[out][i] = q * scale + bias of an MLX 4-bit table in groups of `group`, FP64. +fn dequant(w: []const u32, s: []const f64, b: []const f64, outs: usize, in_dim: usize, group: usize) ![]f64 { + const d = try o.gpa.alloc(f64, outs * in_dim); + for (0..outs) |r| for (0..in_dim) |i| { + const q: f64 = @floatFromInt((w[r * (in_dim / 8) + i / 8] >> @intCast(4 * (i % 8))) & 15); + const g = r * (in_dim / group) + i / group; + d[r * in_dim + i] = q * s[g] + b[g]; + }; + return d; +} + +pub const panic = std.debug.FullPanic(@import("rig.zig").panicDrain); +pub fn main() !void { + stop.install(); + var r = try o.rt.open(); + defer r.deinit(); + var vm = try r.module(kernels.lfm3_mv); + var k = try vm.kernel("pgemm", .{ 16, 16, 1 }); + var t: o.Tally = .{ .name = "lfm3_pgemm" }; + var prng = std.Random.DefaultPrng.init(64); + const rnd = prng.random(); + std.debug.print("pgemm: 4-bit tables in groups of 64 and 32, against FP64; rows alone and in a window\n", .{}); + var buf: [96]u8 = undefined; + const outs: u32 = 1000; // not a multiple of the 64-wide tile + const ld: u32 = outs + 24; + const off: u32 = 16; + for ([_]u32{ 2048, 7680 }) |in_dim| for ([_]u32{ 64, 32 }) |group| { + const wh = try o.gpa.alloc(u32, outs * in_dim / 8); + for (wh) |*v| v.* = rnd.int(u32); + const w = try r.alloc(wh.len * 4); + try r.upload(w, std.mem.sliceAsBytes(wh)); + const sh = try o.gpa.alloc(u16, outs * in_dim / group); + const bh = try o.gpa.alloc(u16, sh.len); + o.randBf(sh, rnd, 0.02); + o.randBf(bh, rnd, 0.1); + const s = try o.upload(&r, sh); + const b = try o.upload(&r, bh); + const wd = try dequant(wh, try o.widen(o.gpa, sh), try o.widen(o.gpa, bh), outs, in_dim, group); + const rows: u32 = 130; + const xh = try o.gpa.alloc(u16, rows * in_dim); + o.randBf(xh, rnd, 1.0); + const x = try o.upload(&r, xh); + const xf = try o.widen(o.gpa, xh); + const want = try o.gpa.alloc(f64, rows * outs); + for (0..rows) |i| for (0..outs) |j| { + var acc: f64 = 0; + for (0..in_dim) |c| acc += xf[i * in_dim + c] * wd[j * in_dim + c]; + want[i * outs + j] = acc; + }; + const y = try r.alloc(rows * ld * 2); + const one = try r.alloc(rows * outs * 2); + for (0..rows) |i| try o.run(&k, .{ (outs + 63) / 64, 1, 1 }, .{ w, s, b, sub(x, i * in_dim * 2), sub(one, i * outs * 2), in_dim, outs, @as(u32, 1), outs, @as(u32, 0), group }); + const alone = try o.fetch(&r, one, rows * outs); + t.check(try std.fmt.bufPrint(&buf, "groups of {d}, in {d}, rows one at a time", .{ group, in_dim }), o.relErr(alone, want)); + for ([_]u32{ 5, 64, 130 }) |m| { + try o.run(&k, .{ (outs + 63) / 64, (m + 63) / 64, 1 }, .{ w, s, b, x, y, in_dim, outs, m, ld, off, group }); + const got = try o.fetch(&r, y, m * ld); + var eq = true; + for (0..m) |i| eq = eq and std.mem.eql(u16, got[i * ld + off ..][0..outs], alone[i * outs ..][0..outs]); + t.same(try std.fmt.bufPrint(&buf, "groups of {d}, in {d}, a window of {d:>3} rows (stride, offset) = alone", .{ group, in_dim, m }), eq); + } + // planted control: the table read in the other group size must not match + try o.run(&k, .{ (outs + 63) / 64, 1, 1 }, .{ w, s, b, x, one, in_dim, outs, @as(u32, 1), outs, @as(u32, 0), 96 - group }); + t.control(try std.fmt.bufPrint(&buf, "groups of {d} read as {d}, in {d}", .{ group, 96 - group, in_dim }), o.relErr(try o.fetch(&r, one, outs), want[0..outs])); + }; + try t.finish(); +} + +fn sub(b: o.Buf, off: usize) o.Buf { + return .{ .rt = b.rt, .ptr = @ptrFromInt(@intFromPtr(b.ptr.?) + off), .len = b.len - off }; +}