From f14e1eb15c7a66bdb146c2c6764f7a18c050b038 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 12 Nov 2026 08:16:12 +0100 Subject: [PATCH 01/78] qwen27 metal: drop prompt_glue.metal and prompt_state.metal, whose six kernels nothing launches q27_prompt_product, _silu, _decay, _qk, q27_prompt_conv and q27_prompt_recur were reached only through prompt_gpu.Kernels and prompt_state_gpu.Kernels, which nothing builds with pipelines. The host records in prompt_glue.zig, prompt_gpu.zig, prompt_state.zig and prompt_state_gpu.zig, the build and runtime source lists, the runtime analysis line and the prompt.metal write in tests/qwen27/run.zig go with them. --- zig/build/qwen27.zig | 2 +- zig/kernels/metal/qwen27/prompt_glue.metal | 92 -------------------- zig/kernels/metal/qwen27/prompt_state.metal | 82 ----------------- zig/runtime_sources/qwen27.zig | 2 +- zig/src/families/qwen27/prompt_glue.zig | 50 ----------- zig/src/families/qwen27/prompt_gpu.zig | 71 --------------- zig/src/families/qwen27/prompt_state.zig | 60 ------------- zig/src/families/qwen27/prompt_state_gpu.zig | 38 -------- zig/src/families/qwen27/qwen27.zig | 6 -- zig/tests/qwen27/run.zig | 1 - zig/tests/qwen27/runtime.zig | 1 - 11 files changed, 2 insertions(+), 403 deletions(-) delete mode 100644 zig/kernels/metal/qwen27/prompt_glue.metal delete mode 100644 zig/kernels/metal/qwen27/prompt_state.metal delete mode 100644 zig/src/families/qwen27/prompt_glue.zig delete mode 100644 zig/src/families/qwen27/prompt_gpu.zig delete mode 100644 zig/src/families/qwen27/prompt_state.zig delete mode 100644 zig/src/families/qwen27/prompt_state_gpu.zig diff --git a/zig/build/qwen27.zig b/zig/build/qwen27.zig index 88f939e7984..2faa0416cf6 100644 --- a/zig/build/qwen27.zig +++ b/zig/build/qwen27.zig @@ -16,7 +16,7 @@ pub const Shared = struct { }; /// The 27B's Metal sources, embedded by name (zig/kernels/metal/qwen27). -const parts = [_][]const u8{ "xsum", "norm_xs", "norm_input", "mlp_xs", "mlp_ane", "glue", "copy", "moe", "moe_rows", "attention_mpp", "attention_io", "attention_prompt", "attention_prompt_sg", "prompt_glue", "prompt_state" }; +const parts = [_][]const u8{ "xsum", "norm_xs", "norm_input", "mlp_xs", "mlp_ane", "glue", "copy", "moe", "moe_rows", "attention_mpp", "attention_io", "attention_prompt", "attention_prompt_sg" }; /// The shared modules, imported into `engine` (whose `core` is itself) for the server and the release binary. pub fn engineModules(b: *std.Build, target: std.Build.ResolvedTarget, metal: *std.Build.Module, sources: *std.Build.Module, engine: *std.Build.Module) Shared { diff --git a/zig/kernels/metal/qwen27/prompt_glue.metal b/zig/kernels/metal/qwen27/prompt_glue.metal deleted file mode 100644 index 3ff271a61ab..00000000000 --- a/zig/kernels/metal/qwen27/prompt_glue.metal +++ /dev/null @@ -1,92 +0,0 @@ -// Prompt intermediates round where the declared BF16 graph rounds, rather than at the lane fusion's final store. -#include -using namespace metal; - -struct Q27PromptProduct { uint width, rows, mode, pad; }; -struct Q27PromptQk { uint rows, heads, width, stride, offset; float eps, scale; uint pad; }; -struct Q27PromptDecay { uint rows, heads, weight_flags, pad; }; - -inline float q27_prompt_probability(float value) { - const float magnitude = metal::abs(value); - const float denominator = precise::exp(magnitude) + 1.0f; - const float tail = 1.0f / denominator; - if (value < 0.0f) return tail; - return 1.0f - tail; -} - -inline float q27_prompt_softplus(float value) { - const float small = metal::exp(-metal::abs(value)); - const float shifted = 1.0f + small; - float logarithm = small; - if (shifted != 1.0f) { - const float difference = shifted - 1.0f; - const float ratio = metal::log(shifted) / difference; - logarithm = ratio * small; - } - return metal::max(value, 0.0f) + logarithm; -} - -kernel void q27_prompt_product(device const bfloat* gate [[buffer(0)]], - device const bfloat* factor [[buffer(1)]], - constant Q27PromptProduct& p [[buffer(2)]], - device bfloat* result [[buffer(3)]], - uint at [[thread_position_in_grid]]) { - if (at >= p.rows * p.width) return; - const float x = float(gate[at]); - const float probability = q27_prompt_probability(x); - float activated; - if (p.mode == 0) activated = float(bfloat(x * float(bfloat(probability)))); - else if (p.mode == 1) activated = x * probability; - else activated = x / (1.0f + metal::exp(-x)); - result[at] = bfloat(activated * float(factor[at])); -} - -kernel void q27_prompt_silu(device const bfloat* input [[buffer(0)]], - constant Q27PromptProduct& p [[buffer(1)]], - device bfloat* result [[buffer(2)]], - uint at [[thread_position_in_grid]]) { - if (at >= p.rows * p.width) return; - const float value = float(input[at]); - const bfloat probability = bfloat(q27_prompt_probability(value)); - result[at] = bfloat(value * float(probability)); -} - -kernel void q27_prompt_decay(device const bfloat* a [[buffer(0)]], - device const bfloat* b [[buffer(1)]], - device const void* a_log [[buffer(2)]], - device const void* time_bias [[buffer(3)]], - constant Q27PromptDecay& p [[buffer(4)]], - device float* decay [[buffer(5)]], - device bfloat* beta [[buffer(6)]], - uint at [[thread_position_in_grid]]) { - if (at >= p.rows * p.heads) return; - const uint head = at % p.heads; - const bool wide_bias = (p.weight_flags & 4u) != 0; - const float bias = wide_bias ? ((device const float*)time_bias)[head] : float(((device const bfloat*)time_bias)[head]); - const float logarithm = (p.weight_flags & 2u) != 0 ? ((device const float*)a_log)[head] : float(((device const bfloat*)a_log)[head]); - const float sum = float(a[at]) + bias; - const float shifted = wide_bias ? sum : float(bfloat(sum)); - const float raw = q27_prompt_softplus(shifted); - const float smooth = wide_bias ? raw : float(bfloat(raw)); - const float speed = precise::exp(logarithm); - decay[at] = precise::exp(-speed * smooth); - beta[at] = bfloat(q27_prompt_probability(float(b[at]))); -} - -kernel void q27_prompt_qk(device const bfloat* activation [[buffer(0)]], - constant Q27PromptQk& p [[buffer(1)]], - device bfloat* result [[buffer(2)]], - uint2 group [[threadgroup_position_in_grid]], - uint lane [[thread_index_in_simdgroup]]) { - const uint begin = group.y * p.stride + p.offset + group.x * p.width + lane * 4; - float values[4]; - float square = 0.0f; - for (uint part = 0; part < 4; ++part) { - values[part] = float(activation[begin + part]); - square += values[part] * values[part]; - } - const float inverse = precise::rsqrt(simd_sum(square) / float(p.width) + p.eps); - const float scale = float(bfloat(p.scale)); - const uint output = (group.y * p.heads + group.x) * p.width + lane * 4; - for (uint part = 0; part < 4; ++part) result[output + part] = bfloat(scale * float(bfloat(values[part] * inverse))); -} diff --git a/zig/kernels/metal/qwen27/prompt_state.metal b/zig/kernels/metal/qwen27/prompt_state.metal deleted file mode 100644 index 1f8808136fc..00000000000 --- a/zig/kernels/metal/qwen27/prompt_state.metal +++ /dev/null @@ -1,82 +0,0 @@ -// Prompt state advances through the core tap order and one fixed per-position reduction topology. -#include -using namespace metal; - -struct Q27PromptConv { uint rows, channels, taps, streams; }; -struct Q27PromptRecur { uint rows, nk, nv, dv, streams, pad0, pad1, pad2; }; - -kernel void q27_prompt_conv(device const bfloat* projected [[buffer(0)]], - device const bfloat* history [[buffer(1)]], - device const bfloat* weight [[buffer(2)]], - constant Q27PromptConv& p [[buffer(3)]], - device bfloat* output [[buffer(4)]], - device bfloat* next_history [[buffer(5)]], - uint3 at [[thread_position_in_grid]]) { - if (at.x >= p.channels || at.y >= p.rows || at.z >= p.streams) return; - const uint keep = p.taps - 1; - const uint old_base = at.z * keep * p.channels; - const uint new_base = at.z * p.rows * p.channels; - float sum = 0.0f; - for (uint tap = 0; tap < p.taps; ++tap) { - const uint row = at.y + tap; - const float value = row < keep ? float(history[old_base + row * p.channels + at.x]) : float(projected[new_base + (row - keep) * p.channels + at.x]); - sum += value * float(weight[at.x * p.taps + tap]); - } - output[new_base + at.y * p.channels + at.x] = bfloat(sum); - if (at.y == 0) { - for (uint tail = 0; tail < keep; ++tail) { - const uint row = p.rows + tail; - next_history[old_base + tail * p.channels + at.x] = row < keep ? history[old_base + row * p.channels + at.x] : projected[new_base + (row - keep) * p.channels + at.x]; - } - } -} - -inline float q27_prompt_fold(thread float* partial) { - for (uint stride = 1; stride < 8; stride *= 2) { - for (uint at = 0; at < 8; at += 2 * stride) partial[at] = partial[at] + partial[at + stride]; - } - float total = partial[0]; - for (uint neighbor = 1; neighbor < 4; neighbor *= 2) total = total + simd_shuffle_xor(total, neighbor); - return total; -} - -kernel void q27_prompt_recur(device const bfloat* query [[buffer(0)]], - device const bfloat* key [[buffer(1)]], - device const bfloat* value [[buffer(2)]], - device const float* decay [[buffer(3)]], - device const bfloat* beta [[buffer(4)]], - device const float* committed [[buffer(5)]], - constant Q27PromptRecur& p [[buffer(6)]], - device bfloat* result [[buffer(7)]], - device float* next_state [[buffer(8)]], - uint3 at [[thread_position_in_grid]], - uint lane [[thread_index_in_simdgroup]]) { - const uint stream = at.z / p.nv; - const uint head = at.z % p.nv; - const uint dimension = at.y * 8 + lane / 4; - const uint first = (lane % 4) * 32; - const uint key_head = head / (p.nv / p.nk); - const uint state_base = ((stream * p.nv + head) * p.dv + dimension) * 128 + first; - float state[32]; - for (uint channel = 0; channel < 32; ++channel) state[channel] = committed[state_base + channel]; - for (uint row = 0; row < p.rows; ++row) { - const uint key_base = ((stream * p.rows + row) * p.nk + key_head) * 128 + first; - const uint head_at = (stream * p.rows + row) * p.nv + head; - float partial[8]; - for (uint i = 0; i < 8; ++i) partial[i] = 0.0f; - for (uint channel = 0; channel < 32; ++channel) { - state[channel] = state[channel] * decay[head_at]; - partial[channel / 4] += state[channel] * float(key[key_base + channel]); - } - const float remembered = q27_prompt_fold(partial); - const float correction = (float(value[head_at * p.dv + dimension]) - remembered) * float(beta[head_at]); - for (uint i = 0; i < 8; ++i) partial[i] = 0.0f; - for (uint channel = 0; channel < 32; ++channel) { - state[channel] = state[channel] + float(key[key_base + channel]) * correction; - partial[channel / 4] += state[channel] * float(query[key_base + channel]); - } - const float output = q27_prompt_fold(partial); - if (lane % 4 == 0) result[head_at * p.dv + dimension] = bfloat(output); - } - for (uint channel = 0; channel < 32; ++channel) next_state[state_base + channel] = state[channel]; -} diff --git a/zig/runtime_sources/qwen27.zig b/zig/runtime_sources/qwen27.zig index 83327917957..1cb40bfe3d3 100644 --- a/zig/runtime_sources/qwen27.zig +++ b/zig/runtime_sources/qwen27.zig @@ -3,7 +3,7 @@ const std = @import("std"); const ks = @import("qwen_runtime_sources"); const Set = @import("set.zig").Set; pub fn add(s: *Set) !void { - inline for (.{ "glue", "copy", "moe", "prompt_glue", "prompt_state" }) |name| + inline for (.{ "glue", "copy", "moe" }) |name| try s.raw("qwen27", name, @field(ks, name)); try s.raw("qwen27", "moe-rows", @import("../src/families/qwen27/moe_gpu.zig").rows_source); const quant = @import("../src/families/qwen27/quant_gpu.zig"); diff --git a/zig/src/families/qwen27/prompt_glue.zig b/zig/src/families/qwen27/prompt_glue.zig deleted file mode 100644 index ce681249b01..00000000000 --- a/zig/src/families/qwen27/prompt_glue.zig +++ /dev/null @@ -1,50 +0,0 @@ -//! Prompt operator records keep scaled Q/K epsilon and staged activation modes separate from decode. -const std = @import("std"); -const Config = @import("config.zig").Config; - -pub const Mode = enum(u32) { staged_bf16 = 0, precise_f32 = 1, divided_control = 2 }; -pub const Product = extern struct { width: u32, rows: u32, mode: Mode, pad: u32 = 0 }; -pub const Qk = extern struct { rows: u32, heads: u32, width: u32, stride: u32, offset: u32, eps: f32, scale: f32, pad: u32 = 0 }; -pub const Decay = extern struct { rows: u32, heads: u32, weight_flags: u32, pad: u32 = 0 }; - -pub fn product(width: usize, rows: usize, mode: Mode) !Product { - if (width == 0 or width > 32768 or rows == 0 or rows > 262144 or @as(u64, width) * rows > 67108864) return error.BadPromptProduct; - return .{ .width = @intCast(width), .rows = @intCast(rows), .mode = mode }; -} - -pub fn qk(c: Config, rows: usize, query: bool) !Qk { - if (c.dk != 128 or c.dv != c.dk or c.k_heads == 0 or c.k_heads > 128 or rows == 0 or rows > 2048) return error.UnsupportedPromptQk; - const width: f64 = @floatFromInt(c.dk); - return .{ .rows = @intCast(rows), .heads = @intCast(c.k_heads), .width = @intCast(c.dk), .stride = @intCast(c.gdnQkvDim()), .offset = if (query) 0 else @intCast(c.k_heads * c.dk), .eps = @floatCast(@as(f64, 1e-6) / width), .scale = @floatCast(if (query) 1.0 / width else 1.0 / std.math.sqrt(width)) }; -} - -pub fn check(a: Qk) !void { - if (a.width != 128 or a.rows == 0 or a.rows > 2048 or a.heads == 0 or a.heads > 128 or a.offset > a.stride or a.heads * a.width > a.stride - a.offset or a.pad != 0 or !std.math.isFinite(a.eps) or a.eps <= 0 or !std.math.isFinite(a.scale) or a.scale <= 0) return error.UnsupportedPromptQk; -} - -pub fn checkDecay(a: Decay) !void { - if (a.rows == 0 or a.rows > 2048 or a.heads == 0 or a.heads > 128 or a.weight_flags > 15 or a.pad != 0) return error.BadPromptDecay; -} - -test "prompt QK epsilon and ABI differ explicitly from the decode epsilon" { - try std.testing.expectEqual(@as(usize, 16), @sizeOf(Product)); - try std.testing.expectEqual(@as(usize, 32), @sizeOf(Qk)); - try std.testing.expectEqual(@as(usize, 16), @sizeOf(Decay)); - try std.testing.expectEqual(@as(usize, 20), @offsetOf(Qk, "eps")); - const a = Qk{ .rows = 35, .heads = 16, .width = 128, .stride = 10240, .offset = 0, .eps = 1e-6 / 128.0, .scale = 1.0 / 128.0 }; - try check(a); - try std.testing.expect(a.eps < 1e-6); - var bad = a; - bad.width = 64; - try std.testing.expectError(error.UnsupportedPromptQk, check(bad)); - bad = a; - bad.offset = 10240; - try std.testing.expectError(error.UnsupportedPromptQk, check(bad)); -} - -test "unsupported prompt product extents refuse without changing activation modes" { - try std.testing.expectEqual(Mode.precise_f32, (try product(128, 48, .precise_f32)).mode); - try std.testing.expectEqual(@as(u32, 6144), (try product(128, 128 * 48, .precise_f32)).rows); - try std.testing.expectError(error.BadPromptProduct, product(0, 1, .staged_bf16)); - try std.testing.expectError(error.BadPromptProduct, product(17408, 65536, .staged_bf16)); -} diff --git a/zig/src/families/qwen27/prompt_gpu.zig b/zig/src/families/qwen27/prompt_gpu.zig deleted file mode 100644 index 675a16d5768..00000000000 --- a/zig/src/families/qwen27/prompt_gpu.zig +++ /dev/null @@ -1,71 +0,0 @@ -//! Caller-owned prompt pipelines compose RMS and activation stages without changing decoder pipelines. -const metal = @import("metal"); -const abi = @import("prompt_glue.zig"); -const glue = @import("glue.zig"); -const Ref = @import("projection.zig").Ref; - -pub const Kernels = struct { - head_norm: metal.Pipeline, - product: metal.Pipeline, - silu: metal.Pipeline, - qk: metal.Pipeline, - decay: metal.Pipeline, - - pub fn rms(k: Kernels, e: metal.ComputeEncoder, x: Ref, gain: Ref, output: Ref, a: glue.Head) !void { - if (a.activation != .bf16 or a.gain != .bf16) return error.UnsupportedPromptStorage; - const threads = try glue.headThreads(a); - if (threads > k.head_norm.maxThreads()) return error.UnsupportedPromptGroup; - e.setPipeline(k.head_norm); - e.setBuffer(x.buffer, x.offset, 0); - e.setBuffer(gain.buffer, gain.offset, 1); - e.setValue(a, 2); - e.setBuffer(output.buffer, output.offset, 3); - e.dispatchGroups(.{ .width = a.rows }, .{ .width = threads }); - e.barrier(); - } - - pub fn activation(k: Kernels, e: metal.ComputeEncoder, gate: Ref, factor: Ref, output: Ref, a: abi.Product) !void { - _ = try abi.product(a.width, a.rows, a.mode); - e.setPipeline(k.product); - bind(e, .{ gate, factor }); - e.setValue(a, 2); - e.setBuffer(output.buffer, output.offset, 3); - e.dispatchThreads(.{ .width = @as(usize, a.width) * a.rows }, .{ .width = 256 }); - e.barrier(); - } - - pub fn convolutionActivation(k: Kernels, e: metal.ComputeEncoder, input: Ref, output: Ref, a: abi.Product) !void { - _ = try abi.product(a.width, a.rows, a.mode); - e.setPipeline(k.silu); - e.setBuffer(input.buffer, input.offset, 0); - e.setValue(a, 1); - e.setBuffer(output.buffer, output.offset, 2); - e.dispatchThreads(.{ .width = @as(usize, a.width) * a.rows }, .{ .width = 256 }); - e.barrier(); - } - - pub fn normalizedQk(k: Kernels, e: metal.ComputeEncoder, activation_rows: Ref, output: Ref, a: abi.Qk) !void { - try abi.check(a); - e.setPipeline(k.qk); - e.setBuffer(activation_rows.buffer, activation_rows.offset, 0); - e.setValue(a, 1); - e.setBuffer(output.buffer, output.offset, 2); - e.dispatchGroups(.{ .width = a.heads, .height = a.rows }, .{ .width = 32 }); - e.barrier(); - } - - pub fn gdnDecay(k: Kernels, e: metal.ComputeEncoder, a_rows: Ref, b_rows: Ref, a_log: Ref, dt: Ref, g: Ref, beta: Ref, p: abi.Decay) !void { - try abi.checkDecay(p); - e.setPipeline(k.decay); - bind(e, .{ a_rows, b_rows, a_log, dt }); - e.setValue(p, 4); - e.setBuffer(g.buffer, g.offset, 5); - e.setBuffer(beta.buffer, beta.offset, 6); - e.dispatchThreads(.{ .width = @as(usize, p.rows) * p.heads }, .{ .width = 256 }); - e.barrier(); - } -}; - -fn bind(e: metal.ComputeEncoder, refs: anytype) void { - inline for (refs, 0..) |r, index| e.setBuffer(r.buffer, r.offset, index); -} diff --git a/zig/src/families/qwen27/prompt_state.zig b/zig/src/families/qwen27/prompt_state.zig deleted file mode 100644 index 6de771d694d..00000000000 --- a/zig/src/families/qwen27/prompt_state.zig +++ /dev/null @@ -1,60 +0,0 @@ -//! Prompt convolution and packed-state contracts refuse unsupported shape or storage before dispatch. -const std = @import("std"); - -pub const Conv = extern struct { rows: u32, channels: u32, taps: u32, streams: u32 }; -pub const Recur = extern struct { rows: u32, nk: u32, nv: u32, dv: u32, streams: u32, pad0: u32 = 0, pad1: u32 = 0, pad2: u32 = 0 }; - -pub fn checkConv(a: Conv) !void { - if (a.rows == 0 or a.rows > 2048 or a.channels == 0 or a.channels > 65536 or a.taps < 2 or a.taps > 8 or a.streams == 0 or a.streams > 128) return error.BadPromptConv; - if (@as(u64, a.rows) * a.channels * a.streams > std.math.maxInt(u32)) return error.BadPromptConv; -} - -pub fn checkRecur(a: Recur) !void { - if (a.rows == 0 or a.rows > 2048 or a.nk == 0 or a.nk > 128 or a.nv == 0 or a.nv > 128 or a.nv % a.nk != 0 or a.dv < 8 or a.dv > 256 or a.dv % 8 != 0 or a.streams == 0 or a.streams > 64 or a.pad0 != 0 or a.pad1 != 0 or a.pad2 != 0) return error.BadPromptRecurrence; -} - -pub fn canonical(partials: [32]f32) f32 { - @setFloatMode(.strict); - var values = partials; - var stride: usize = 1; - while (stride < values.len) : (stride *= 2) { - var at: usize = 0; - while (at < values.len) : (at += 2 * stride) values[at] += values[at + stride]; - } - return values[0]; -} - -pub fn packedFold(partials: [32]f32) f32 { - @setFloatMode(.strict); - var values = partials; - var roots: [4]f32 = undefined; - for (0..4) |lane| { - var stride: usize = 1; - while (stride < 8) : (stride *= 2) { - var at: usize = 0; - while (at < 8) : (at += 2 * stride) values[lane * 8 + at] += values[lane * 8 + at + stride]; - } - roots[lane] = values[lane * 8]; - } - return (roots[0] + roots[1]) + (roots[2] + roots[3]); -} - -test "packed local folds preserve the canonical FP32 tree without reassociation" { - for (0..64) |seed| { - var values: [32]f32 = undefined; - for (&values, 0..) |*v, index| { - const bits: u32 = @intCast((seed * 7919 + index * 65537) & 0x7fffff); - v.* = @bitCast((if ((seed + index) % 3 == 0) @as(u32, 1 << 31) else 0) | (120 + @as(u32, @intCast(index % 16))) << 23 | bits); - } - try std.testing.expectEqual(@as(u32, @bitCast(canonical(values))), @as(u32, @bitCast(packedFold(values)))); - } -} - -test "prompt recurrence and convolution ABI and row bounds are explicit" { - try std.testing.expectEqual(@as(usize, 16), @sizeOf(Conv)); - try std.testing.expectEqual(@as(usize, 32), @sizeOf(Recur)); - try checkConv(.{ .rows = 35, .channels = 10240, .taps = 4, .streams = 1 }); - try checkRecur(.{ .rows = 128, .nk = 16, .nv = 48, .dv = 128, .streams = 1 }); - try std.testing.expectError(error.BadPromptRecurrence, checkRecur(.{ .rows = 128, .nk = 16, .nv = 47, .dv = 128, .streams = 1 })); - try std.testing.expectError(error.BadPromptConv, checkConv(.{ .rows = 2049, .channels = 10240, .taps = 4, .streams = 1 })); -} diff --git a/zig/src/families/qwen27/prompt_state_gpu.zig b/zig/src/families/qwen27/prompt_state_gpu.zig deleted file mode 100644 index 55fddbe661d..00000000000 --- a/zig/src/families/qwen27/prompt_state_gpu.zig +++ /dev/null @@ -1,38 +0,0 @@ -//! Prompt chain output and final FP32 state are encoded without touching decoder replay arithmetic. -const metal = @import("metal"); -const abi = @import("prompt_state.zig"); -const Ref = @import("projection.zig").Ref; - -pub const Kernels = struct { - conv: metal.Pipeline, - recur: metal.Pipeline, - - pub fn convolution(k: Kernels, e: metal.ComputeEncoder, projected: Ref, history: Ref, weight: Ref, output: Ref, next_history: Ref, p: abi.Conv) !void { - try abi.checkConv(p); - e.setPipeline(k.conv); - e.setBuffer(projected.buffer, projected.offset, 0); - e.setBuffer(history.buffer, history.offset, 1); - e.setBuffer(weight.buffer, weight.offset, 2); - e.setValue(p, 3); - e.setBuffer(output.buffer, output.offset, 4); - e.setBuffer(next_history.buffer, next_history.offset, 5); - e.dispatchThreads(.{ .width = p.channels, .height = p.rows, .depth = p.streams }, .{ .width = 256 }); - e.barrier(); - } - - pub fn recurrence(k: Kernels, e: metal.ComputeEncoder, q: Ref, key: Ref, value: Ref, decay: Ref, beta: Ref, committed: Ref, output: Ref, next_state: Ref, p: abi.Recur) !void { - try abi.checkRecur(p); - e.setPipeline(k.recur); - e.setBuffer(q.buffer, q.offset, 0); - e.setBuffer(key.buffer, key.offset, 1); - e.setBuffer(value.buffer, value.offset, 2); - e.setBuffer(decay.buffer, decay.offset, 3); - e.setBuffer(beta.buffer, beta.offset, 4); - e.setBuffer(committed.buffer, committed.offset, 5); - e.setValue(p, 6); - e.setBuffer(output.buffer, output.offset, 7); - e.setBuffer(next_state.buffer, next_state.offset, 8); - e.dispatchThreads(.{ .width = 32, .height = p.dv / 8, .depth = p.nv * p.streams }, .{ .width = 32, .height = 2 }); - e.barrier(); - } -}; diff --git a/zig/src/families/qwen27/qwen27.zig b/zig/src/families/qwen27/qwen27.zig index 6b1c366da2f..de17e8b1f23 100644 --- a/zig/src/families/qwen27/qwen27.zig +++ b/zig/src/families/qwen27/qwen27.zig @@ -30,11 +30,7 @@ pub const moe_gpu = @import("moe_gpu.zig"); pub const projection = @import("projection.zig"); pub const prompt_memory = @import("prompt_memory.zig"); pub const prompt_rows = @import("prompt_rows.zig"); -pub const prompt_glue = @import("prompt_glue.zig"); -pub const prompt_gpu = @import("prompt_gpu.zig"); -pub const prompt_state = @import("prompt_state.zig"); pub const prompt_marks = @import("prompt_marks.zig"); -pub const prompt_state_gpu = @import("prompt_state_gpu.zig"); pub const quant_gpu = @import("quant_gpu.zig"); pub const ane_mlp = @import("ane_mlp.zig"); pub const round_plan = @import("round_plan.zig"); @@ -66,8 +62,6 @@ test { _ = @import("prompt_tail.zig"); _ = round_plan; _ = attention; - _ = prompt_glue; - _ = prompt_state; _ = prompt_marks; _ = config; _ = affine; diff --git a/zig/tests/qwen27/run.zig b/zig/tests/qwen27/run.zig index 77be2f0e586..6513b5fb2e7 100644 --- a/zig/tests/qwen27/run.zig +++ b/zig/tests/qwen27/run.zig @@ -34,7 +34,6 @@ fn emit(a: std.mem.Allocator, io: std.Io, dir: []const u8) !void { try write(a, io, dir, "glue.metal", sources.glue ++ sources.copy); try write(a, io, dir, "attention.metal", sources.glue ++ sources.attention_mpp ++ sources.attention_io); try write(a, io, dir, "gdn.metal", gdn_source.text); - try write(a, io, dir, "prompt.metal", sources.prompt_glue ++ sources.prompt_state); for ([_][]const u8{ "norm_input", "norm_xs", "mlp_xs", "xsum" }) |name| { const template = if (std.mem.eql(u8, name, "norm_input")) sources.norm_input else if (std.mem.eql(u8, name, "norm_xs")) sources.norm_xs else if (std.mem.eql(u8, name, "mlp_xs")) sources.mlp_xs else sources.xsum; const mark = std.mem.indexOf(u8, template, "Q27_SPECIALIZE") orelse return error.BadKernelSource; diff --git a/zig/tests/qwen27/runtime.zig b/zig/tests/qwen27/runtime.zig index 4faba9f4812..c561c8ad1f5 100644 --- a/zig/tests/qwen27/runtime.zig +++ b/zig/tests/qwen27/runtime.zig @@ -13,5 +13,4 @@ pub fn main() void { inline for (.{ &q.gdn_pool_gpu.Pool.init, &q.gdn_pool_gpu.Pool.forward, &q.gdn_pool_gpu.Pool.commit, &q.gdn_pool_gpu.Pool.clearSlot, &q.gdn_hook.Hook.forward, &q.gdn_hook.Hook.commit }) |function| std.mem.doNotOptimizeAway(function); inline for (.{ &q.attention.Attention.init, &q.attention.Attention.preprocess, &q.attention.Attention.encode, &q.attention.Attention.keep, &q.attention.Scratch.init, &q.attention.Scratch.upload, &q.attention.Binding.upload, &q.attention.Binding.uploadKeep }) |function| std.mem.doNotOptimizeAway(function); inline for (.{ &q.decode_round.Runner.init, &q.decode_round.Runner.verify, &q.decode_round.Runner.keep, &q.decode_round.Runner.reset }) |function| std.mem.doNotOptimizeAway(function); - inline for (.{ &q.prompt_gpu.Kernels.rms, &q.prompt_gpu.Kernels.activation, &q.prompt_gpu.Kernels.normalizedQk, &q.prompt_gpu.Kernels.gdnDecay, &q.prompt_state_gpu.Kernels.convolution, &q.prompt_state_gpu.Kernels.recurrence }) |function| std.mem.doNotOptimizeAway(function); } From e57053373dfb27afed8111d9b69c1ae89a09763d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Thu, 12 Nov 2026 08:17:05 +0100 Subject: [PATCH 02/78] core metal: drop tf_lane_rope and Attention.rotate, which only lanes_see calls No family calls rotate(); the shared attention's runtime callers (Gemma 4 and Qwen27 DFlash) rope their rows elsewhere. Shape loses rotary_dim and rope, which only the kernel read, and lanes_see feeds raw K and Q straight to append and attention. --- zig/kernels/metal/core/shared_attention.metal | 28 ------ zig/src/core/shared_attention.zig | 43 ++------ zig/src/families/gemma4/dflash.zig | 2 +- .../qwen27/dflash/runtime_backend.zig | 2 +- zig/tests/lanes_see.zig | 98 ++++--------------- 5 files changed, 30 insertions(+), 143 deletions(-) diff --git a/zig/kernels/metal/core/shared_attention.metal b/zig/kernels/metal/core/shared_attention.metal index 126aef9685c..11ae5ce8910 100644 --- a/zig/kernels/metal/core/shared_attention.metal +++ b/zig/kernels/metal/core/shared_attention.metal @@ -149,34 +149,6 @@ inline void tf_context_row(const device bfloat* Q, const device bfloat* K, const tf_context_row(Q, K, V, (const device uchar*)nullptr, O, 0, head, args.keys, 0, as_type(args.scale_bits), lane); } -// Position is supplied per row. Caller-supplied inverse frequencies preserve the model's RoPE convention. -[[kernel]] void tf_lane_rope(const device bfloat* X [[buffer(0)]], const device uint* positions [[buffer(1)]], - const device float* frequencies [[buffer(2)]], constant tf_context_args& args [[buffer(3)]], - device bfloat* Y [[buffer(4)]], uint index [[thread_position_in_grid]]) { - if ((size_t)index >= (size_t)args.rows * args.heads * TF_D) return; - const uint d = index % TF_D; -#if TF_RD > 0 - if (d >= TF_RD) { Y[index] = X[index]; return; } -#if TF_INTERLEAVED - const uint pair = d / 2; - const uint first = pair * 2, second = first + 1; - const bool high = (d & 1) != 0; -#else - const uint pair = d % (TF_RD / 2); - const uint first = pair, second = first + TF_RD / 2; - const bool high = d >= TF_RD / 2; -#endif - const uint row = index / (args.heads * TF_D); - const size_t base = index - d; - const float angle = float(positions[row]) * frequencies[pair]; - const float c = precise::cos(angle), s = precise::sin(angle); - const float x = float(X[base + first]), y = float(X[base + second]); - Y[index] = bfloat(high ? fma(x, s, y*c) : fma(x, c, -y*s)); -#else - Y[index] = X[index]; -#endif -} - // Only newly produced rows are written. Prompt and previous-round cache slots are never copied or rewritten. [[kernel]] void tf_lane_append_kv(const device bfloat* K [[buffer(0)]], const device bfloat* V [[buffer(1)]], constant uint* slots [[buffer(2)]], constant tf_context_args& args [[buffer(3)]], diff --git a/zig/src/core/shared_attention.zig b/zig/src/core/shared_attention.zig index f5d537300e4..f508c61475f 100644 --- a/zig/src/core/shared_attention.zig +++ b/zig/src/core/shared_attention.zig @@ -1,23 +1,20 @@ -//! Shared-context attention, per-row RoPE and cache append; byte masks alone decide what each row sees. +//! Shared-context attention and cache append; byte masks alone decide what each row sees. const std = @import("std"); const mtl = @import("metal"); const sources = @import("kernel_sources"); pub const Output = enum { f32, bf16 }; -pub const Rope = enum { split_half, interleaved }; pub const Shape = struct { query_heads: usize, kv_heads: usize, head_dim: usize, - rotary_dim: usize, - rope: Rope = .split_half, capacity: usize, output: Output = .f32, pub fn validate(s: Shape) !void { if (s.query_heads == 0 or s.kv_heads == 0 or s.query_heads % s.kv_heads != 0 or s.query_heads > 128 or s.head_dim == 0 or s.head_dim > 256 or s.head_dim % 32 != 0 or - s.rotary_dim > s.head_dim or s.rotary_dim % 2 != 0 or s.capacity == 0 or s.capacity > std.math.maxInt(u32)) return error.InvalidAttentionShape; + s.capacity == 0 or s.capacity > std.math.maxInt(u32)) return error.InvalidAttentionShape; _ = try std.math.mul(usize, s.capacity, try std.math.mul(usize, s.kv_heads, s.head_dim)); } }; @@ -30,8 +27,8 @@ const Args = extern struct { rows: u32, keys: u32 = 0, mask_stride: u32 = 0, hea pub fn source(a: std.mem.Allocator, s: Shape) ![]u8 { try s.validate(); - return std.fmt.allocPrint(a, "#define TF_QH {d}\n#define TF_KVH {d}\n#define TF_D {d}\n#define TF_RD {d}\n#define TF_INTERLEAVED {d}\n#define TF_OUT {s}\n{s}", .{ - s.query_heads, s.kv_heads, s.head_dim, s.rotary_dim, @intFromBool(s.rope == .interleaved), if (s.output == .f32) "float" else "bfloat", sources.core_shared_attention, + return std.fmt.allocPrint(a, "#define TF_QH {d}\n#define TF_KVH {d}\n#define TF_D {d}\n#define TF_OUT {s}\n{s}", .{ + s.query_heads, s.kv_heads, s.head_dim, if (s.output == .f32) "float" else "bfloat", sources.core_shared_attention, }); } @@ -39,7 +36,6 @@ pub const Attention = struct { shape: Shape, attend: mtl.Pipeline, decode: mtl.Pipeline, - rope: mtl.Pipeline, append: mtl.Pipeline, split: mtl.Pipeline, merge: mtl.Pipeline, @@ -55,8 +51,6 @@ pub const Attention = struct { errdefer attend.deinit(); const decode = try mtl.Pipeline.init(device, lib, "tf_core_decode_attention", false); errdefer decode.deinit(); - const rope = try mtl.Pipeline.init(device, lib, "tf_lane_rope", false); - errdefer rope.deinit(); const append = try mtl.Pipeline.init(device, lib, "tf_lane_append_kv", false); errdefer append.deinit(); const split = try mtl.Pipeline.init(device, lib, "tf_shared_lane_attention_split", false); @@ -65,12 +59,11 @@ pub const Attention = struct { errdefer merge.deinit(); if (attend.simdWidth() != 32 or attend.maxThreads() < 32 or split.simdWidth() != 32 or merge.simdWidth() != 32) return error.UnsupportedAttentionPipeline; const part = try device.buffer(16 * s.query_heads * chunksOf(s.capacity) * (s.head_dim + 2) * 4, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); - return .{ .shape = s, .attend = attend, .decode = decode, .rope = rope, .append = append, .split = split, .merge = merge, .part = part }; + return .{ .shape = s, .attend = attend, .decode = decode, .append = append, .split = split, .merge = merge, .part = part }; } pub fn deinit(p: Attention) void { p.attend.deinit(); p.decode.deinit(); - p.rope.deinit(); p.append.deinit(); p.split.deinit(); p.merge.deinit(); @@ -145,25 +138,6 @@ pub const Attention = struct { e.dispatchGroups(mtl.Size.of(p.shape.query_heads, 1, 1), mtl.Size.of(32, 1, 1)); } - /// positions u32[rows], frequencies f32[rotary_dim/2], storage bf16[rows,heads,D]; prompt chunks too. - pub fn rotate(p: Attention, e: mtl.ComputeEncoder, rows: usize, heads: usize, x: Ref, positions: Ref, frequencies: Ref, y: Ref) !void { - if (rows == 0 or rows > p.shape.capacity or (heads != p.shape.query_heads and heads != p.shape.kv_heads)) return error.InvalidRopeView; - const elements = try std.math.mul(usize, rows, heads * p.shape.head_dim); - if (elements > std.math.maxInt(u32)) return error.InvalidRopeView; - try check(x, elements * 2, 2); - try check(y, elements * 2, 2); - try check(positions, rows * 4, 4); - try check(frequencies, @max(p.shape.rotary_dim / 2 * 4, 4), 4); - try disjoint(x, y); - try disjoint(y, positions); - try disjoint(y, frequencies); - e.setPipeline(p.rope); - for ([_]Ref{ x, positions, frequencies }, 0..) |r, i| e.setBuffer(r.buf, r.off, i); - e.setValue(Args{ .rows = @intCast(rows), .heads = @intCast(heads) }, 3); - e.setBuffer(y.buf, y.off, 4); - e.dispatchThreads(mtl.Size.of(elements, 1, 1), mtl.Size.of(128, 1, 1)); - } - /// Prepared K and unmodified V go directly to unique physical slots; caller orders subsequent attention. pub fn appendKv(p: Attention, e: mtl.ComputeEncoder, slots: []const u32, k: Ref, v: Ref, cache_k: Ref, cache_v: Ref) !void { if (slots.len == 0 or slots.len > 16) return error.InvalidAppendView; @@ -202,19 +176,16 @@ fn disjoint(a: Ref, b: Ref) !void { if (a.buf.id == b.buf.id) return error.AliasedAttentionBuffer; } -test "attention shapes keep grouped-query and rotary dimensions explicit" { +test "attention shapes keep grouped-query dimensions explicit" { var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); - const good = Shape{ .query_heads = 4, .kv_heads = 2, .head_dim = 64, .rotary_dim = 32, .capacity = 128 }; + const good = Shape{ .query_heads = 4, .kv_heads = 2, .head_dim = 64, .capacity = 128 }; const text = try source(arena.allocator(), good); try std.testing.expect(std.mem.indexOf(u8, text, "#define TF_QH 4\n#define TF_KVH 2\n") != null); var bad = good; bad.query_heads = 3; try std.testing.expectError(error.InvalidAttentionShape, bad.validate()); bad = good; - bad.rotary_dim = 33; - try std.testing.expectError(error.InvalidAttentionShape, bad.validate()); - bad = good; bad.head_dim = 31; try std.testing.expectError(error.InvalidAttentionShape, bad.validate()); } diff --git a/zig/src/families/gemma4/dflash.zig b/zig/src/families/gemma4/dflash.zig index addeedc41fd..7774b78e150 100644 --- a/zig/src/families/gemma4/dflash.zig +++ b/zig/src/families/gemma4/dflash.zig @@ -109,7 +109,7 @@ pub const Drafter = struct { d.ops = try ops.Ops.init(target.device); errdefer d.ops.deinit(); const keys_max = @max(options.capacity + block_rows, c.window - 1 + block_rows); - d.attention = try attention.Attention.init(gpa, target.device, .{ .query_heads = c.heads, .kv_heads = c.kv_heads, .head_dim = c.head_dim, .rotary_dim = c.head_dim, .capacity = keys_max, .output = .bf16 }); + d.attention = try attention.Attention.init(gpa, target.device, .{ .query_heads = c.heads, .kv_heads = c.kv_heads, .head_dim = c.head_dim, .capacity = keys_max, .output = .bf16 }); errdefer d.attention.deinit(); d.picks = try round.Ops.init(target.device); errdefer d.picks.deinit(); diff --git a/zig/src/families/qwen27/dflash/runtime_backend.zig b/zig/src/families/qwen27/dflash/runtime_backend.zig index 51c2a4d3dbc..0fee0ed3906 100644 --- a/zig/src/families/qwen27/dflash/runtime_backend.zig +++ b/zig/src/families/qwen27/dflash/runtime_backend.zig @@ -77,7 +77,7 @@ pub const Backend = struct { try c.check(); const primitives = try core.draft_ops.Ops.init(target.device); errdefer primitives.deinit(); - const attention = try core.shared_attention.Attention.init(a, target.device, .{ .query_heads = c.heads, .kv_heads = c.kv_heads, .head_dim = c.head_dim, .rotary_dim = c.head_dim, .capacity = c.window + 16, .output = .bf16 }); + const attention = try core.shared_attention.Attention.init(a, target.device, .{ .query_heads = c.heads, .kv_heads = c.kv_heads, .head_dim = c.head_dim, .capacity = c.window + 16, .output = .bf16 }); errdefer attention.deinit(); const topk = try core.bf16_topk_gpu.Ops.init(target.device); errdefer topk.deinit(); diff --git a/zig/tests/lanes_see.zig b/zig/tests/lanes_see.zig index 60705e85782..27301d7e572 100644 --- a/zig/tests/lanes_see.zig +++ b/zig/tests/lanes_see.zig @@ -1,4 +1,4 @@ -//! Shared context, data masks and lane-local rotary positions on small synthetic buffers only. +//! Shared context and data masks on small synthetic buffers only. const std = @import("std"); const mtl = @import("metal"); const shared = @import("shared_attention"); @@ -41,30 +41,6 @@ fn guards(b: mtl.Buffer) !void { for (b.contents()[b.length() - pad .. b.length()]) |x| if (x != 0xa5) return error.GuardOverwrite; } -fn checkRope(s: shared.Shape, input: []const u16, output: []const u16, positions: []const u32, freq: []const f32, heads: usize) !void { - for (0..positions.len) |r| for (0..heads) |h| for (0..s.head_dim) |d| { - const at = (r * heads + h) * s.head_dim + d; - if (d >= s.rotary_dim) { - if (input[at] != output[at]) return error.RopeTailChanged; - continue; - } - const pair = if (s.rope == .interleaved) d / 2 else d % (s.rotary_dim / 2); - const first = if (s.rope == .interleaved) pair * 2 else pair; - const second = if (s.rope == .interleaved) first + 1 else first + s.rotary_dim / 2; - const base = at - d; - const x: f64 = fromBf(input[base + first]); - const y: f64 = fromBf(input[base + second]); - const angle: f32 = @as(f32, @floatFromInt(positions[r])) * freq[pair]; - const c = @cos(@as(f64, angle)); - const sn = @sin(@as(f64, angle)); - const high = if (s.rope == .interleaved) d % 2 != 0 else d >= s.rotary_dim / 2; - const want = if (high) x * sn + y * c else x * c - y * sn; - const got: f64 = fromBf(output[at]); - // One final bf16 cast; allow its rounding and the fp32 trig/FMA rounding only. - if (@abs(got - want) > 0.004 * @abs(want) + 0.00001) return error.RopeMismatch; - }; -} - fn cpuAttention(s: shared.Shape, q: []const u16, k: []const u16, v: []const u16, mask: []const u8, keys: usize, row: usize, head: usize, scale: f32, out: []f64) !void { const scores = try a.alloc(f64, keys); defer a.free(scores); @@ -96,12 +72,12 @@ fn outputValue(b: mtl.Buffer, index: usize, output: shared.Output) f32 { return if (output == .f32) values(b, f32, index + 1)[index] else fromBf(values(b, u16, index + 1)[index]); } -const Counts = struct { ordinary: usize = 0, bf16_values: usize = 0, max_error: f64 = 0, max_fraction: f64 = 0, cases: usize = 0, rows: usize = 0, cpu_values: usize = 0, rope_values: usize = 0, prompt_bytes: usize = 0 }; +const Counts = struct { ordinary: usize = 0, bf16_values: usize = 0, max_error: f64 = 0, max_fraction: f64 = 0, cases: usize = 0, rows: usize = 0, cpu_values: usize = 0, prompt_bytes: usize = 0 }; -fn run(device: mtl.Device, queue: mtl.Queue, n: usize, prompt: usize, dims: [4]usize, mode: shared.Rope, output: shared.Output, counts: *Counts) !void { +fn run(device: mtl.Device, queue: mtl.Queue, n: usize, prompt: usize, dims: [3]usize, output: shared.Output, counts: *Counts) !void { const pool = mtl.objc.Pool.push(); defer pool.pop(); - const s = shared.Shape{ .query_heads = dims[0], .kv_heads = dims[1], .head_dim = dims[2], .rotary_dim = dims[3], .rope = mode, .capacity = prompt + n * 4, .output = output }; + const s = shared.Shape{ .query_heads = dims[0], .kv_heads = dims[1], .head_dim = dims[2], .capacity = prompt + n * 4, .output = output }; const p = try shared.Attention.init(a, device, s); defer p.deinit(); var c = try kv.Cache.init(a, .{ .lanes = n, .prompt_capacity = prompt, .round_capacity = 4, .kv_heads = s.kv_heads, .head_dim = s.head_dim }); @@ -117,37 +93,20 @@ fn run(device: mtl.Device, queue: mtl.Queue, n: usize, prompt: usize, dims: [4]u const max_rows = @max(n, prompt); const raw_k = try buffer(device, max_rows * width * 2); defer raw_k.deinit(); - const prepared_k = try buffer(device, max_rows * width * 2); - defer prepared_k.deinit(); const raw_v = try buffer(device, max_rows * width * 2); defer raw_v.deinit(); - const positions = try buffer(device, max_rows * 4); - defer positions.deinit(); - const frequencies = try buffer(device, @max(s.rotary_dim / 2 * 4, 4)); - defer frequencies.deinit(); - const freq = values(frequencies, f32, @max(s.rotary_dim / 2, 1)); - for (freq, 0..) |*v, d| v.* = 1.0 / @as(f32, @floatFromInt(1 + d * 7)); if (prompt > 0) { const raw = values(raw_k, u16, prompt * width); const val = values(raw_v, u16, prompt * width); - const pos = values(positions, u32, prompt); for (raw, 0..) |*v, i| v.* = data(i, 1); for (val, 0..) |*v, i| v.* = data(i, 2); - for (pos, 0..) |*v, i| v.* = @intCast(i); - const cb = queue.commandBuffer(); - const e = cb.compute(.serial); - try p.rotate(e, prompt, s.kv_heads, ref(raw_k, 0), ref(positions, 0), ref(frequencies, 0), ref(prepared_k, 0)); - e.end(); - try finish(cb); - try checkRope(s, raw, values(prepared_k, u16, prompt * width), pos, freq, s.kv_heads); - counts.rope_values += prompt * width; - @memcpy(values(ck, u16, prompt * width), values(prepared_k, u16, prompt * width)); + @memcpy(values(ck, u16, prompt * width), raw); @memcpy(values(cv, u16, prompt * width), val); const ids = try a.alloc(u32, prompt); defer a.free(ids); for (ids, 0..) |*v, i| v.* = @intCast(10 + i); - try c.appendPrompt(ids, values(prepared_k, u16, prompt * width), val); + try c.appendPrompt(ids, raw, val); } const prompt_k = try a.dupe(u16, values(ck, u16, prompt * width)); defer a.free(prompt_k); @@ -160,11 +119,9 @@ fn run(device: mtl.Device, queue: mtl.Queue, n: usize, prompt: usize, dims: [4]u var slots: [16]u32 = undefined; var owners: [16]usize = undefined; var used: usize = 0; - const pos = values(positions, u32, n); for (0..n) |lane| if (active & (@as(u16, 1) << @intCast(lane)) != 0) { owners[used] = lane; slots[used] = @intCast(try c.slot(ticket, lane)); - pos[used] = ticket.positions[lane]; for (0..width) |d| { values(raw_k, u16, n * width)[used * width + d] = data(d, 3 + round * 17 + lane); values(raw_v, u16, n * width)[used * width + d] = data(d, 5 + round * 19 + lane); @@ -173,26 +130,20 @@ fn run(device: mtl.Device, queue: mtl.Queue, n: usize, prompt: usize, dims: [4]u }; const cb = queue.commandBuffer(); const e = cb.compute(.serial); - try p.rotate(e, used, s.kv_heads, ref(raw_k, 0), ref(positions, 0), ref(frequencies, 0), ref(prepared_k, 0)); - e.barrier(); - try p.appendKv(e, slots[0..used], ref(prepared_k, 0), ref(raw_v, 0), ref(ck, 0), ref(cv, 0)); + try p.appendKv(e, slots[0..used], ref(raw_k, 0), ref(raw_v, 0), ref(ck, 0), ref(cv, 0)); e.end(); try finish(cb); - try checkRope(s, values(raw_k, u16, used * width), values(prepared_k, u16, used * width), pos[0..used], freq, s.kv_heads); - counts.rope_values += used * width; - for (0..used) |i| try c.put(ticket, owners[i], @intCast(100 + round * 16 + owners[i]), values(prepared_k, u16, used * width)[i * width ..][0..width], values(raw_v, u16, used * width)[i * width ..][0..width]); + for (0..used) |i| try c.put(ticket, owners[i], @intCast(100 + round * 16 + owners[i]), values(raw_k, u16, used * width)[i * width ..][0..width], values(raw_v, u16, used * width)[i * width ..][0..width]); try std.testing.expectEqualSlices(u16, prompt_k, values(ck, u16, prompt * width)); try std.testing.expectEqualSlices(u16, prompt_v, values(cv, u16, prompt * width)); for (0..used) |i| { const off = @as(usize, slots[i]) * width; - try std.testing.expectEqualSlices(u16, values(prepared_k, u16, used * width)[i * width ..][0..width], values(ck, u16, s.capacity * width)[off..][0..width]); + try std.testing.expectEqualSlices(u16, values(raw_k, u16, used * width)[i * width ..][0..width], values(ck, u16, s.capacity * width)[off..][0..width]); try std.testing.expectEqualSlices(u16, values(raw_v, u16, used * width)[i * width ..][0..width], values(cv, u16, s.capacity * width)[off..][0..width]); } if (round < 2) try c.commit(ticket) else last = ticket; } const view = try c.current(last); - const raw_q = try buffer(device, n * qwidth * 2); - defer raw_q.deinit(); const q = try buffer(device, n * qwidth * 2); defer q.deinit(); const masks = try buffer(device, n * view.keys); @@ -202,10 +153,7 @@ fn run(device: mtl.Device, queue: mtl.Queue, n: usize, prompt: usize, dims: [4]u defer out.deinit(); const serial = try buffer(device, n * qwidth * elem); defer serial.deinit(); - const raw = values(raw_q, u16, n * qwidth); - for (raw, 0..) |*v, i| v.* = data(i, 73); - const qp = values(positions, u32, n); - @memcpy(qp, last.positions[0..n]); + for (values(q, u16, n * qwidth), 0..) |*v, i| v.* = data(i, 73); const mask = values(masks, u8, n * view.keys); try c.visibility(view, n, mask); if (n > 1) { @@ -219,15 +167,11 @@ fn run(device: mtl.Device, queue: mtl.Queue, n: usize, prompt: usize, dims: [4]u const scale: f32 = 1.0 / @sqrt(@as(f32, @floatFromInt(s.head_dim))); const cb = queue.commandBuffer(); const e = cb.compute(.serial); - try p.rotate(e, n, s.query_heads, ref(raw_q, 0), ref(positions, 0), ref(frequencies, 0), ref(q, 0)); - e.barrier(); try p.encode(e, n, view.keys, view.keys, scale, ref(q, 0), ref(ck, 0), ref(cv, 0), ref(masks, 0), ref(out, 0)); for (0..n) |row| try p.encode(e, 1, view.keys, view.keys, scale, ref(q, row * qwidth * 2), ref(ck, 0), ref(cv, 0), ref(masks, row * view.keys), ref(serial, row * qwidth * elem)); if (n == 1) try p.ordinary(e, view.keys, scale, ref(q, 0), ref(ck, 0), ref(cv, 0), ref(serial, 0)); e.end(); try finish(cb); - try checkRope(s, raw, values(q, u16, n * qwidth), qp, freq, s.query_heads); - counts.rope_values += n * qwidth; try std.testing.expectEqualSlices(u8, out.contents()[pad..][0 .. n * qwidth * elem], serial.contents()[pad..][0 .. n * qwidth * elem]); const expected = try a.alloc(f64, s.head_dim); defer a.free(expected); @@ -259,12 +203,12 @@ fn run(device: mtl.Device, queue: mtl.Queue, n: usize, prompt: usize, dims: [4]u try std.testing.expectEqual(fromBf(values(cv, u16, s.capacity * width)[((last.base + n - 1) * s.kv_heads + kh) * s.head_dim + d]), outputValue(out, head * s.head_dim + d, output)); try std.testing.expectEqual(@as(f32, 0), outputValue(out, ((n - 1) * s.query_heads + head) * s.head_dim + d, output)); }; - for ([_]mtl.Buffer{ ck, cv, raw_k, prepared_k, raw_v, positions, frequencies, raw_q, q, masks, out, serial }) |b| try guards(b); + for ([_]mtl.Buffer{ ck, cv, raw_k, raw_v, q, masks, out, serial }) |b| try guards(b); if (n == 1) counts.ordinary += 1; counts.prompt_bytes += prompt * width * 4; counts.rows += n; counts.cases += 1; - std.debug.print("PASS N{d} P{d} QH{d}/KVH{d} D{d} RoPE{d} {s} {s}: data masks, peer-current visibility, serial-row equality, prompt unchanged\n", .{ n, prompt, s.query_heads, s.kv_heads, s.head_dim, s.rotary_dim, @tagName(mode), @tagName(output) }); + std.debug.print("PASS N{d} P{d} QH{d}/KVH{d} D{d} {s}: data masks, peer-current visibility, serial-row equality, prompt unchanged\n", .{ n, prompt, s.query_heads, s.kv_heads, s.head_dim, @tagName(output) }); } /// Compare split attention to serial attention and fp64 at draft shapes and trained-block masks. @@ -272,7 +216,7 @@ fn splitCase(device: mtl.Device, queue: mtl.Queue, keys: usize, counts: *Counts) const pool = mtl.objc.Pool.push(); defer pool.pop(); const rows = 16; - const s = shared.Shape{ .query_heads = 32, .kv_heads = 8, .head_dim = 128, .rotary_dim = 128, .capacity = 2064, .output = .f32 }; + const s = shared.Shape{ .query_heads = 32, .kv_heads = 8, .head_dim = 128, .capacity = 2064, .output = .f32 }; const p = try shared.Attention.init(a, device, s); defer p.deinit(); const qw = s.query_heads * s.head_dim; @@ -339,14 +283,14 @@ pub fn main(_: std.process.Init) !void { const queue = try device.queue(); defer queue.deinit(); var counts: Counts = .{}; - for ([_][4]usize{ .{ 2, 1, 32, 0 }, .{ 4, 2, 64, 32 }, .{ 8, 2, 128, 128 }, .{ 2, 2, 256, 64 } }) |dims| for ([_]shared.Rope{ .split_half, .interleaved }) |mode| for (1..17) |n| { - try run(device, queue, n, if (n % 3 == 0) 0 else if (n % 3 == 1) 5 else 33, dims, mode, .f32, &counts); - }; - for ([_][4]usize{ .{ 2, 1, 32, 0 }, .{ 4, 2, 64, 32 }, .{ 8, 2, 128, 128 }, .{ 2, 2, 256, 64 } }) |dims| for ([_]shared.Rope{ .split_half, .interleaved }) |mode| { - for ([_]usize{ 0, 1, 33 }) |prompt| try run(device, queue, 1, prompt, dims, mode, .f32, &counts); - try run(device, queue, 1, 5, dims, mode, .bf16, &counts); - try run(device, queue, 8, 5, dims, mode, .bf16, &counts); + for ([_][3]usize{ .{ 2, 1, 32 }, .{ 4, 2, 64 }, .{ 8, 2, 128 }, .{ 2, 2, 256 } }) |dims| for (1..17) |n| { + try run(device, queue, n, if (n % 3 == 0) 0 else if (n % 3 == 1) 5 else 33, dims, .f32, &counts); }; + for ([_][3]usize{ .{ 2, 1, 32 }, .{ 4, 2, 64 }, .{ 8, 2, 128 }, .{ 2, 2, 256 } }) |dims| { + for ([_]usize{ 0, 1, 33 }) |prompt| try run(device, queue, 1, prompt, dims, .f32, &counts); + try run(device, queue, 1, 5, dims, .bf16, &counts); + try run(device, queue, 8, 5, dims, .bf16, &counts); + } for ([_]usize{ 300, 1100, 2063 }) |keys| try splitCase(device, queue, keys, &counts); - std.debug.print("PASS {d} cases, {d} ordinary N=1 comparisons, {d} batch/serial rows, {d} fp32 CPU attention values, {d} bf16 values, {d} RoPE values, {d} prompt bytes preserved, max fp32 error {d:.9}, max bound fraction {d:.6}\n", .{ counts.cases, counts.ordinary, counts.rows, counts.cpu_values, counts.bf16_values, counts.rope_values, counts.prompt_bytes, counts.max_error, counts.max_fraction }); + std.debug.print("PASS {d} cases, {d} ordinary N=1 comparisons, {d} batch/serial rows, {d} fp32 CPU attention values, {d} bf16 values, {d} prompt bytes preserved, max fp32 error {d:.9}, max bound fraction {d:.6}\n", .{ counts.cases, counts.ordinary, counts.rows, counts.cpu_values, counts.bf16_values, counts.prompt_bytes, counts.max_error, counts.max_fraction }); } From 77be1352e161119792f294a75cbdbc3d5d1a48bf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 13 Nov 2026 00:13:06 +0100 Subject: [PATCH 03/78] runtime sources: drop the Rope loop and rotary fields core.zig still passed to shared_attention; changelog line --- CHANGELOG.md | 1 + zig/runtime_sources/core.zig | 6 +++--- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 716bd8c6a93..fbe77abbd46 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -82,6 +82,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - The prompt cache can keep a bounded ladder of intermediate checkpoints under a long prompt's history mark, so a later, compacted prompt resumes inside its retained prefix instead of the shared cut. On one GB10 with `nvidia/Qwen3.8-Flash-Next-NVFP4`, a 104K-token request that kept its first 100K resumed 65,536 tokens instead of 2, cutting TTFT from 50.7 s to 19.9 s with the same tokens; a cold 165K prompt paid 0.9 s more, and the 8 GiB cache held two more states. A follow-up after that turn resumed 65,536 tokens too, because the compacted turn evicted the larger states. `TF_PROMPT_INTERMEDIATE=1` enables it for Flash Next on CUDA; it stays off by default (#1909). Its checkpoints go after a conversation's other redundant copies but before any conversation's newest state, and take only the room left after the endpoint, the rewind state and other requests' promised keeps. Thanks to @javierbits. +- Metal: kernels nothing launches are removed: Qwen3.8-27B's `prompt_glue.metal` and `prompt_state.metal` (six `q27_prompt_*` kernels) and core's `tf_lane_rope` with `Attention.rotate`; saved Qwen3.8-27B prompt states on Metal are recomputed once after upgrading, since their identity covers the kernel sources (#3065, part of #2251). Thanks to @jschmied. - Tests: main's nightly is green again. The Flash Next learn, prompt and compare tools build again on Linux (#2585 had put one file in two modules), and the dense serving and serving-identity pins take the changes #2579, #2585, #2285, #2205 and #2789 made, each named beside its pin. Serving is unchanged (#3002). - HIP: Qwen3.8-Flash-Next's MTP head takes a prompt chunk on the prompt tiles, not a decode window at a time, so cold prompts finish about 5% sooner at 8K and 32K on a Radeon 8060S (gfx1151). Replies are unchanged; the head's caches and drafts are not bit-identical to before (#1878). Thanks to @gilby. diff --git a/zig/runtime_sources/core.zig b/zig/runtime_sources/core.zig index dbc6cb491fd..6472609c1bc 100644 --- a/zig/runtime_sources/core.zig +++ b/zig/runtime_sources/core.zig @@ -46,10 +46,10 @@ fn projections(s: *Set) !void { try s.fragments("core-lane", "prepared", try lane.source(s.a, l)); }; const attention = @import("../src/core/shared_attention.zig"); - for (std.enums.values(attention.Rope)) |rope| for (std.enums.values(attention.Output)) |out| { - const text = try attention.source(s.a, .{ .query_heads = 8, .kv_heads = 2, .head_dim = 128, .rotary_dim = 64, .capacity = 4096, .rope = rope, .output = out }); + for (std.enums.values(attention.Output)) |out| { + const text = try attention.source(s.a, .{ .query_heads = 8, .kv_heads = 2, .head_dim = 128, .capacity = 4096, .output = out }); try s.add("core", "shared-attention", text, .{ .language = mtl.CompileOptions.mlx().language, .math_functions = .precise }); - }; + } const route = @import("../src/core/moe_route.zig"); for ([_]route.Shape{ .{ .hidden = 4096, .experts = 288, .topk = 8 }, .{ .hidden = 4096, .experts = 128, .topk = 8, .out_bf16 = true }, .{ .hidden = 6144, .experts = 128, .topk = 4 } }) |shape| try s.raw("core", "route", try route.source(s.a, shape)); From 55dcb9800d1d29af4ce43bd34c21f0bbda86756e Mon Sep 17 00:00:00 2001 From: Ash <71890694+ashhart@users.noreply.github.com> Date: Thu, 12 Nov 2026 09:05:30 +0000 Subject: [PATCH 04/78] deepseek_v41: tf-dsv41-assets builds ReleaseSafe whatever -Doptimize says dsv41_cuda.module pinned the family, kernel and runtime modules to Debug, and dsv41-assets built tf-dsv41-assets on it, so the AOT planner always ran unoptimized. Defilan measured one aot-needs enumeration at --context 524288 on a GB10: about 45 minutes as built, under 7 with the family modules at ReleaseSafe, planning the same 138 variants (#2270). module and runtime keep their Debug signatures for every other caller; moduleAt and runtimeAt take the mode, and dsv41-assets builds the tool on moduleAt(.ReleaseSafe), so its runtime and kernel modules go to ReleaseSafe too. ReleaseSafe keeps the safety checks, and no code under the tool reads the optimize mode, so the plan is unchanged. test-dsv41-assets-host stays on the Debug module. Refs #2270. Co-authored-by: Defilan <7034124+Defilan@users.noreply.github.com> --- zig/build/dsv41_assets.zig | 5 +++-- zig/build/dsv41_cuda.zig | 17 ++++++++++++----- 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/zig/build/dsv41_assets.zig b/zig/build/dsv41_assets.zig index 06c612faa61..aa950c2ffa5 100644 --- a/zig/build/dsv41_assets.zig +++ b/zig/build/dsv41_assets.zig @@ -2,7 +2,8 @@ const std = @import("std"); pub fn steps(b: *std.Build) void { const step = b.step("dsv41-assets", "Install CPU-only DeepSeek token-map and AOT-plan generators"); - const asset_module = @import("dsv41_cuda.zig").module(b, b.graph.host, "asset_main.zig", false); + // ReleaseSafe whatever -Doptimize says: a 512K aot-needs pass took ~45 min in Debug on a GB10, under 7 here (#2270) + const asset_module = @import("dsv41_cuda.zig").moduleAt(b, b.graph.host, .ReleaseSafe, "asset_main.zig", false); const assets = b.addExecutable(.{ .name = "tf-dsv41-assets", .root_module = asset_module }); step.dependOn(&b.addInstallArtifact(assets, .{}).step); const token_map = b.addExecutable(.{ .name = "tf-dsv41-token-map", .root_module = b.createModule(.{ @@ -12,6 +13,6 @@ pub fn steps(b: *std.Build) void { .link_libc = true, }) }); step.dependOn(&b.addInstallArtifact(token_map, .{}).step); - const tests = b.addRunArtifact(b.addTest(.{ .root_module = asset_module })); + const tests = b.addRunArtifact(b.addTest(.{ .root_module = @import("dsv41_cuda.zig").module(b, b.graph.host, "asset_main.zig", false) })); b.step("test-dsv41-assets-host", "DeepSeek asset commands refuse GPU operations and malformed arguments").dependOn(&tests.step); } diff --git a/zig/build/dsv41_cuda.zig b/zig/build/dsv41_cuda.zig index 6b64f7d380e..b3e19f0cc64 100644 --- a/zig/build/dsv41_cuda.zig +++ b/zig/build/dsv41_cuda.zig @@ -2,11 +2,14 @@ const std = @import("std"); const shared = @import("dsv41_shared.zig"); const units = @import("dsv41_units.zig"); pub fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, images: bool) *std.Build.Module { + return runtimeAt(b, target, .Debug, images); +} +pub fn runtimeAt(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, images: bool) *std.Build.Module { const options = @import("cuda_runtime.zig").kernelOptions(b, images, null); - const cuda = b.createModule(.{ .root_source_file = b.path("zig/src/cuda/root.zig"), .target = target, .optimize = .Debug, .link_libc = true }); + const cuda = b.createModule(.{ .root_source_file = b.path("zig/src/cuda/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); cuda.addImport("gpu_memory", @import("gpu_memory.zig").module(b, target)); cuda.addOptions("kernel_options", options); - cuda.addImport("stagger", b.createModule(.{ .root_source_file = b.path("zig/src/core/stagger.zig"), .target = target, .optimize = .Debug })); + cuda.addImport("stagger", b.createModule(.{ .root_source_file = b.path("zig/src/core/stagger.zig"), .target = target, .optimize = optimize })); if (images) { const empty = b.addWriteFiles().add("compile-only.fatbin", "SYNTHETIC-COMPILE-ONLY"); const all = @import("cuda.zig").kernels; @@ -15,16 +18,20 @@ pub fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, images: bool) *s return cuda; } pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, file: []const u8, images: bool) *std.Build.Module { - const cuda = runtime(b, target, images); + return moduleAt(b, target, .Debug, file, images); +} +/// `module` with the runtime, kernel and family modules at `optimize`. +pub fn moduleAt(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, file: []const u8, images: bool) *std.Build.Module { + const cuda = runtimeAt(b, target, optimize, images); const opts = b.addOptions(); opts.addOption(bool, "with_kernels", images); opts.addOption(u32, "family_scope", 41); - const dk = b.createModule(.{ .root_source_file = b.path("zig/src/families/deepseek_v41/kernels.zig"), .target = target, .optimize = .Debug, .link_libc = true }); + const dk = b.createModule(.{ .root_source_file = b.path("zig/src/families/deepseek_v41/kernels.zig"), .target = target, .optimize = optimize, .link_libc = true }); dk.addImport("cuda", cuda); dk.addOptions("dsv41_kernel_options", opts); const empty = b.addWriteFiles().add("compile-only-dsv41.fatbin", "SYNTHETIC-COMPILE-ONLY"); for (units.units ++ units.optional) |k| dk.addAnonymousImport(b.fmt("dsv41_fatbin_{s}", .{k.name}), .{ .root_source_file = empty }); - return assemble(b, target, .Debug, file, cuda, dk); + return assemble(b, target, optimize, file, cuda, dk); } /// The family's root `file` over the runtime `cuda` and kernel module `dk`, every family module at `optimize`. pub fn assemble(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, file: []const u8, cuda: *std.Build.Module, dk: *std.Build.Module) *std.Build.Module { From b23d43318827d1a08e208ae3c259bd24323b6eac Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 12:01:02 +0000 Subject: [PATCH 05/78] changelog: #3007 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index fbe77abbd46..c3007ea6f1f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- DeepSeek-V4.1 on CUDA: `zig build dsv41-assets` now builds `tf-dsv41-assets` at ReleaseSafe whatever `-Doptimize` says, so `aot-needs` plans the same Triton variants about 15 times faster (713 s against 45 s at a 65,536-token context on a Mac; a GB10 run at 524,288 tokens went from about 45 minutes to under 7) (#3007, found on #2270). Thanks to @Defilan. + - Docs: the CUDA memory plan says what each setting decides and which file reads it, including that `--context` is capped at the model's own window and `--context 0` takes it (#3028). Thanks to @Bizuayeu. - Qwen3.8-27B with DFlash2 on M5 Macs: one greedy request at the default 16-row window uses the long-copy driver again, so copy-heavy replies such as a whole-file rewrite run at about 792 tok/s instead of 241 on an M5 Ultra (#3013). Every copied token is still verified by the target, so replies are the same tokens. Sampled requests, concurrent streams, wider windows and M1-M4 Macs keep the lane loop; `TF_QWEN27_LANES=1` keeps every request on lanes (#3069). From 3bef2aca8321efee2cd7147f73e9a27c683b5900 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Fri, 30 Oct 2026 11:30:14 +0000 Subject: [PATCH 06/78] cuda experts: a bf16 format, a checkpoint's own weights in mma fragment blocks (relu^2 up, fp32 or bf16 down) --- zig/build/cuda_kernels.zig | 1 + zig/kernels/cuda/experts_bf16.cu | 114 +++++++++++++++ zig/src/cuda/experts.zig | 12 +- zig/src/cuda/experts_bf16.zig | 94 ++++++++++++ zig/src/cuda/kernels.zig | 1 + zig/tests/cuda/bf16_experts_tests.zig | 201 ++++++++++++++++++++++++++ zig/tests/cuda/main.zig | 2 + 7 files changed, 421 insertions(+), 4 deletions(-) create mode 100644 zig/kernels/cuda/experts_bf16.cu create mode 100644 zig/src/cuda/experts_bf16.zig create mode 100644 zig/tests/cuda/bf16_experts_tests.zig diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index c0c8b503336..950d2d5586f 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -97,6 +97,7 @@ pub const kernels = [_]Kernel{ .{ .name = "exlr8", .flags = &.{ "-O3", "--fmad=false" } }, // ours: EXLR8 lane projections, the Metal kernels' order .{ .name = "exlr8_dense", .flags = &.{ "-O3", "--fmad=false" } }, // ours: EXLR8 dense matrices by row, the Metal order .{ .name = "nvfp4_experts_tile", .flags = &.{"-O3"}, .min_sm = 80 }, + .{ .name = "experts_bf16", .flags = &.{"-O3"} }, // ours: grouped bf16 experts on the plan, mma fragment blocks .{ .name = "nvfp4_experts", .flags = &.{"-O3"} }, // nvfp4/experts.cu's device code, the NVFP4 expert kernels .{ .name = "nvfp4_experts32", .flags = &.{"-O3"} }, // ours: GLM's routed experts at block 32 with csf0 scales .{ .name = "rank_sum", .flags = &.{"-O3"} }, // ours: the rank-order sum of a reduce-scatter (rank_scatter.zig) diff --git a/zig/kernels/cuda/experts_bf16.cu b/zig/kernels/cuda/experts_bf16.cu new file mode 100644 index 00000000000..df6d756ec18 --- /dev/null +++ b/zig/kernels/cuda/experts_bf16.cu @@ -0,0 +1,114 @@ +// Ours: grouped bf16 experts on mma-fragment blocks; a row's sums chain in input order, whatever its tile or form. + +#include +#include + +#include "experts.cuh" + +namespace tf_bf16_experts { + +constexpr int WV = 4; // uint4 a lane a block + +__device__ __forceinline__ uint4 ld_w(const uint4* p) { return tf_common::ld_nc(p); } + +// RT row tiles of 16 pairs share each block's fragments; a unit is one item's 32-column block. +template +__device__ __forceinline__ void body(const __nv_bfloat16* __restrict__ X, int x_stride, int slots, + const uint4* __restrict__ W, int KG, int NB, const int* __restrict__ items, + const int* __restrict__ counts, const int* __restrict__ members, + void* __restrict__ out, int N, int skip) { + const int lane = threadIdx.x & 31, gq = lane >> 2, t = lane & 3; + const int units = __ldg(counts) * NB; + for (int unit = blockIdx.x * WARPS + (threadIdx.x >> 5); unit < units; unit += gridDim.x * WARPS) { + const int it = unit / NB, cb = unit - it * NB; + const int e = __ldg(items + 3 * it), first = __ldg(items + 3 * it + 1), cnt = __ldg(items + 3 * it + 2); + if (e == skip) continue; + const uint4* blk = W + ((size_t)e * NB + cb) * (size_t)KG * (32 * WV) + lane; + for (int r0 = 0; r0 < cnt; r0 += 16 * RT) { + bool v[RT][2]; + int pr[RT][2]; + const __nv_bfloat16* x[RT][2]; +#pragma unroll + for (int r = 0; r < RT; ++r) +#pragma unroll + for (int h = 0; h < 2; ++h) { + const int i = r0 + 16 * r + gq + 8 * h; + v[r][h] = i < cnt; + pr[r][h] = v[r][h] ? __ldg(members + first + i) : 0; + x[r][h] = X + (size_t)(slots ? pr[r][h] / slots : pr[r][h]) * x_stride + 4 * t; + } + float acc[1][RT][NTW][4]; +#pragma unroll + for (int r = 0; r < RT; ++r) +#pragma unroll + for (int j = 0; j < NTW; ++j) acc[0][r][j][0] = acc[0][r][j][1] = acc[0][r][j][2] = acc[0][r][j][3] = 0.f; + const uint2 zero = make_uint2(0u, 0u); + for (int g = 0; g < KG; ++g) { + uint4 w[NTW]; +#pragma unroll + for (int j = 0; j < NTW; ++j) w[j] = ld_w(blk + (size_t)g * (32 * WV) + 32 * j); +#pragma unroll + for (int r = 0; r < RT; ++r) { + uint2 xa[2], xb[2]; +#pragma unroll + for (int q = 0; q < 2; ++q) { + xa[q] = v[r][0] ? __ldg(reinterpret_cast(x[r][0] + 32 * g + 16 * q)) : zero; + xb[q] = v[r][1] ? __ldg(reinterpret_cast(x[r][1] + 32 * g + 16 * q)) : zero; + } +#pragma unroll + for (int q = 0; q < 2; ++q) +#pragma unroll + for (int j = 0; j < NTW; ++j) + mma(acc[0][r][j], xa[q].x, xb[q].x, xa[q].y, xb[q].y, q ? w[j].z : w[j].x, q ? w[j].w : w[j].y); + } + } +#pragma unroll + for (int r = 0; r < RT; ++r) + epilogue(acc, r, out, N, cb * COLS + 2 * t, pr[r][0], pr[r][1], v[r][0], v[r][1], 0.f); + } + } +} + +} // namespace tf_bf16_experts + +// Decode (items of up to 16 pairs) and prompt (up to 64) entries per epilogue: the NVFP4 entries' arguments, no scales. +#define TF_BF16_EXPERTS(NAME, EPI, RT, WARPS) \ + extern "C" __global__ void __launch_bounds__(WARPS * 32) \ + NAME(const __nv_bfloat16* __restrict__ X, int x_stride, int slots, const uint4* __restrict__ W, int KG, int NB, \ + const int* __restrict__ items, const int* __restrict__ counts, const int* __restrict__ members, \ + void* __restrict__ out, int N, int skip) { \ + tf_bf16_experts::body(X, x_stride, slots, W, KG, NB, items, counts, members, out, N, skip); \ + } +TF_BF16_EXPERTS(tf_bf16_experts_up, 1, 1, 4) +TF_BF16_EXPERTS(tf_bf16_experts_down_f32, 0, 1, 4) +TF_BF16_EXPERTS(tf_bf16_experts_down_bf16, 3, 1, 4) +TF_BF16_EXPERTS(tf_bf16_experts_prompt_up, 1, 4, 4) +TF_BF16_EXPERTS(tf_bf16_experts_prompt_down_f32, 0, 4, 4) +TF_BF16_EXPERTS(tf_bf16_experts_prompt_down_bf16, 3, 4, 4) + +// The repack, one thread a uint4: matrix e's [k][n] (in_out, x @ W) or [n][k] (nn.Linear) rows into its blocks. +extern "C" __global__ void tf_bf16_experts_pack(const __nv_bfloat16* __restrict__ src, long long e_stride, int ld, + int n0, int k0, int in_out, int n, int k, uint4* __restrict__ dst) { + const int e = blockIdx.y; + const long long per = (long long)(n / 32) * (k / 32) * 32 * tf_bf16_experts::WV; + const long long q = (long long)blockIdx.x * blockDim.x + threadIdx.x; + if (q >= per) return; + const int lane = (int)(q % 32), j = (int)(q / 32 % tf_bf16_experts::WV); + const long long b = q / (32 * tf_bf16_experts::WV); + const int kg = k / 32, g = (int)(b % kg), cb = (int)(b / kg); + const int row = n0 + cb * 32 + 8 * j + (lane >> 2), col = k0 + g * 32 + 4 * (lane & 3); + const __nv_bfloat16* s = src + e * e_stride; + uint16_t h[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int c = col + (i & 3) + 16 * (i >> 2); + const __nv_bfloat16 v = in_out ? s[(size_t)c * ld + row] : s[(size_t)row * ld + c]; + h[i] = __bfloat16_as_ushort(v); + } + uint4 o; + o.x = h[0] | (uint32_t)h[1] << 16; + o.y = h[2] | (uint32_t)h[3] << 16; + o.z = h[4] | (uint32_t)h[5] << 16; + o.w = h[6] | (uint32_t)h[7] << 16; + dst[e * per + q] = o; +} diff --git a/zig/src/cuda/experts.zig b/zig/src/cuda/experts.zig index 3d70527c1d8..3f7bb4de962 100644 --- a/zig/src/cuda/experts.zig +++ b/zig/src/cuda/experts.zig @@ -14,15 +14,16 @@ const Function = module.Function; const Stream = stream_.Stream; const Plan = grouped.Plan; const Affine3Experts = @import("experts3.zig").Affine3Experts; +pub const Bf16Experts = @import("experts_bf16.zig").Bf16Experts; -pub const Format = enum { affine4, affine3, fp8g, nvfp4, nvfp4_csf0, exl3, iq2xxs }; +pub const Format = enum { affine4, affine3, fp8g, nvfp4, nvfp4_csf0, exl3, iq2xxs, bf16 }; /// One layer's experts on the device as each format's make() lays them out; a family stacks its shared ones last. pub const Layer = struct { format: Format, exl3: ?@import("exl3_experts.zig").Layer = null, - up: u64, // affine4 [E][NI/32][D/64][288] int32; fp8g [E][NI/32][D/32][2][32][2][4]; nvfp4 [E][NI/32][D/32][2][144]; nvfp4_csf0 [E][NI/32] blocks (experts32.zig); iq2xxs [E][NI/16][D/256][1056] bytes (ggml.zig) - down: u64, // affine4 [E][D/32][NI/64][288]; fp8g [E][D/32][NI/32][1][32][2][4]; nvfp4 [E][D/32][NI/32][1][144]; nvfp4_csf0 [E][D/32] blocks; iq2xxs [E][D/16][NI/256][1056] + up: u64, // affine4 [E][NI/32][D/64][288] int32; fp8g [E][NI/32][D/32][2][32][2][4]; nvfp4 [E][NI/32][D/32][2][144]; nvfp4_csf0 [E][NI/32] blocks (experts32.zig); iq2xxs [E][NI/16][D/256][1056] bytes (ggml.zig); bf16 [E][NI/32][D/32][4][32] uint4 + down: u64, // affine4 [E][D/32][NI/64][288]; fp8g [E][D/32][NI/32][1][32][2][4]; nvfp4 [E][D/32][NI/32][1][144]; nvfp4_csf0 [E][D/32] blocks; iq2xxs [E][D/16][NI/256][1056]; bf16 [E][D/32][NI/32][4][32] up_scale: u64 = 0, // fp8g [E][2][NI/128][D/128] fp32; nvfp4 [E][2] global scales; affine4's are in the blocks down_scale: u64 = 0, // fp8g [E][1][D/128][NI/128]; nvfp4 [E][1] width: usize, // NI @@ -52,6 +53,7 @@ pub const Grouped = struct { nvfp4: ?*const Nvfp4Experts = null, nvfp4_csf0: ?*const @import("experts32.zig").Nvfp4b32Experts = null, iq2xxs: ?*const Iq2xxsExperts = null, + bf16: ?*const Bf16Experts = null, /// Decode rows: `items` bounds the plan, `pairs` its picks; `tile` 16 or 64 (EXL3 only); `skip` an expert or -1. pub fn decode(g: Grouped, s: Stream, epi: Epi, in: Rows, l: Layer, p: Plan, out: u64, items: usize, pairs: usize, tile: usize, skip: i32) !void { @@ -77,6 +79,7 @@ pub const Grouped = struct { .nvfp4 => return (g.nvfp4 orelse return error.FormatNotLoaded).run(s, epi, in, l, p, out, items, skip), .nvfp4_csf0 => return (g.nvfp4_csf0 orelse return error.FormatNotLoaded).run(s, epi, in, l, p, out, items, skip), .iq2xxs => return iq2Run(g.iq2xxs orelse return error.FormatNotLoaded, s, false, epi, in, l, p, out, items, skip), + .bf16 => return (g.bf16 orelse return error.FormatNotLoaded).run(s, false, epi, in, l, p, out, items, skip), } } @@ -104,6 +107,7 @@ pub const Grouped = struct { .nvfp4 => return (g.nvfp4 orelse return error.FormatNotLoaded).runPrompt(s, epi, in, l, p, out, items, skip), .nvfp4_csf0 => return (g.nvfp4_csf0 orelse return error.FormatNotLoaded).runPrompt(s, epi, in, l, p, out, items, skip), .iq2xxs => return iq2Run(g.iq2xxs orelse return error.FormatNotLoaded, s, true, epi, in, l, p, out, items, skip), + .bf16 => return (g.bf16 orelse return error.FormatNotLoaded).run(s, true, epi, in, l, p, out, items, skip), } } }; @@ -111,7 +115,7 @@ pub const Grouped = struct { /// Pairs a prompt item holds, per format (the Python engines' prompt tiles). pub fn promptTile(f: Format) usize { return switch (f) { - .affine4, .affine3, .fp8g, .nvfp4, .nvfp4_csf0, .exl3, .iq2xxs => 64, + .affine4, .affine3, .fp8g, .nvfp4, .nvfp4_csf0, .exl3, .iq2xxs, .bf16 => 64, }; } diff --git a/zig/src/cuda/experts_bf16.zig b/zig/src/cuda/experts_bf16.zig new file mode 100644 index 00000000000..76ed554c2f7 --- /dev/null +++ b/zig/src/cuda/experts_bf16.zig @@ -0,0 +1,94 @@ +//! bf16 experts (experts_bf16.cu): a checkpoint's bf16 weights in mma fragment blocks, relu^2 up, fp32 or bf16 down. +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); +const kernels = @import("kernels.zig"); +const grouped = @import("grouped.zig"); +const driver = @import("driver.zig"); +const experts = @import("experts.zig"); + +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; +const Plan = grouped.Plan; +const Epi = experts.Epi; +const Rows = experts.Rows; +const Layer = experts.Layer; + +pub const Bf16Experts = struct { + mod: Module, + fns: [2][3]Function, // decode, prompt; by Epi + resident: [3]usize, // decode blocks the grid may hold: per SM times SMs + pack_fn: Function, + + pub const cols = 32; // output columns a block + pub const warps = 4; + /// Bytes of a (32 columns, 32 inputs) block. + pub const block = 32 * 32 * 2; + + pub const symbols = struct { + pub const decode = [3][:0]const u8{ "tf_bf16_experts_up", "tf_bf16_experts_down_f32", "tf_bf16_experts_down_bf16" }; + pub const prompt = [3][:0]const u8{ "tf_bf16_experts_prompt_up", "tf_bf16_experts_prompt_down_f32", "tf_bf16_experts_prompt_down_bf16" }; + pub const pack = "tf_bf16_experts_pack"; + }; + + pub fn load(d: *const driver.Driver, sms: usize) !Bf16Experts { + if (!kernels.available) return error.BuiltWithoutKernels; + var mod = try Module.load(d, kernels.experts_bf16); + errdefer mod.unload(); + var e: Bf16Experts = .{ .mod = mod, .fns = undefined, .resident = undefined, .pack_fn = try mod.function(symbols.pack) }; + for (symbols.decode, symbols.prompt, 0..) |dn, pn, i| { + e.fns[0][i] = try mod.function(dn); + e.fns[1][i] = try mod.function(pn); + e.resident[i] = @max(1, try e.fns[0][i].occupancy(warps * 32, 0)) * sms; + } + return e; + } + + pub fn unload(e: *Bf16Experts) void { + e.mod.unload(); + } + + /// One call on a plan of at most `items` items: 16 pairs an item in decode, 64 in a prompt. + pub fn run(e: *const Bf16Experts, s: Stream, prompt: bool, epi: Epi, in: Rows, l: Layer, p: Plan, out: u64, items: usize, skip: i32) !void { + if (!l.relu2 or l.swiglu or l.limit != 0) return error.OptionNotBuilt; + const up = epi == .up; + const kg = (if (up) l.dims else l.width) / 32; + const nb = (if (up) l.width else l.dims) / cols; + const ei = @backingInt(epi); + const units = items * nb; + const grid = if (prompt) (units + warps - 1) / warps else @min((units + warps - 1) / warps, e.resident[ei]); + if (grid < 1) return; + var a: launch_.Args = .{}; + a.add(in.x); + a.add(@as(i32, @intCast(in.stride))); + a.add(@as(i32, @intCast(if (up) in.slots else 0))); + a.add(if (up) l.up else l.down); + a.add(@as(i32, @intCast(kg))); + a.add(@as(i32, @intCast(nb))); + for ([_]u64{ p.items, p.counts, p.members, out }) |v| a.add(v); + a.add(@as(i32, @intCast(if (up) l.width else l.dims))); + a.add(skip); + try launch_.launch(e.fns[@intFromBool(prompt)][ei], .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = warps * 32, .y = 1, .z = 1 } }, s, &a); + } + + /// Bytes of `count` experts' blocks for n outputs from k inputs. + pub fn bytes(count: usize, n: usize, k: usize) usize { + return count * (n / cols) * (k / 32) * block; + } + + /// A source matrix: experts `stride` apart, rows of `ld`, its first output and input, [k][n] (x @ W) or [n][k]. + pub const Source = struct { stride: usize = 0, ld: usize, n0: usize = 0, k0: usize = 0, in_out: bool }; + + /// `count` bf16 matrices of n outputs from k inputs at `src` into their blocks at `dst` (bytes(count, n, k)). + pub fn pack(e: *const Bf16Experts, s: Stream, src: u64, from: Source, dst: u64, count: usize, n: usize, k: usize) !void { + if (n % cols != 0 or k % 32 != 0) return error.UnexpectedShape; + const per = bytes(1, n, k) / 16; + var a: launch_.Args = .{}; + a.add(src); + a.add(@as(i64, @intCast(from.stride))); + for ([_]usize{ from.ld, from.n0, from.k0, @intFromBool(from.in_out), n, k }) |v| a.add(@as(i32, @intCast(v))); + a.add(dst); + try launch_.launch(e.pack_fn, .{ .grid = .{ .x = @intCast((per + 255) / 256), .y = @intCast(count), .z = 1 }, .block = .{ .x = 256, .y = 1, .z = 1 } }, s, &a); + } +}; diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 1ff9f6a4489..c8eb892ca12 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -51,6 +51,7 @@ pub const fp8_experts: []const u8 = if (available) &Blob("fatbin_fp8_experts").b pub const exlr8: []const u8 = if (available) &Blob("fatbin_exlr8").bytes else &.{}; pub const exlr8_dense: []const u8 = if (available) &Blob("fatbin_exlr8_dense").bytes else &.{}; pub const nvfp4_experts_tile: []const u8 = if (available) &Blob("fatbin_nvfp4_experts_tile").bytes else &.{}; +pub const experts_bf16: []const u8 = if (available) &Blob("fatbin_experts_bf16").bytes else &.{}; pub const nvfp4_experts: []const u8 = if (available) &Blob("fatbin_nvfp4_experts").bytes else &.{}; pub const nvfp4_experts32: []const u8 = if (available) &Blob("fatbin_nvfp4_experts32").bytes else &.{}; pub const rank_sum: []const u8 = if (available) &Blob("fatbin_rank_sum").bytes else &.{}; diff --git a/zig/tests/cuda/bf16_experts_tests.zig b/zig/tests/cuda/bf16_experts_tests.zig new file mode 100644 index 00000000000..323178e9476 --- /dev/null +++ b/zig/tests/cuda/bf16_experts_tests.zig @@ -0,0 +1,201 @@ +//! Grouped bf16 experts against float64: the repack, up and downs, a row's bits the same at any row count and form. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const Gpu = check.Gpu; + +const grouped = cuda.grouped; +const ex_ = cuda.experts; +const Bf16 = ex_.Bf16Experts; + +const E = 6; // routed experts; the shared expert's halves are experts E and E + 1 +const D = 256; +const W = 96; +const K = 3; // routed picks a row; a row has K + 2 pairs + +fn bf(x: f32) u16 { + const b: u32 = @bitCast(x); + return @intCast((b + 0x7FFF + ((b >> 16) & 1)) >> 16); +} + +fn f(h: u16) f64 { + return @as(f32, @bitCast(@as(u32, h) << 16)); +} + +const Host = struct { + routed_up: []u16, // [E][D][W]: x @ W, as the stacked checkpoint tensors + routed_down: []u16, // [E][W][D] + shared_up: []u16, // [2W][D]: nn.Linear + shared_down: []u16, // [D][2W] + + /// Expert e's weight from input i to output o, up or down, as the kernels must read it. + fn at(h: Host, up: bool, e: usize, o: usize, i: usize) f64 { + if (e < E) return f(if (up) h.routed_up[(e * D + i) * W + o] else h.routed_down[(e * W + i) * D + o]); + const half = e - E; + return f(if (up) h.shared_up[(half * W + o) * D + i] else h.shared_down[o * 2 * W + half * W + i]); + } +}; + +pub fn experts(gpu: Gpu) !void { + const a = gpu.gpa; + var rng = std.Random.DefaultPrng.init(0x7f16); + const r = rng.random(); + var host: Host = undefined; + inline for (.{ "routed_up", "routed_down", "shared_up", "shared_down" }, .{ E * D * W, E * W * D, 2 * W * D, D * 2 * W }) |name, n| { + @field(host, name) = try a.alloc(u16, n); + for (@field(host, name)) |*v| v.* = bf((r.float(f32) - 0.5) * 0.25); + } + defer inline for (.{ "routed_up", "routed_down", "shared_up", "shared_down" }) |name| a.free(@field(host, name)); + + const sms: usize = @intCast(try gpu.ctx.attribute(.multiprocessor_count)); + var kx = try Bf16.load(gpu.d, sms); + defer kx.unload(); + var mod = try cuda.Module.load(gpu.d, cuda.kernels.experts); + defer mod.unload(); + const router = try grouped.Router.resolve(mod); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + + const e = E + 2; + var up = try cuda.DeviceBuffer.alloc(gpu.d, Bf16.bytes(e, W, D)); + defer up.free(); + var down = try cuda.DeviceBuffer.alloc(gpu.d, Bf16.bytes(e, D, W)); + defer down.free(); + var src: [4]cuda.DeviceBuffer = undefined; + inline for (.{ "routed_up", "routed_down", "shared_up", "shared_down" }, 0..) |name, i| src[i] = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(@field(host, name))); + defer for (&src) |*b| b.free(); + const one_up = Bf16.bytes(1, W, D); + const one_down = Bf16.bytes(1, D, W); + try kx.pack(s, src[0].ptr, .{ .stride = D * W, .ld = W, .in_out = true }, up.ptr, E, W, D); + try kx.pack(s, src[1].ptr, .{ .stride = W * D, .ld = D, .in_out = true }, down.ptr, E, D, W); + for (0..2) |h| { + try kx.pack(s, src[2].ptr, .{ .ld = D, .n0 = h * W, .in_out = false }, up.ptr + (E + h) * one_up, 1, W, D); + try kx.pack(s, src[3].ptr, .{ .ld = 2 * W, .k0 = h * W, .in_out = false }, down.ptr + (E + h) * one_down, 1, D, W); + } + try s.synchronize(); + const layer: ex_.Layer = .{ .format = .bf16, .up = up.ptr, .down = down.ptr, .width = W, .dims = D, .experts = e, .relu2 = true }; + const g: ex_.Grouped = .{ .bf16 = &kx }; + + // 37 rows: three 16-pair decode items for the busiest expert and a partial prompt tile; then each row alone + const rows = 37; + const slots = K + 2; + const x = try a.alloc(u16, rows * D); + defer a.free(x); + for (x) |*v| v.* = bf(r.float(f32) * 2 - 1); + const picks = try a.alloc(i32, rows * slots); + defer a.free(picks); + for (0..rows) |row| { + var chosen: [K]i32 = undefined; + for (0..K) |k| { + var p: i32 = undefined; + while (true) { + p = @intCast(r.uintLessThan(usize, E)); + if (std.mem.indexOfScalar(i32, chosen[0..k], p) == null) break; + } + chosen[k] = p; + picks[row * slots + k] = p; + } + picks[row * slots + K] = E; + picks[row * slots + K + 1] = E + 1; + } + + // the float64 reference: act = bf16(relu(bf16(x . up))^2), y = act . down + const act_ref = try a.alloc(f64, rows * slots * W); + defer a.free(act_ref); + const y_ref = try a.alloc(f64, rows * slots * D); + defer a.free(y_ref); + const mag = try a.alloc(f64, rows * slots * D); + defer a.free(mag); + for (0..rows * slots) |pr| { + const ex: usize = @intCast(picks[pr]); + const row = pr / slots; + for (0..W) |o| { + var acc: f64 = 0; + for (0..D) |i| acc += f(x[row * D + i]) * host.at(true, ex, o, i); + const u = @max(f(bf(@floatCast(acc))), 0); + act_ref[pr * W + o] = f(bf(@floatCast(u * u))); + } + for (0..D) |o| { + var acc: f64 = 0; + var m: f64 = 0; + for (0..W) |i| { + const t = act_ref[pr * W + i] * host.at(false, ex, o, i); + acc += t; + m += @abs(t); + } + y_ref[pr * D + o] = acc; + mag[pr * D + o] = m; + } + } + + var dx = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(x)); + defer dx.free(); + var whole: [3][]u8 = undefined; // decode form, all rows: act, y fp32, y bf16 + var got = [3]bool{ false, false, false }; + defer for (whole, got) |w, ok| if (ok) a.free(w); + for ([_]bool{ false, true }) |prompt| { + const run = try oneCase(gpu, router, g, s, layer, dx.ptr, picks, rows, slots, prompt); + defer for (run) |b| a.free(b); + const act_h = std.mem.bytesAsSlice(u16, run[0]); + var worst: f64 = 0; + for (act_h, act_ref) |gv, want| worst = @max(worst, @abs(f(gv) - want) / @max(@abs(want), 1e-3)); + try check.expect(worst <= 1.0 / 64.0, "bf16 experts {s}: up off by {e} (relative)", .{ if (prompt) "prompt" else "decode", worst }); + const y32 = std.mem.bytesAsSlice(f32, run[1]); + var worst_y: f64 = 0; + for (y32, y_ref, mag) |gv, want, m| worst_y = @max(worst_y, @abs(@as(f64, gv) - want) / @max(m, 1e-6)); + try check.expect(worst_y <= 1e-5, "bf16 experts {s}: fp32 down off by {e} of its terms' magnitude", .{ if (prompt) "prompt" else "decode", worst_y }); + if (!prompt) { + for (run, &whole, &got) |b, *w, *ok| { + w.* = try a.dupe(u8, b); + ok.* = true; + } + } else for (run, whole, [_][]const u8{ "prompt up", "prompt fp32 down", "prompt bf16 down" }) |b, w, what| try check.sameBytes(what, b, w); + check.pass("bf16 experts: {d} rows x {d} pairs ({s}) up within {e}, fp32 down within {e} of the float64 reference", .{ rows, slots, if (prompt) "prompt, items of 64" else "decode, items of 16", worst, worst_y }); + } + // each row alone gives the bytes it gave beside the others + for ([_]usize{ 0, 17, rows - 1 }) |row| { + const run = try oneCase(gpu, router, g, s, layer, dx.ptr + row * D * 2, picks[row * slots ..][0..slots], 1, slots, false); + defer for (run) |b| a.free(b); + for (run, whole, [_]usize{ W * 2, D * 4, D * 2 }, [_][]const u8{ "row up", "row fp32 down", "row bf16 down" }) |b, w, per, what| try check.sameBytes(what, b, w[row * slots * per ..][0 .. slots * per]); + } + check.pass("bf16 experts: a row alone, beside 36 others and in the prompt form gives the same bytes", .{}); +} + +/// Plan, up and both downs over `rows` rows; returns act, fp32 y and bf16 y. +fn oneCase(gpu: Gpu, router: grouped.Router, g: ex_.Grouped, s: cuda.Stream, l: ex_.Layer, x: u64, picks: []const i32, rows: usize, slots: usize, prompt: bool) ![3][]u8 { + const pairs = rows * slots; + const tile: usize = if (prompt) ex_.promptTile(.bf16) else 16; + var dpicks = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(picks)); + defer dpicks.free(); + const bound = grouped.maxItems(pairs, l.experts, tile); + var bufs: [5]cuda.DeviceBuffer = undefined; + const big = pairs > grouped.small; + const sizes = [5]usize{ pairs * 4, bound * 12, 8, (if (big) pairs else 1) * 4, (if (big) (pairs + 1023) / 1024 * l.experts else 1) * 4 }; + var made: usize = 0; + defer for (bufs[0..made]) |*b| b.free(); + for (sizes, &bufs) |n, *b| { + b.* = try cuda.DeviceBuffer.alloc(gpu.d, n); + made += 1; + } + const plan: grouped.Plan = .{ .members = bufs[0].ptr, .items = bufs[1].ptr, .counts = bufs[2].ptr, .rank = bufs[3].ptr, .hist = bufs[4].ptr }; + try router.route(s, dpicks.ptr, pairs, l.experts, tile, plan); + var act = try cuda.DeviceBuffer.alloc(gpu.d, pairs * l.width * 2); + defer act.free(); + var y32 = try cuda.DeviceBuffer.alloc(gpu.d, pairs * l.dims * 4); + defer y32.free(); + var y16 = try cuda.DeviceBuffer.alloc(gpu.d, pairs * l.dims * 2); + defer y16.free(); + for ([_]*cuda.DeviceBuffer{ &act, &y32, &y16 }) |b| try b.fill8(0xff, s.handle); // a call that writes nothing fails + const call = if (prompt) &ex_.Grouped.prompt else &ex_.Grouped.decode; + try call(g, s, .up, .{ .x = x, .stride = l.dims, .slots = slots }, l, plan, act.ptr, bound, pairs, tile, -1); + try call(g, s, .down_f32, .{ .x = act.ptr, .stride = l.width }, l, plan, y32.ptr, bound, pairs, tile, -1); + try call(g, s, .down_bf16, .{ .x = act.ptr, .stride = l.width }, l, plan, y16.ptr, bound, pairs, tile, -1); + try s.synchronize(); + var out: [3][]u8 = undefined; + out[0] = try check.download(gpu, act); + errdefer gpu.gpa.free(out[0]); + out[1] = try check.download(gpu, y32); + errdefer gpu.gpa.free(out[1]); + out[2] = try check.download(gpu, y16); + return out; +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index f9f39206584..ebd3f880790 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -215,6 +215,7 @@ const usage = \\ fp8-lane block-FP8 projections against the Python lane matmul's bytes (oracle/fp8_lane.py) \\ fp8-experts grouped block-FP8 experts against recorded fixture bytes \\ nvfp4-experts grouped NVFP4 experts against recorded fixture bytes + \\ bf16-experts grouped bf16 experts against float64: repack, up, downs, rows and forms byte for byte \\ projection-f32 FP32 projection outputs against FP64, rounding and row/reduction invariance \\ mxfp8 MXFP8 strided projections against FP64 and row-count invariance \\ quant-ref FP8 and NVFP4 projections and experts against float64 references (synthetic) @@ -430,6 +431,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "fp8-experts")) return fp8_experts_tests.experts(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "nvfp4")) return nvfp4_tests.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "nvfp4-experts")) return nvfp4_experts_tests.experts(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "bf16-experts")) return @import("bf16_experts_tests.zig").experts(gpu); if (std.mem.eql(u8, cmd, "projection-f32")) return projection_f32_tests.run(gpu); if (std.mem.eql(u8, cmd, "mxfp8")) return mxfp8_tests.run(gpu); if (std.mem.eql(u8, cmd, "quant-ref")) return quant_ref_tests.run(gpu); From ce603be2905987903a40a5deb77a282e54053fd9 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Fri, 30 Oct 2026 11:30:14 +0000 Subject: [PATCH 07/78] kernel contract: the bf16 experts (experts_bf16.cu), relu^2 up, fp32 and bf16 down, gates 1-5 with two controls --- docs/kernel-contract.md | 4 + zig/kernels/cuda/contract_controls_set.cu | 1 + .../cuda/contract_experts_bf16_controls.cuh | 24 ++ .../core/contract/experts_bf16_records.zig | 37 ++ zig/src/core/contract/registry_joins.zig | 5 +- zig/src/core/contract/registry_sets_test.zig | 1 + zig/tests/cuda/contract_experts_bf16.zig | 319 ++++++++++++++++++ zig/tests/cuda/contract_family_sets.zig | 5 + zig/tests/cuda/contract_limits.zig | 2 +- zig/tests/cuda/main.zig | 4 +- 10 files changed, 398 insertions(+), 4 deletions(-) create mode 100644 zig/kernels/cuda/contract_experts_bf16_controls.cuh create mode 100644 zig/src/core/contract/experts_bf16_records.zig create mode 100644 zig/tests/cuda/contract_experts_bf16.zig diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 6b4d03873ff..520a89bc5d5 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -401,6 +401,10 @@ CUDA also checks the affine-3 experts' split decode pair, bit-equal to the pair one dropping a last round of K groups short of the warps, one chaining the groups in reverse). CUDA also checks Flash Next's NVFP4 experts on prompt tiles with the decode form's sums (SwiGLU gate-up and fp32 down, each bit-equal to the decode instance it replaces; the scaled-code prompt tile is the planted control). +CUDA also checks the bf16 experts (`experts_bf16.cu`, a checkpoint's own bf16 weights: Nemotron 3.5 Transcribe's routed +and shared experts): relu^2 up, fp32 down and bf16 down, in the decode and prompt forms, against float64 (two controls: +the fp32 down launched one 32-input group short, and a down summing in bf16). nvcc 13.2's prompt forms and 13.3's +decode forms on sm_120 and sm_121 take other registers than 13.0's and have their own records. CUDA also checks Nemotron's pocket (`nemotron_ops.cu`, the learned neurons the living-weights learner adds beside the shared expert): relu^2 up into act rows and down added into the shared expert's slot, a block row a row and tiled, the down to fp32 and bf16, against float64, the neurons past the count NaN so a read of one shows (two controls: the up's diff --git a/zig/kernels/cuda/contract_controls_set.cu b/zig/kernels/cuda/contract_controls_set.cu index e4608a5b9f7..3624ac10b03 100644 --- a/zig/kernels/cuda/contract_controls_set.cu +++ b/zig/kernels/cuda/contract_controls_set.cu @@ -3,3 +3,4 @@ #include "contract_controls.cu" #include "contract_js_controls.cuh" #include "contract_dense_controls.cuh" +#include "contract_experts_bf16_controls.cuh" diff --git a/zig/kernels/cuda/contract_experts_bf16_controls.cuh b/zig/kernels/cuda/contract_experts_bf16_controls.cuh new file mode 100644 index 00000000000..1af8126ec29 --- /dev/null +++ b/zig/kernels/cuda/contract_experts_bf16_controls.cuh @@ -0,0 +1,24 @@ +// Gate 4's control for experts_bf16.cu: its fp32 down with each pair's sum rounded to bf16 at every input. +#pragma once + +extern "C" __global__ void tf_ctl_bf16_experts_bf16_sums(const __nv_bfloat16* X, int x_stride, const uint4* W, int KG, int NB, + const int* items, const int* counts, const int* members, float* out, + int N) { + const int col = blockIdx.x * blockDim.x + threadIdx.x; + if (col >= N) return; + const int cb = col / 32, j = (col % 32) / 8; + for (int it = blockIdx.y; it < counts[0]; it += gridDim.y) { + const int e = items[3 * it], first = items[3 * it + 1], cnt = items[3 * it + 2]; + const __nv_bfloat16* blk = reinterpret_cast(W + (static_cast(e) * NB + cb) * KG * 128); + for (int i = 0; i < cnt; ++i) { + const int pr = members[first + i]; + float s = 0.0f; + for (int k = 0; k < KG * 32; ++k) { + const int q = k % 16, lane = 4 * (col % 8) + q / 4; + const float w = __bfloat162float(blk[(static_cast(k / 32) * 128 + 32 * j + lane) * 8 + 4 * (k % 32 / 16) + q % 4]); + s = __bfloat162float(__float2bfloat16_rn(s + tf_ctl_x(X, static_cast(pr) * x_stride + k) * w)); + } + out[static_cast(pr) * N + col] = s; + } + } +} diff --git a/zig/src/core/contract/experts_bf16_records.zig b/zig/src/core/contract/experts_bf16_records.zig new file mode 100644 index 00000000000..7b4c4bf7784 --- /dev/null +++ b/zig/src/core/contract/experts_bf16_records.zig @@ -0,0 +1,37 @@ +//! The bf16 experts' CUDA launches (experts_bf16.cu): decode and prompt rows by epilogue, and the repack. +const registry = @import("registry.zig"); + +/// nvcc 13.0's usage, the compiler of the release fatbins. +pub const cuda = [_]registry.CudaRecord{ + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_up", .symbol = "tf_bf16_experts_up", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 76, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_down_f32", .symbol = "tf_bf16_experts_down_f32", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 76, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_down_bf16", .symbol = "tf_bf16_experts_down_bf16", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 76, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 }, .{ .regs = 72, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_prompt_up", .symbol = "tf_bf16_experts_prompt_up", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_prompt_down_f32", .symbol = "tf_bf16_experts_prompt_down_f32", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_prompt_down_bf16", .symbol = "tf_bf16_experts_prompt_down_bf16", .image = "experts_bf16", .threads = 128 }, .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 }, .{ .regs = 162, .static = 0, .bound = 128 } } }, + .{ .launch = .{ .name = "experts_bf16 tf_bf16_experts_pack", .symbol = "tf_bf16_experts_pack", .image = "experts_bf16", .threads = 256 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 40, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 }, .{ .regs = 36, .static = 0, .bound = 0 } } }, +}; + +/// nvcc 13.2's prompt forms on sm_120 and sm_121 (164 registers where 13.0 takes 162); the rest compile as 13.0's. +pub const cuda_132 = [_]registry.Override{ + .{ .name = "experts_bf16 tf_bf16_experts_prompt_up", .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 } } }, + .{ .name = "experts_bf16 tf_bf16_experts_prompt_down_f32", .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 } } }, + .{ .name = "experts_bf16 tf_bf16_experts_prompt_down_bf16", .usage = .{ .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 168, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 }, .{ .regs = 164, .static = 0, .bound = 128 } } }, +}; + +/// nvcc 13.3's decode forms on sm_120 and sm_121 (70 registers where 13.0 and 13.2 take 72). +pub const cuda_133 = [_]registry.Override{ + .{ .name = "experts_bf16 tf_bf16_experts_up", .usage = .{ .{ .regs = 76, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 70, .static = 0, .bound = 128 }, .{ .regs = 70, .static = 0, .bound = 128 } } }, + .{ .name = "experts_bf16 tf_bf16_experts_down_f32", .usage = .{ .{ .regs = 76, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 70, .static = 0, .bound = 128 }, .{ .regs = 70, .static = 0, .bound = 128 } } }, + .{ .name = "experts_bf16 tf_bf16_experts_down_bf16", .usage = .{ .{ .regs = 76, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 84, .static = 0, .bound = 128 }, .{ .regs = 70, .static = 0, .bound = 128 }, .{ .regs = 70, .static = 0, .bound = 128 } } }, +}; + +test "bf16 experts CUDA: every launch is in the live registry and fits every SM" { + const t = @import("std").testing; + for (cuda) |r| { + var found: usize = 0; + for (registry.all_cuda) |active| found += @intFromBool(@import("std").mem.eql(u8, active.launch.name, r.launch.name)); + try t.expectEqual(@as(usize, 1), found); + try t.expectEqual(@as(usize, 0), try registry.checkCuda(t.allocator, &.{r}, true)); + } +} diff --git a/zig/src/core/contract/registry_joins.zig b/zig/src/core/contract/registry_joins.zig index 615190a1678..782618b89bd 100644 --- a/zig/src/core/contract/registry_joins.zig +++ b/zig/src/core/contract/registry_joins.zig @@ -13,6 +13,7 @@ pub const cuda_after = @import("registry_media.zig").cuda ++ @import("modelopt_lane_records.zig").cuda ++ @import("nvfp4_fact_records.zig").cuda ++ @import("nvfp4_experts32_records.zig").cuda ++ + @import("experts_bf16_records.zig").cuda ++ @import("lane_half_records.zig").cuda ++ @import("pocket_records.zig").cuda ++ @import("minimax_cuda_records.zig").cuda ++ @@ -30,6 +31,7 @@ pub const nvcc_132_after = @import("registry_learning.zig").cuda_132 ++ @import("qwen_served_records.zig").cuda_132 ++ @import("flashnext_tile_records.zig").cuda_132 ++ @import("nvfp4_fact_records.zig").cuda_132 ++ + @import("experts_bf16_records.zig").cuda_132 ++ @import("pocket_records.zig").cuda_132 ++ @import("minimax_cuda_records.zig").cuda_132; pub const nvcc_133_after = @import("flashnext_tile_records.zig").cuda_133 ++ @@ -37,7 +39,8 @@ pub const nvcc_133_after = @import("flashnext_tile_records.zig").cuda_133 ++ @import("nvcc_133_records.zig").cuda_133 ++ @import("minimax_cuda_records.zig").cuda_133 ++ @import("nvfp4_experts32_records.zig").cuda_133 ++ - @import("paging_pending.zig").cuda_133; + @import("paging_pending.zig").cuda_133 ++ + @import("experts_bf16_records.zig").cuda_133; pub const cuda_pending = @import("qwen_whole_records.zig").records ++ @import("exl3_prompt_pending.zig").records ++ @import("cuda_lane_pending.zig").records ++ diff --git a/zig/src/core/contract/registry_sets_test.zig b/zig/src/core/contract/registry_sets_test.zig index 743c4b1a8cf..43979dd1ab3 100644 --- a/zig/src/core/contract/registry_sets_test.zig +++ b/zig/src/core/contract/registry_sets_test.zig @@ -54,4 +54,5 @@ test { _ = @import("exl3_grouped_records.zig"); _ = @import("proposal_sampling_records.zig"); _ = @import("pocket_records.zig"); + _ = @import("experts_bf16_records.zig"); } diff --git a/zig/tests/cuda/contract_experts_bf16.zig b/zig/tests/cuda/contract_experts_bf16.zig new file mode 100644 index 00000000000..4eb411f98b9 --- /dev/null +++ b/zig/tests/cuda/contract_experts_bf16.zig @@ -0,0 +1,319 @@ +//! bf16 experts (experts_bf16.cu) under the contract on Transcribe's shapes, decode and prompt forms; two controls. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("check.zig"); +const rig_ = @import("contract_rig.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const grouped = cuda.grouped; +const ex_ = cuda.experts; +const Bf16 = ex_.Bf16Experts; + +pub const Kind = enum { up, down_f32, down_bf16 }; + +/// A planted fault: the fp32 down launched one 32-input group short, or contract_controls.cu's bf16-summing down. +pub const Fault = enum { none, short, bf16_sums }; + +pub const Entry = struct { name: []const u8, kind: Kind, prompt: bool, fault: Fault = .none }; + +pub const entries = [_]Entry{ + .{ .name = "cuda bf16 experts up tf_bf16_experts_up (relu^2)", .kind = .up, .prompt = false }, + .{ .name = "cuda bf16 experts down tf_bf16_experts_down_f32", .kind = .down_f32, .prompt = false }, + .{ .name = "cuda bf16 experts down tf_bf16_experts_down_bf16", .kind = .down_bf16, .prompt = false }, + .{ .name = "cuda bf16 experts prompt up tf_bf16_experts_prompt_up (relu^2)", .kind = .up, .prompt = true }, + .{ .name = "cuda bf16 experts prompt down tf_bf16_experts_prompt_down_f32", .kind = .down_f32, .prompt = true }, + .{ .name = "cuda bf16 experts prompt down tf_bf16_experts_prompt_down_bf16", .kind = .down_bf16, .prompt = true }, + .{ .name = "control: bf16 experts fp32 down launched one 32-input group short", .kind = .down_f32, .prompt = false, .fault = .short }, + .{ .name = "control: bf16 experts fp32 down summing in bf16", .kind = .down_f32, .prompt = false, .fault = .bf16_sums }, +}; + +const experts = 4; +const slots = 8; // pairs a token row: Transcribe's 6 routed picks and its shared expert's two halves + +/// An mma k16 step a link of each row's chain, plus one step's error over its own products (attention_range's 36). +fn depth(s: c.Shape) usize { + return s.k / 16 + core.attention_range.mma_step; +} + +/// Transcribe's routed experts, up 2,688 to 1,856 and back; decode windows of 1 to 32 rows, prompt items of 2 and 9. +const ups = [2][]const c.Shape{ &.{ .{ .rows = 1, .n = 1856, .k = 2688 }, .{ .rows = 32, .n = 1856, .k = 2688 } }, &.{ .{ .rows = 40, .n = 1856, .k = 2688 }, .{ .rows = 72, .n = 1856, .k = 2688 } } }; +const downs = [2][]const c.Shape{ &.{ .{ .rows = 1, .n = 2688, .k = 1856 }, .{ .rows = 32, .n = 2688, .k = 1856 } }, &.{ .{ .rows = 40, .n = 2688, .k = 1856 }, .{ .rows = 72, .n = 2688, .k = 1856 } } }; + +pub fn spec(e: Entry) c.Spec { + const served = (if (e.kind == .up) ups else downs)[@intFromBool(e.prompt)]; + const tiles = [3]c.Tile{ .{ .size = (if (e.prompt) ex_.promptTile(.bf16) else 16) / slots }, .{ .size = 64, .step = 32, .min = 32 }, .{ .size = 64, .step = 32, .min = 32 } }; + return switch (e.kind) { + // bf16(acc)'s rounding enters the square twice + .up => .{ .name = e.name, .accum = .f32, .depth = depth, .ops = 1, .mids = 2, .out = .bf16, .tiles = tiles, .served = served, .sums = .{ .f16 = 30, .bf16 = 30, .why = "a sum that collapses below 0 gives relu^2 0, which misses by want itself, and the bf16 roundings already allow 3 u of it: at most 1 / (3 u) = 85x is provable" } }, + .down_f32 => .{ .name = e.name, .accum = .f32, .depth = depth, .out = .f32, .tiles = tiles, .served = served }, + .down_bf16 => .{ .name = e.name, .accum = .f32, .depth = depth, .out = .bf16, .tiles = tiles, .served = served }, + }; +} + +/// shape.rows token rows, `slots` pairs each over `experts` experts; up reads a token's row, down a pair's. +pub const Experts = struct { + rig: rig_.Rig, + entry: Entry, + kx: *const Bf16, + router: grouped.Router, + ctl: ?cuda.Function = null, + shape: c.Shape = undefined, + x: []u16 = &.{}, + w: []u16 = &.{}, // [experts][n][k], as nn.Linear stores them + picks: []i32 = &.{}, + xd: u64 = 0, + wd: u64 = 0, + plan: grouped.Plan = undefined, + out: u64 = 0, + + pub fn deinit(self: *Experts) void { + self.free(); + self.rig.deinit(); + } + + fn free(self: *Experts) void { + const a = self.rig.gpu.gpa; + a.free(self.x); + a.free(self.w); + a.free(self.picks); + self.x = &.{}; + self.w = &.{}; + self.picks = &.{}; + } + + pub fn kernel(self: *Experts) g.Kernel { + const control: ?g.Gate = switch (self.entry.fault) { + .none => null, + .short => .shapes, + .bf16_sums => .sums, + }; + return .{ .spec = spec(self.entry), .ctx = self, .vt = &vt, .control = control }; + } + + const vt = g.hooks(Experts); + + fn up(self: *const Experts) bool { + return self.entry.kind == .up; + } + + fn pairs(self: *const Experts) usize { + return self.shape.rows * slots; + } + + fn xRows(self: *const Experts) usize { + return if (self.up()) self.shape.rows else self.pairs(); + } + + fn tile(self: *const Experts) usize { + return if (self.entry.prompt) ex_.promptTile(.bf16) else 16; + } + + fn items(self: *const Experts) usize { + return grouped.maxItems(self.pairs(), experts, self.tile()); + } + + /// Random bf16 rows and weights; non-finite: NaN, +Inf and -Inf in the first rows; stress: mirrored integers. + pub fn prepare(self: *Experts, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.n % 32 != 0 or shape.k == 0 or shape.k % 32 != 0) return false; + self.free(); + self.rig.reset(); + self.shape = shape; + const a = self.rig.gpu.gpa; + const n = shape.n; + const k = shape.k; + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + self.x = try a.alloc(u16, self.xRows() * k); + self.w = try a.alloc(u16, experts * n * k); + self.picks = try a.alloc(i32, self.pairs()); + for (self.x) |*v| v.* = c.toBf16(rng.float(f32) * 2 - 1); + for (self.w) |*v| v.* = c.toBf16((rng.float(f32) - 0.5) * 0.25); + for (self.picks, 0..) |*p, i| p.* = @intCast(if (i < experts) i else rng.uintLessThan(usize, experts)); + switch (fill) { + .random => {}, + .nonfinite => { + // a row holding NaN, one holding +Inf, one holding both infinities + const planted = [_]struct { usize, usize, f32 }{ .{ 0, k / 3, std.math.nan(f32) }, .{ 1, k / 2, std.math.inf(f32) }, .{ 2, 1, std.math.inf(f32) }, .{ 2, k - 1, -std.math.inf(f32) } }; + for (planted) |p| if (p[0] < self.xRows()) { + self.x[p[0] * k + p[1]] = c.toBf16(p[2]); + }; + }, + .stress => self.stress(rng), + } + self.xd = try self.rig.input(std.mem.sliceAsBytes(self.x)); + const src = try self.rig.scratch(self.w.len * 2); + try (cuda.DeviceBuffer{ .d = self.rig.gpu.d, .ptr = src, .len = self.w.len * 2 }).upload(0, std.mem.sliceAsBytes(self.w)); + const blocks = Bf16.bytes(experts, n, k); + const zeros = try a.alloc(u8, @max(blocks, self.pairs() * 4, self.items() * 12)); + defer a.free(zeros); + @memset(zeros, 0); + self.wd = try self.rig.input(zeros[0..blocks]); + try self.kx.pack(self.rig.s, src, .{ .stride = n * k, .ld = k, .in_out = false }, self.wd, experts, n, k); + const ids = try self.rig.input(std.mem.sliceAsBytes(self.picks)); + const hist = (self.pairs() + 1023) / 1024; + self.plan = .{ .members = try self.rig.input(zeros[0 .. self.pairs() * 4]), .items = try self.rig.input(zeros[0 .. self.items() * 12]), .counts = try self.rig.input(zeros[0..8]), .rank = try self.rig.scratch(self.pairs() * 4), .hist = try self.rig.scratch(hist * experts * 4) }; + try self.router.route(self.rig.s, ids, self.pairs(), experts, self.tile(), self.plan); + try self.rig.s.synchronize(); + self.out = try self.rig.output(self.pairs() * n * c.size(spec(self.entry).out)); + return true; + } + + /// Gate 4: mirrored integer weights and inputs, every partial sum an exact integer under 2^20, the row's sum small. + fn stress(self: *Experts, rng: std.Random) void { + const k = self.shape.k; + const h = k / 2; + for (0..experts * self.shape.n) |j| for (0..h) |i| { + const v: f32 = @floatFromInt(1 + rng.uintLessThan(u32, 2)); + self.w[j * k + i] = c.toBf16(v); + self.w[j * k + h + i] = c.toBf16(v); + }; + for (0..self.xRows()) |r| for (0..h) |i| { + const v: f32 = @floatFromInt(4 * (1 + 2 * rng.uintLessThan(u32, 64))); + self.x[r * k + i] = c.toBf16(v); + self.x[r * k + h + i] = c.toBf16(if (i < 8) 4 - v else -v); + }; + } + + pub fn run(self: *Experts) !void { + const n = self.shape.n; + const k = self.shape.k; + const in: ex_.Rows = .{ .x = self.xd, .stride = k, .slots = if (self.up()) slots else 0 }; + switch (self.entry.fault) { + .none => { + const epi: ex_.Epi = switch (self.entry.kind) { + .up => .up, + .down_f32 => .down_f32, + .down_bf16 => .down_bf16, + }; + const l: ex_.Layer = .{ .format = .bf16, .up = self.wd, .down = self.wd, .width = if (self.up()) n else k, .dims = if (self.up()) k else n, .experts = experts, .relu2 = true }; + try self.kx.run(self.rig.s, self.entry.prompt, epi, in, l, self.plan, self.out, self.items(), -1); + }, + .short => try self.launch(self.kx.fns[0][1], k / 32 - 1, .{ .x = @intCast((self.items() * n / Bf16.cols + Bf16.warps - 1) / Bf16.warps) }, Bf16.warps * 32), + .bf16_sums => try self.launch(self.ctl.?, k / 32, .{ .x = @intCast((n + 127) / 128), .y = @intCast(self.items()) }, 128), + } + try self.rig.s.synchronize(); + } + + /// A down kernel's launch with its own input groups: experts_bf16.cu's arguments, less `slots` for the control. + fn launch(self: *Experts, f: cuda.Function, kg: usize, grid: struct { x: u32, y: u32 = 1 }, threads: u32) !void { + var a: cuda.Args = .{}; + a.add(self.xd); + a.add(@as(i32, @intCast(self.shape.k))); + if (self.entry.fault == .short) a.add(@as(i32, 0)); + a.add(self.wd); + a.add(@as(i32, @intCast(kg))); + a.add(@as(i32, @intCast(self.shape.n / Bf16.cols))); + for ([_]u64{ self.plan.items, self.plan.counts, self.plan.members, self.out }) |v| a.add(v); + a.add(@as(i32, @intCast(self.shape.n))); + if (self.entry.fault == .short) a.add(@as(i32, -1)); + try cuda.launch.launch(f, .{ .grid = .{ .x = grid.x, .y = grid.y, .z = 1 }, .block = .{ .x = threads } }, self.rig.s, &a); + } + + pub fn count(self: *Experts) usize { + return self.pairs() * self.shape.n; + } + + /// Pair p's products with column j of its expert, in input order. + fn terms(self: *const Experts, p: usize, j: usize, t: []f64) void { + const k = self.shape.k; + const xr = (if (self.up()) p / slots else p) * k; + const wr = (@as(usize, @intCast(self.picks[p])) * self.shape.n + j) * k; + for (t, 0..) |*o, i| o.* = c.fromBf16(self.x[xr + i]) * c.fromBf16(self.w[wr + i]); + } + + /// relu^2 for up, as float64 keeps a NaN; down's sum as it is. + fn epilogue(self: *const Experts, sum: f64) f64 { + if (!self.up()) return sum; + const u = if (std.math.isNan(sum)) sum else @max(sum, 0); + return u * u; + } + + pub fn reference(self: *Experts, want: []f64, mag: []f64) !void { + const t = try self.rig.gpu.gpa.alloc(f64, self.shape.k); + defer self.rig.gpu.gpa.free(t); + for (0..self.pairs()) |p| for (0..self.shape.n) |j| { + self.terms(p, j, t); + var sum: f64 = 0; + var abs: f64 = 0; + for (t) |v| { + sum += v; + abs += @abs(v); + } + const o = p * self.shape.n + j; + want[o] = self.epilogue(sum); + // a sum's error reaches relu^2 through its derivative 2 u + mag[o] = if (self.up()) 2 * @sqrt(want[o]) * abs else abs; + }; + } + + pub fn simulate(self: *Experts, acc: c.Accum, order: g.Order, out: []f64) !void { + const t = try self.rig.gpu.gpa.alloc(f64, self.shape.k); + defer self.rig.gpu.gpa.free(t); + for (0..self.pairs()) |p| for (0..self.shape.n) |j| { + self.terms(p, j, t); + out[p * self.shape.n + j] = self.epilogue(contract.rows.sumIn(acc, order, t)); + }; + } +}; + +/// Every registration, on the heap: the gates hold pointers to them. +pub const Set = struct { + gpa: std.mem.Allocator, + kx: Bf16, + mod: cuda.Module, + ctl: cuda.Module, + items: [entries.len]Experts = undefined, + + const controls = struct { + pub const bytes align(16) = @embedFile("fatbin_contract_controls").*; + }; + + /// The set, its kernels appended to `list` (each control after the kernels it breaks). + pub fn init(gpu: check.Gpu, s: cuda.Stream, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel)) !*Set { + const set = try gpu.gpa.create(Set); + errdefer gpu.gpa.destroy(set); + const sms: usize = @intCast(try gpu.ctx.attribute(.multiprocessor_count)); + var kx = try Bf16.load(gpu.d, sms); + errdefer kx.unload(); + var mod = try cuda.Module.load(gpu.d, cuda.kernels.experts); + errdefer mod.unload(); + var ctl = try cuda.Module.load(gpu.d, &controls.bytes); + errdefer ctl.unload(); + set.* = .{ .gpa = gpu.gpa, .kx = kx, .mod = mod, .ctl = ctl }; + const router = try grouped.Router.resolve(set.mod); + const sums = try set.ctl.function("tf_ctl_bf16_experts_bf16_sums"); + for (&set.items, entries) |*x, e| x.* = .{ .rig = .{ .gpu = gpu, .s = s }, .entry = e, .kx = &set.kx, .router = router, .ctl = if (e.fault == .bf16_sums) sums else null }; + for (&set.items) |*x| try list.append(a, x.kernel()); + return set; + } + + pub fn deinit(set: *Set) void { + for (&set.items) |*x| x.deinit(); + set.ctl.unload(); + set.mod.unload(); + set.kx.unload(); + set.gpa.destroy(set); + } +}; + +/// The bf16 experts' gates 1-4 and their controls alone. +pub fn run(gpu: check.Gpu) !void { + const a = gpu.gpa; + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var list: std.ArrayList(g.Kernel) = .empty; + defer list.deinit(a); + const set = try Set.init(gpu, s, a, &list); + defer set.deinit(); + var t: g.Tally = .{}; + try g.runAll(a, list.items, 0x62663136, &t); + std.debug.print("SUMMARY bf16 experts: {d} kernels, {d} pass gates 1-4, {d} fail, {d} known findings; {d} controls: {d} red as planted, {d} missed\n", .{ t.kernels, t.passed, t.failed, t.known, t.controls, t.caught, t.missed }); + try check.expect(t.problems() == 0, "bf16-experts contract: {d} gate failures, {d} controls missed", .{ t.failed, t.missed }); +} + +test "every control is named one, and only the controls plant a fault" { + for (entries) |e| try std.testing.expectEqual(e.fault != .none, std.mem.startsWith(u8, e.name, "control: ")); +} diff --git a/zig/tests/cuda/contract_family_sets.zig b/zig/tests/cuda/contract_family_sets.zig index 666968e5ba2..ed1bdc8d345 100644 --- a/zig/tests/cuda/contract_family_sets.zig +++ b/zig/tests/cuda/contract_family_sets.zig @@ -10,6 +10,7 @@ const Grouped = @import("contract_gb10_many.zig").Registrations; pub const Sets = struct { gpa: std.mem.Allocator, + bf16_experts: *@import("contract_experts_bf16.zig").Set, pocket: *@import("contract_pocket.zig").Set, minimax: *@import("contract_minimax.zig").Set, moss: *@import("contract_moss_image.zig").Set, @@ -23,6 +24,8 @@ pub const Sets = struct { /// Each set's kernels appended to `list`, in this order. pub fn init(gpu: check.Gpu, s: cuda.Stream, a: std.mem.Allocator, list: *std.ArrayList(g.Kernel), affine4: *cuda.qlinear.Affine4Kernels, cap: u32) !Sets { + const bf16_experts = try @import("contract_experts_bf16.zig").Set.init(gpu, s, a, list); + errdefer bf16_experts.deinit(); const pocket = try @import("contract_pocket.zig").Set.init(gpu, s, a, list); errdefer pocket.deinit(); const minimax = try @import("contract_minimax.zig").Set.init(gpu, s, a, list); @@ -57,6 +60,7 @@ pub const Sets = struct { const joint_schema = try @import("contract_joint_schema.zig").Set.init(gpu, s, a, list); return .{ .gpa = gpu.gpa, + .bf16_experts = bf16_experts, .pocket = pocket, .minimax = minimax, .moss = moss, @@ -84,5 +88,6 @@ pub const Sets = struct { sets.moss.deinit(); sets.minimax.deinit(); sets.pocket.deinit(); + sets.bf16_experts.deinit(); } }; diff --git a/zig/tests/cuda/contract_limits.zig b/zig/tests/cuda/contract_limits.zig index 02f9ef56dc2..1f970b29532 100644 --- a/zig/tests/cuda/contract_limits.zig +++ b/zig/tests/cuda/contract_limits.zig @@ -30,7 +30,7 @@ pub fn image(name: []const u8) error{UnknownContractImage}![]const u8 { if (std.mem.eql(u8, name, "h3")) return &h3_image.bytes; inline for (.{ "qwen4_projection", "qwen4_hc", "qwen4_gdn", "qwen4_ple", "gemma4_ple" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); - inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "nvfp4_experts32", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows", "nemotron_ops", "minimax_m3", "minimax_m3_attn", "moss_image", "lfm2_moe", "lfm2", "dots_mocr" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); + inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "nvfp4_experts32", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows", "nemotron_ops", "minimax_m3", "minimax_m3_attn", "experts_bf16", "moss_image", "lfm2_moe", "lfm2", "dots_mocr" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); if (std.mem.eql(u8, name, "contract_controls")) return if (options.controls) &controls.bytes else &.{}; return error.UnknownContractImage; } diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index ebd3f880790..7c113ec4aec 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -215,7 +215,7 @@ const usage = \\ fp8-lane block-FP8 projections against the Python lane matmul's bytes (oracle/fp8_lane.py) \\ fp8-experts grouped block-FP8 experts against recorded fixture bytes \\ nvfp4-experts grouped NVFP4 experts against recorded fixture bytes - \\ bf16-experts grouped bf16 experts against float64: repack, up, downs, rows and forms byte for byte + \\ bf16-experts [contract] grouped bf16 experts against float64 (repack, up, downs, rows, forms); contract: gates 1-4, controls red \\ projection-f32 FP32 projection outputs against FP64, rounding and row/reduction invariance \\ mxfp8 MXFP8 strided projections against FP64 and row-count invariance \\ quant-ref FP8 and NVFP4 projections and experts against float64 references (synthetic) @@ -431,7 +431,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "fp8-experts")) return fp8_experts_tests.experts(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "nvfp4")) return nvfp4_tests.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "nvfp4-experts")) return nvfp4_experts_tests.experts(gpu, try arg(rest, 0)); - if (std.mem.eql(u8, cmd, "bf16-experts")) return @import("bf16_experts_tests.zig").experts(gpu); + if (std.mem.eql(u8, cmd, "bf16-experts")) return if (rest.len == 0) @import("bf16_experts_tests.zig").experts(gpu) else if (std.mem.eql(u8, rest[0], "contract")) @import("contract_experts_bf16.zig").run(gpu) else error.UnknownCommand; if (std.mem.eql(u8, cmd, "projection-f32")) return projection_f32_tests.run(gpu); if (std.mem.eql(u8, cmd, "mxfp8")) return mxfp8_tests.run(gpu); if (std.mem.eql(u8, cmd, "quant-ref")) return quant_ref_tests.run(gpu); From 268e756a82ad7b19039bd455558965ab36c397a1 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Mon, 2 Nov 2026 19:06:01 +0000 Subject: [PATCH 08/78] cuda experts: bf16 run refuses a width or dims off the 32-wide blocks, as pack does, with a test --- zig/src/cuda/experts_bf16.zig | 1 + zig/src/cuda/experts_host_test.zig | 15 +++++++++++++++ 2 files changed, 16 insertions(+) diff --git a/zig/src/cuda/experts_bf16.zig b/zig/src/cuda/experts_bf16.zig index 76ed554c2f7..f585ba0875a 100644 --- a/zig/src/cuda/experts_bf16.zig +++ b/zig/src/cuda/experts_bf16.zig @@ -52,6 +52,7 @@ pub const Bf16Experts = struct { /// One call on a plan of at most `items` items: 16 pairs an item in decode, 64 in a prompt. pub fn run(e: *const Bf16Experts, s: Stream, prompt: bool, epi: Epi, in: Rows, l: Layer, p: Plan, out: u64, items: usize, skip: i32) !void { if (!l.relu2 or l.swiglu or l.limit != 0) return error.OptionNotBuilt; + if (l.width % 32 != 0 or l.dims % 32 != 0) return error.UnexpectedShape; // whole blocks, as pack takes const up = epi == .up; const kg = (if (up) l.dims else l.width) / 32; const nb = (if (up) l.width else l.dims) / cols; diff --git a/zig/src/cuda/experts_host_test.zig b/zig/src/cuda/experts_host_test.zig index 635eea49d2f..65f4a7f7952 100644 --- a/zig/src/cuda/experts_host_test.zig +++ b/zig/src/cuda/experts_host_test.zig @@ -38,3 +38,18 @@ test "the clamp-after order: exl3 takes it on both paths, every other format ref try std.testing.expectError(error.OptionNotBuilt, g.prompt(s, .up, in, l, p, 0, 1, 1, 64, -1)); }; } + +test "bf16 experts refuse dims or widths off the 32-wide blocks before any launch" { + const b: experts.Bf16Experts = undefined; // never reached: the refusal comes first + const g: Grouped = .{ .bf16 = &b }; + const s: Stream = undefined; + const p: Plan = .{ .members = 0, .items = 0, .counts = 0, .rank = 0, .hist = 0 }; + const in: experts.Rows = .{ .x = 0, .stride = 64 }; + for ([_][2]usize{ .{ 48, 64 }, .{ 64, 48 }, .{ 2688, 1880 } }) |wd| { + const l: Layer = .{ .format = .bf16, .up = 0, .down = 0, .width = wd[0], .dims = wd[1], .experts = 2, .relu2 = true }; + for ([_]experts.Epi{ .up, .down_f32, .down_bf16 }) |epi| { + try std.testing.expectError(error.UnexpectedShape, g.decode(s, epi, in, l, p, 0, 1, 1, 16, -1)); + try std.testing.expectError(error.UnexpectedShape, g.prompt(s, epi, in, l, p, 0, 1, 1, 64, -1)); + } + } +} From 1f7137b686327e68111d06d4c8287095c1a5974d Mon Sep 17 00:00:00 2001 From: Ash <71890694+ashhart@users.noreply.github.com> Date: Thu, 12 Nov 2026 23:54:19 +0000 Subject: [PATCH 09/78] tests: re-pin experts.zig and kernels.zig for #1630 #1630 adds the bf16 experts format to Grouped (experts.zig) and registers the experts_bf16 image (kernels.zig); no existing kernel source changed. Both rows are recomputed on main after #3011, which re-pinned experts.zig for #2843. --- zig/tests/fixtures/gradient-serving.tsv | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zig/tests/fixtures/gradient-serving.tsv b/zig/tests/fixtures/gradient-serving.tsv index 60b83c49883..3d3c09e7466 100644 --- a/zig/tests/fixtures/gradient-serving.tsv +++ b/zig/tests/fixtures/gradient-serving.tsv @@ -89,11 +89,11 @@ zig/src/cuda/cublaslt.zig c6e3b8892cad23bd08c7d9843da8117670e3a881e0708a89cdf72f zig/src/cuda/driver.zig f20856bd06afe9365ab9ca1bc39800ccb260ade2704197c033e1344a0c8dd742 raw commits e9bf05ffd, b923e554c, b2c75a98a; series patches 14-runtime-foundation.patch; HELD: must be re-qualified by the LW device gates on the 1.1.1 candidate zig/src/cuda/embed_rows.zig 09fb2e612f43297c19d2cb77d45139dd98212a20075da0a15069a70544f72f6b raw zig/src/cuda/exlr8.zig 1e2164096ae949a8a9a4903df82c7c0c184a6f5e1ba5ba4423038a42b0633fc6 raw -zig/src/cuda/experts.zig 1450c0f7874f0cb5b9ae3a8bebf845f283214863b60cf401a26fdb1898ee7271 raw main 3768bc0168 (#2843: grouped EXL3 experts launched and sized at the caller's tile and pair count; no kernel source change) +zig/src/cuda/experts.zig 8bfc2f15a742cceba87dd7eb8e80f3f1b4e3f7eb323b557eb1bb4bf5a9fed440 raw main 3768bc0168 (#2843: grouped EXL3 experts launched and sized at the caller's tile and pair count; no kernel source change), #1630 (a bf16 format: its Grouped member, decode and prompt arms and prompt tile; other formats unchanged) zig/src/cuda/fp8.zig a9caf312cb451dc95572b493e69ae211dd1e8b5be5b2c48cae3febd5d9327143 raw series patches 65-pr-1024.patch; HELD: must be re-qualified by the LW device gates on the 1.1.1 candidate zig/src/cuda/graph.zig 17639aaef0298c934b3f5ad3511925169c096ae57ede8ebee655faa838844f7b raw zig/src/cuda/grouped.zig 5fa1f39394b701eab59383bc7ed214ea3e17652a1229236df49f5281e6e764b5 raw commits e1be0806d; HELD: must be re-qualified by the LW device gates on the 1.1.1 candidate -zig/src/cuda/kernels.zig db8a02c5b1469baa67e29b915b3e63cee0d0ec8b76009b9472b9307dadf12ebd raw main 42735f94f0 (#2335: the mla image registered), f35900e60c (#2709: the dots_mocr image registered) +zig/src/cuda/kernels.zig 34361df431dc6419474f5460158fd29bcba90f2a3792092b25308b0cb413fd69 raw main 42735f94f0 (#2335: the mla image registered), f35900e60c (#2709: the dots_mocr image registered), #1630 (the experts_bf16 image registered) zig/src/cuda/launch.zig 9ec64848834d82ddcea85a6123e975c95a5ae2bf85c2a71772244794b0bcd705 raw series patches 14-runtime-foundation.patch; HELD: must be re-qualified by the LW device gates on the 1.1.1 candidate zig/src/cuda/memory.zig 7d2a91bfca8a08ae59df16ddf30d22226aba773cc814dd106930840ab344fc52 raw commits 5dcf77d7c, f630f1a90, e9bf05ffd, 881cccaa4, a34cb6ded; series patches 14-runtime-foundation.patch; HELD: must be re-qualified by the LW device gates on the 1.1.1 candidate zig/src/cuda/modelopt.zig 648c9ba58197af0ce905f8792383502fc44a52c45d1413fe9e21d28d04c22253 raw From af2c43769b0fd55ece852b7408a7db4df58326eb Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 12:01:04 +0000 Subject: [PATCH 10/78] changelog: #1630 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index c3007ea6f1f..b66d1e88107 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- CUDA kernels: a bf16 format for grouped experts runs a checkpoint's own bf16 expert weights in mma fragment blocks, with relu² up and fp32 or bf16 down, in decode and prompt forms; no served checkpoint selects it yet (Nemotron 3.5 Transcribe, #1631, will), and every other expert format's dispatch is unchanged (#1630). Thanks to @BobClawblaw. + - DeepSeek-V4.1 on CUDA: `zig build dsv41-assets` now builds `tf-dsv41-assets` at ReleaseSafe whatever `-Doptimize` says, so `aot-needs` plans the same Triton variants about 15 times faster (713 s against 45 s at a 65,536-token context on a Mac; a GB10 run at 524,288 tokens went from about 45 minutes to under 7) (#3007, found on #2270). Thanks to @Defilan. - Docs: the CUDA memory plan says what each setting decides and which file reads it, including that `--context` is capped at the model's own window and `--context 0` takes it (#3028). Thanks to @Bizuayeu. From 3da2b3bd65a942e3bb3dd5c8335e52810c579ec6 Mon Sep 17 00:00:00 2001 From: Ash <71890694+ashhart@users.noreply.github.com> Date: Thu, 12 Nov 2026 23:56:28 +0000 Subject: [PATCH 11/78] glm: the DFlash2 drafter's startup line and docs give its licence as Apache-2.0, as its model card now says --- THIRD_PARTY_NOTICES.md | 3 +-- docs/recipes/glm-5.3-flash-cuda.md | 2 +- docs/recipes/glm-5.3-flash.md | 4 ++-- zig/src/families/glm/cuda_native.zig | 2 +- zig/src/families/glm/dflash.zig | 2 +- zig/src/families/glm/dflash_load.zig | 2 +- 6 files changed, 7 insertions(+), 8 deletions(-) diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index c894ea02bd1..7a72107791f 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -292,8 +292,7 @@ below); `tools/zig/bolt_ggml.py` reads them from that file. The IQ2_XXS and Q8_0 Bolt's GGML checkpoints (`zig/kernels/cuda/ggml.cu`, `zig/src/cuda/ggml.zig`) are TensorFold's own. TensorFold ships no model weights. The `z-lab/Qwen3.8-27B-DFlash2` model card states Apache-2.0. -The optional `incoai/GLM-5.3-Flash-DFlash2` model card states CC BY-NC-ND 4.0, for non-commercial use -without derivatives; only the Python engine reads that drafter. Each checkpoint keeps its own license. +The optional `incoai/GLM-5.3-Flash-DFlash2` model card states Apache-2.0 too. Each checkpoint keeps its own license. ## Structured output: libtfgrammar diff --git a/docs/recipes/glm-5.3-flash-cuda.md b/docs/recipes/glm-5.3-flash-cuda.md index 8d5efda9499..93f671bd8fe 100644 --- a/docs/recipes/glm-5.3-flash-cuda.md +++ b/docs/recipes/glm-5.3-flash-cuda.md @@ -162,7 +162,7 @@ concurrent one its run alone. | DFlash2 | Prose, 16 streams | 526.0 | 583.6 | One-stream runs are `tools/bench_openai.py` at 512 tokens and stream runs `tools/bench_concurrent.py` at 256. DFlash2's -weights are CC BY-NC-ND 4.0. After uses `TF_GLM_DRAFT_CONFIDENCE=0.4` and before keeps the default 0.7. Cold prompts +weights are Apache-2.0. After uses `TF_GLM_DRAFT_CONFIDENCE=0.4` and before keeps the default 0.7. Cold prompts are unchanged at about 4,430 tok/s from 8K to 131K tokens, and a repeated code reply reaches about 850 tok/s with windows of up to 32 rows. Other GPUs and rank counts are not covered here. diff --git a/docs/recipes/glm-5.3-flash.md b/docs/recipes/glm-5.3-flash.md index 538adf74bdc..f2e7a68c365 100644 --- a/docs/recipes/glm-5.3-flash.md +++ b/docs/recipes/glm-5.3-flash.md @@ -42,8 +42,8 @@ tensorfold serve TensorFold/GLM-5.3-Flash-MLX-4bit-MTP --tp 2 --rank 0 --master ``` Rank 1 starts first and rank 0 serves HTTP. Use the same context and drafting settings on both ranks. -The optional `incoai/GLM-5.3-Flash-DFlash2` model has CC BY-NC-ND 4.0 terms; pull it on both ranks only -when those terms fit the intended use. The CLI uses it automatically once it has been pulled. +The optional `incoai/GLM-5.3-Flash-DFlash2` model is Apache-2.0; pull it on both ranks to use it. +The CLI uses it automatically once it has been pulled. Without it, the engine uses MTP drafts; `--drafter none` explicitly selects MTP-only drafting. Give both ranks the same drafter setting. `--no-drafts` disables all drafting for the serial reference. A checkpoint with neither an MTP head nor a supplied DFlash2 model is refused unless drafts are disabled. diff --git a/zig/src/families/glm/cuda_native.zig b/zig/src/families/glm/cuda_native.zig index 25efdd50087..904d7cb72aa 100644 --- a/zig/src/families/glm/cuda_native.zig +++ b/zig/src/families/glm/cuda_native.zig @@ -52,7 +52,7 @@ pub const Options = struct { mtp_confidence: ?f64 = null, // --mtp-confidence: rank 0's lanes stop an MTP chain where its chance falls under it learn_conversations: bool = false, // --learn-conversations: refused with TF_GLM_REQUEST_STARTS=1 (cuda_env) }; -/// --drafter: DFlash2 (incoai/GLM-5.3-Flash-DFlash2, non-commercial weights) drafts instead of the MTP head. +/// --drafter: DFlash2 (incoai/GLM-5.3-Flash-DFlash2, Apache-2.0 weights) drafts instead of the MTP head. pub const takes_drafter = true; /// The smallest room among the other ranks and whose it is. pub const Peer = struct { bytes: u64, rank: usize }; diff --git a/zig/src/families/glm/dflash.zig b/zig/src/families/glm/dflash.zig index 8fb578f3d45..aa039ffc299 100644 --- a/zig/src/families/glm/dflash.zig +++ b/zig/src/families/glm/dflash.zig @@ -1,4 +1,4 @@ -//! DFlash2 drafts (--drafter): incoai/GLM-5.3-Flash-DFlash2 (CC BY-NC-ND 4.0, read unmodified), after z-lab (MIT). +//! DFlash2 drafts (--drafter): incoai/GLM-5.3-Flash-DFlash2 (Apache-2.0, read unmodified), after z-lab (MIT). const std = @import("std"); const cuda = @import("cuda"); const mm = @import("cuda_mm.zig"); diff --git a/zig/src/families/glm/dflash_load.zig b/zig/src/families/glm/dflash_load.zig index 5726ec05fae..a2096579776 100644 --- a/zig/src/families/glm/dflash_load.zig +++ b/zig/src/families/glm/dflash_load.zig @@ -206,7 +206,7 @@ pub fn load(gpa: std.mem.Allocator, io: std.Io, e: anytype, dir: []const u8, ran try @import("dflash_load_owner.zig").acquire(gpa, dr, &tail, Tail, .{ .book = df.VOCAB * df.RANK, .ctx = df.MAX_CTX, .cand = candidates.host, .hidden = df.MAX_M * df.RANK }); dr.log = try @import("sys").environment.read(u64, "TF_GLM_DFLASH2_LOG", dr.log); if (std.c.getenv("TF_GLM_DFLASH2_PROF") != null) dr.prof = io; - std.log.info("dflash2: drafter from {s} (CC BY-NC-ND 4.0: non-commercial use only), linears at {d} bits, {d:.2} GiB on the device", .{ dir, if (bits == 0) 16 else bits, @as(f64, @floatFromInt(l.at)) / (1 << 30) }); + std.log.info("dflash2: drafter from {s} (Apache-2.0), linears at {d} bits, {d:.2} GiB on the device", .{ dir, if (bits == 0) 16 else bits, @as(f64, @floatFromInt(l.at)) / (1 << 30) }); return dr; } From a509a43fbb893095c9efcfe0c1c7daf1462eab46 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 12:01:06 +0000 Subject: [PATCH 12/78] changelog: #3084 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index b66d1e88107..8975e9de0af 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- GLM-5.3-Flash: the DFlash2 drafter's startup line and docs give its licence as Apache-2.0, as its model card now says (#3084, for #1583). + - CUDA kernels: a bf16 format for grouped experts runs a checkpoint's own bf16 expert weights in mma fragment blocks, with relu² up and fp32 or bf16 down, in decode and prompt forms; no served checkpoint selects it yet (Nemotron 3.5 Transcribe, #1631, will), and every other expert format's dispatch is unchanged (#1630). Thanks to @BobClawblaw. - DeepSeek-V4.1 on CUDA: `zig build dsv41-assets` now builds `tf-dsv41-assets` at ReleaseSafe whatever `-Doptimize` says, so `aot-needs` plans the same Triton variants about 15 times faster (713 s against 45 s at a 65,536-token context on a Mac; a GB10 run at 524,288 tokens went from about 45 minutes to under 7) (#3007, found on #2270). Thanks to @Defilan. From b5853bf03f45d33a6ea6ad360c9ee41aef6107ab Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Wed, 11 Nov 2026 04:08:34 +0000 Subject: [PATCH 13/78] metal: dispatch generic prompt lanes as independent row blocks Give tf_lane the reg path's 16-row grid axis and runtime sums stride. Each block keeps the same group products, FMAs and K reduction; cooperative and older-chip encodes retain 16 rows. Qwen27 selects a whole-prompt key for non-reg widths behind the existing Ultra gate. Decode and mixed stacks stay short. Compare each width with main's frozen short kernel, including ragged tails, producer sums, strided outputs and three controls. Pin real quant dispatch counts with main-red host regressions. Refs #2818 --- zig/kernels/metal/core/lane_projection.metal | 17 +- zig/qwen27_quant_metal.zig | 17 +- zig/qwen27_quant_test.zig | 86 +++++ zig/src/core/lane_projection.zig | 23 +- zig/src/families/qwen27/quant_gpu.zig | 8 +- zig/tests/core_lane_long.zig | 153 ++++++++ zig/tests/core_lane_reg.zig | 13 +- zig/tests/fixtures/lane_prompt/README.md | 7 + .../lane_prompt/main_fb8fdf3_lane.metal | 356 ++++++++++++++++++ 9 files changed, 653 insertions(+), 27 deletions(-) create mode 100644 zig/tests/core_lane_long.zig create mode 100644 zig/tests/fixtures/lane_prompt/README.md create mode 100644 zig/tests/fixtures/lane_prompt/main_fb8fdf3_lane.metal diff --git a/zig/kernels/metal/core/lane_projection.metal b/zig/kernels/metal/core/lane_projection.metal index 2fedb0c13c2..05b55977cd1 100644 --- a/zig/kernels/metal/core/lane_projection.metal +++ b/zig/kernels/metal/core/lane_projection.metal @@ -141,12 +141,8 @@ inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, in #endif #if TF_PRECOMPUTE_SUMS -#if TF_BITS == 8 #define TF_SUM_VALUE(row, g) ((row) < M ? XS[(g) * mdims[1] + (row)] : 0.0f) #else -#define TF_SUM_VALUE(row, g) XS[(g) * 16 + (row)] -#endif -#else #define TF_SUM_VALUE(row, g) tf_lane_xsum(X, row, g, M) #endif @@ -178,7 +174,14 @@ inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, in const device float* XS [[buffer(5)]], #endif uint sgi [[simdgroup_index_in_threadgroup]], uint lanei [[thread_index_in_simdgroup]], - uint tgx [[threadgroup_position_in_grid]]) { + uint2 tg [[threadgroup_position_in_grid]]) { + const int row0 = int(tg.y) * 16, rows = mdims[0] - row0; + if (rows <= 0) return; + X += (size_t)row0 * TF_K; + Y += (size_t)row0 * TF_YS; +#if TF_PRECOMPUTE_SUMS + XS += row0; +#endif const int lane = int(lanei), sg = int(sgi); const short qid = short(lane) >> 2; const short fm = (qid & 4) | ((short(lane) >> 1) & 3); @@ -187,8 +190,8 @@ inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, in #else const short fn = ((qid & 2) | (short(lane) & 1)) * 4; #endif - const int M = mdims[0]; - const int tile = tf_tile(int(tgx)); + const int M = min(rows, 16); + const int tile = tf_tile(int(tg.x)); const int n0 = tile * NT; const int g_begin = TF_G0 + (sg * TF_GN) / TF_SK, g_end = TF_G0 + ((sg + 1) * TF_GN) / TF_SK; constexpr auto desc = matmul2d_descriptor(16, NT, GS, false, true, false, matmul2d_descriptor::mode::multiply); diff --git a/zig/qwen27_quant_metal.zig b/zig/qwen27_quant_metal.zig index 02a97ba26d5..9684fb4bba7 100644 --- a/zig/qwen27_quant_metal.zig +++ b/zig/qwen27_quant_metal.zig @@ -32,9 +32,11 @@ pub const Record = struct { released: usize = 0, live_bytes: usize = 0, fail_pipeline: bool = false, + tensor_units: bool = true, current: usize = 0, rows: u32 = 0, - launches: [256]struct { pipeline: usize, rows: u32 } = undefined, + launches: [256]struct { pipeline: usize, rows: u32, grid: Size, stride: i32 } = undefined, + sums_stride: i32 = 0, count: usize = 0, pub fn device(r: *Record) Device { return .{ .r = r }; @@ -65,8 +67,8 @@ pub const Buffer = struct { }; pub const Device = struct { r: *Record, - pub fn tensorUnits(_: Device) bool { - return true; + pub fn tensorUnits(d: Device) bool { + return d.r.tensor_units; } pub fn generation(_: Device) usize { return 17; @@ -111,10 +113,13 @@ pub const ComputeEncoder = struct { } pub fn setBuffer(_: ComputeEncoder, _: Buffer, _: usize, _: usize) void {} pub fn setValue(e: ComputeEncoder, v: anytype, at: usize) void { - if (@TypeOf(v) == [2]i32 and at == 3) e.r.rows = @intCast(v[0]); + if (@TypeOf(v) == [2]i32 and at == 3) { + e.r.rows = @intCast(v[0]); + e.r.sums_stride = v[1]; + } } - pub fn dispatchGroups(e: ComputeEncoder, _: Size, _: Size) void { - e.r.launches[e.r.count] = .{ .pipeline = e.r.current, .rows = e.r.rows }; + pub fn dispatchGroups(e: ComputeEncoder, grid: Size, _: Size) void { + e.r.launches[e.r.count] = .{ .pipeline = e.r.current, .rows = e.r.rows, .grid = grid, .stride = e.r.sums_stride }; e.r.count += 1; } pub fn dispatchThreads(_: ComputeEncoder, _: Size, _: Size) void {} diff --git a/zig/qwen27_quant_test.zig b/zig/qwen27_quant_test.zig index 73d432c8705..1c151d8a718 100644 --- a/zig/qwen27_quant_test.zig +++ b/zig/qwen27_quant_test.zig @@ -132,3 +132,89 @@ test "qwen27 quant: preparation skips projections that cannot take a decode lane try t.expectEqual(@as(u32, 0), k.lanes.count()); try t.expectEqual(@as(usize, 0), r.libraries); } + +const wide_formats = [_]lane_projection.Format{ + .{ .bits = 8, .group = 64 }, .{ .bits = 6, .group = 32 }, .{ .bits = 5, .group = 64 }, + .{ .bits = 3, .group = 32 }, .{ .bits = 2, .group = 64 }, .{ .bits = 8, .group = 128 }, +}; + +test "qwen27 quant: non-reg layouts admit whole prompt chunks without changing source" { + for (wide_formats) |format| { + const long = lane_projection.Layout{ .n = 64, .k = 5120, .format = format, .rows = 2048 }; + try long.validate(); + var short = long; + short.rows = 16; + const text = try lane_projection.source(a, long); + defer a.free(text); + const old = try lane_projection.source(a, short); + defer a.free(old); + try t.expectEqualStrings(old, text); + for ([_]usize{ 0, 15, 17, 2064 }) |rows| { + short.rows = rows; + try t.expectError(error.InvalidSchedule, short.validate()); + } + } +} + +test "qwen27 quant: each non-reg whole prompt is one projection dispatch" { + for (wide_formats) |format| for ([_]u32{ 17, 1065, 2048 }) |rows| { + var r: mtl.Record = .{}; + var k = quant.Kernels{ .allocator = a, .device = r.device(), .long_prompt = true, .prompt = true }; + defer k.deinit(); + var p = matrix(&r, 64); + p.bits = format.bits; + p.group = @intCast(format.group); + const ref = quant.Ref{ .buffer = r.virtual() }; + const producer: ?quant.Sums = if (format.bits == 8 and format.group == 64) .{ .ref = ref, .stride = 2080 } else null; + try k.quant(r.encoder(), p, ref, producer, ref, ref, rows); + try t.expectEqual(@as(usize, 1), r.count); + try t.expectEqual(rows, r.launches[0].rows); + try t.expectEqual(@as(usize, (rows + 15) / 16), r.launches[0].grid.height); + try t.expectEqual(@as(i32, if (producer != null) 2080 else @intCast((rows + 15) / 16 * 16)), r.launches[0].stride); + var entries = k.lanes.valueIterator(); + const entry = entries.next().?; + try t.expectEqual(@as(u32, 2048), entry.max_rows); + if (entry.sums_buffer) |b| try t.expectEqual(@as(usize, 2048 * (5120 / format.group) * 4), b.length()); + // A decode or verify row still uses the short key and its original 16-row steps. + k.prompt = false; + r.count = 0; + try k.quant(r.encoder(), p, ref, producer, ref, ref, 17); + try t.expectEqual(@as(usize, 2), r.count); + try t.expectEqual(@as(u32, 16), r.launches[0].rows); + try t.expectEqual(@as(u32, 1), r.launches[1].rows); + try t.expectEqual(@as(usize, 1), r.launches[0].grid.height); + try t.expectEqual(@as(u32, 2), k.lanes.count()); + k.prompt = true; + k.long_prompt = false; + r.count = 0; + try k.quant(r.encoder(), p, ref, producer, ref, ref, 2048); + try t.expectEqual(@as(usize, 128), r.count); + for (r.launches[0..r.count]) |call| { + try t.expectEqual(@as(u32, 16), call.rows); + try t.expectEqual(@as(usize, 1), call.grid.height); + } + }; +} + +test "qwen27 quant: cooperative and older-chip projections keep their 16-row limits" { + var r: mtl.Record = .{}; + var layout = lane_projection.Layout{ .n = 64, .k = 512, .format = .{ .bits = 4, .group = 64 }, .cooperative = true }; + const coop = try lane_projection.Projection.init(a, r.device(), layout); + defer coop.deinit(); + try t.expectEqual(@as(u32, 16), coop.max_rows); + const ref = lane_projection.Ref{ .buf = r.virtual() }; + try coop.encode(r.encoder(), ref, ref, ref, ref, 16); + try t.expectEqual(@as(usize, 1), r.launches[0].grid.height); + try t.expectError(error.InvalidRows, coop.encode(r.encoder(), ref, ref, ref, ref, 17)); + layout.rows = 2048; + try t.expectError(error.InvalidSchedule, layout.validate()); + layout.cooperative = false; + layout.format = .{ .bits = 6, .group = 32 }; + r.tensor_units = false; + const old = try lane_projection.Projection.init(a, r.device(), layout); + defer old.deinit(); + try t.expectEqual(@as(u32, 16), old.max_rows); + try t.expectError(error.InvalidRows, old.encode(r.encoder(), ref, ref, ref, ref, 17)); + const over_budget = lane_projection.Layout{ .n = 2048, .k = 5120, .format = .{ .bits = 8, .group = 128 }, .sk = 8 }; + try t.expectError(error.ThreadgroupMemoryLimit, over_budget.validate()); +} diff --git a/zig/src/core/lane_projection.zig b/zig/src/core/lane_projection.zig index a12f9d52f2b..a27141818d9 100644 --- a/zig/src/core/lane_projection.zig +++ b/zig/src/core/lane_projection.zig @@ -1,4 +1,4 @@ -//! Affine lane projections for 1-16 bf16 rows on Metal. No family or model dependency. +//! Affine lane projections in independent 16-row blocks on Metal. No family or model dependency. const std = @import("std"); const mtl = @import("metal"); const ks = @import("kernel_sources"); @@ -36,7 +36,7 @@ pub const Layout = struct { cooperative: bool = false, /// 4-bit g64 codes in reg_layout, unpacked into tensor-op registers; every word equals the cooperative kernel's. reg: bool = false, - /// Rows one reg encode may take, `reg_rows` to `max_reg_rows`: it sizes the sums scratch, never the kernel source. + /// Rows one encode may take: it sizes scratch, never the kernel source; cooperative encodes still take 16. rows: usize = reg_rows, /// Elements between output rows (default N): a member of a fused stack writes its columns of the fused rows. y_stride: ?usize = null, @@ -50,7 +50,9 @@ pub const Layout = struct { if (l.n == 0 or l.n % 16 != 0 or l.k == 0 or l.k % l.format.group != 0 or l.n > std.math.maxInt(i32) / 16 or l.k > std.math.maxInt(i32) / 16) return error.InvalidShape; if (l.sk == 0 or l.sk > 32 or (l.pf != 1 and l.pf != 2)) return error.InvalidSchedule; - if (l.rows != reg_rows and (!l.reg or l.rows < reg_rows or l.rows % 16 != 0 or l.rows > max_reg_rows)) return error.InvalidSchedule; + if (l.cooperative) { + if (l.rows != reg_rows) return error.InvalidSchedule; + } else if (l.rows < @as(usize, if (l.reg) reg_rows else 16) or l.rows % 16 != 0 or l.rows > max_reg_rows) return error.InvalidSchedule; const g = l.groupCut(); if (g[1] == 0 or g[0] >= l.k / l.format.group or g[1] > l.k / l.format.group - g[0]) return error.InvalidGroupCut; if (l.groups != null and l.output != .f32) return error.PartialOutputNeedsFloat; @@ -238,12 +240,15 @@ pub const Projection = struct { sums_buffer: ?mtl.Buffer = null, threads_per_slice: usize = 32, reg: bool = false, - /// Rows one encode takes at most: the layout's rows for reg, 16 otherwise. + cooperative: bool = false, + /// Rows one encode takes at most: the layout's rows on tensor units, 16 on cooperative or older-chip paths. max_rows: u32 = reg_rows, pub fn init(a: std.mem.Allocator, device: mtl.Device, l: Layout) !Projection { try l.validate(); - if (!device.tensorUnits() and (l.format.bits != 6 or l.format.group != 32 or l.cooperative or l.reg or l.precompute_sums)) return error.TensorUnitsRequired; + const units = device.tensorUnits(); + if (!units and (l.format.bits != 6 or l.format.group != 32 or l.cooperative or l.reg or l.precompute_sums)) return error.TensorUnitsRequired; + const capacity: usize = if (l.cooperative or !units) 16 else l.rows; const text = try source(a, l); defer a.free(text); const adapted = try frags.source(device, a, text); @@ -258,9 +263,9 @@ pub const Projection = struct { if (l.precompute_sums) { sums_pipe = try mtl.Pipeline.init(device, lib, "tf_lane_sums", false); errdefer sums_pipe.?.deinit(); - sums_buffer = try device.buffer(@as(usize, if (l.reg) l.rows else 16) * (l.k / l.format.group) * 4, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + sums_buffer = try device.buffer(capacity * (l.k / l.format.group) * 4, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); } - return .{ .pipe = p, .external_sums = l.precompute_sums and l.format.bits == 8 and l.format.group == 64 and !l.cooperative, .n = l.n, .k = l.k, .group_count = l.k / l.format.group, .weight_bytes = l.weightWords() * 4, .metadata_bytes = l.metadataElements() * 2, .output_bytes = if (l.output == .f32) 4 else 2, .y_stride = l.y_stride orelse l.n, .tile_count = l.tiles(), .sk = l.sk, .sums_pipe = sums_pipe, .sums_buffer = sums_buffer, .threads_per_slice = if (l.cooperative) 64 else 32, .reg = l.reg, .max_rows = @intCast(if (l.reg) l.rows else 16) }; + return .{ .pipe = p, .external_sums = l.precompute_sums and l.format.bits == 8 and l.format.group == 64 and !l.cooperative, .n = l.n, .k = l.k, .group_count = l.k / l.format.group, .weight_bytes = l.weightWords() * 4, .metadata_bytes = l.metadataElements() * 2, .output_bytes = if (l.output == .f32) 4 else 2, .y_stride = l.y_stride orelse l.n, .tile_count = l.tiles(), .sk = l.sk, .sums_pipe = sums_pipe, .sums_buffer = sums_buffer, .threads_per_slice = if (l.cooperative) 64 else 32, .reg = l.reg, .cooperative = l.cooperative, .max_rows = @intCast(capacity) }; } pub fn deinit(p: Projection) void { @@ -306,7 +311,7 @@ pub const Projection = struct { e.setValue([2]i32{ @intCast(rows), @intCast(stride) }, 3); e.setBuffer(y.buf, y.off, 4); if (given) |g| e.setBuffer(g.ref.buf, g.ref.off, 5) else if (p.sums_buffer) |buffer| e.setBuffer(buffer, 0, 5); - e.dispatchGroups(mtl.Size.of(p.tile_count, if (p.reg) blocks else 1, 1), mtl.Size.of(p.threads_per_slice * p.sk, 1, 1)); + e.dispatchGroups(mtl.Size.of(p.tile_count, if (p.cooperative) 1 else blocks, 1), mtl.Size.of(p.threads_per_slice * p.sk, 1, 1)); if (given == null and p.sums_buffer != null) e.barrier(); } }; @@ -341,7 +346,7 @@ test "a reg layout's rows size its sums scratch, never its kernel source" { } var plain = base; plain.reg = false; - plain.rows = 256; + plain.rows = max_reg_rows + 16; try std.testing.expectError(error.InvalidSchedule, plain.validate()); } diff --git a/zig/src/families/qwen27/quant_gpu.zig b/zig/src/families/qwen27/quant_gpu.zig index 12ae470de2e..be2149254ae 100644 --- a/zig/src/families/qwen27/quant_gpu.zig +++ b/zig/src/families/qwen27/quant_gpu.zig @@ -13,7 +13,7 @@ pub const Sums = struct { ref: Ref, stride: u32 }; pub const Key = struct { kind: Kind, n: u32 = 0, k: u32, slices: u32 = 0, tile: u32 = 0, tmr: u32 = 0, edge: bool = false, skip: u32 = 0, cut: [2]u32 = .{ 0, 0 }, bits: u8 = 4, group: u32 = 64, long: bool = false }; pub const Kind = enum { projection, xsum, norm, input_norm, mlp, ane_input, mlp_ane, add_ane }; -/// Whether prompt chunks take one reg dispatch a projection (`Kernels.long_prompt`): Ultras with tensor units only. +/// Whether prompt chunks take one dispatch a projection (`Kernels.long_prompt`): Ultras with tensor units only. pub fn longPrompt(name: []const u8, tensor_units: bool) bool { return tensor_units and std.mem.endsWith(u8, name, " Ultra"); } @@ -56,7 +56,7 @@ pub const Kernels = struct { pair_twins: std.AutoHashMapUnmanaged([2]usize, bool) = .empty, /// Set around prompt chunks: prompt rows take projection.promptSplit, decoded rows the decode split. prompt: bool = false, - /// Set on `longPrompt` chips: prompt chunks past 128 rows take one reg dispatch a projection, on their own key. + /// Set on `longPrompt` chips: chunks past a short entry's capacity take one dispatch on a separate long key. long_prompt: bool = false, lanes: std.AutoHashMapUnmanaged(Key, lane.Projection) = .empty, cache: std.AutoHashMapUnmanaged(Key, mtl.Pipeline) = .empty, @@ -174,7 +174,7 @@ pub const Kernels = struct { .output = if (cut[1] > 0) .f32 else .bf16, .precompute_sums = reg or (p.bits == 8 and p.group == 64), .reg = reg, - .rows = if (long) lane.max_reg_rows else lane.reg_rows, + .rows = if (long) lane.max_reg_rows else if (reg) lane.reg_rows else 16, }); k.joined(entry.value_ptr); } @@ -246,7 +246,7 @@ pub const Kernels = struct { } const reg = p.bits == 4 and p.group == 64; if (!reg and (skip > 0 or cut[1] > 0)) return error.BadCoreProjectionRows; - const long = reg and longKey(k.long_prompt, k.prompt, rows); + const long = if (reg) longKey(k.long_prompt, k.prompt, rows) else k.long_prompt and k.prompt and rows > 16; const sums: ?Sums = if (reg or (p.bits == 8 and p.group == 64)) xs else null; const entry = try k.laneEntry(p, sk, skip, cut, long); const step = entry.value_ptr.max_rows; diff --git a/zig/tests/core_lane_long.zig b/zig/tests/core_lane_long.zig new file mode 100644 index 00000000000..502aa57f3da --- /dev/null +++ b/zig/tests/core_lane_long.zig @@ -0,0 +1,153 @@ +//! Whole prompt lanes at each width against main's frozen 16-row kernel, including strided oQ member formats. +const std = @import("std"); +const mtl = @import("metal"); +const lane = @import("core_lane"); +const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; +const original = @embedFile("fixtures/lane_prompt/main_fb8fdf3_lane.metal"); + +fn bf16(v: f32) u16 { + const b: u32 = @bitCast(v); + return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +fn pipe(device: mtl.Device, text: []const u8) !mtl.Pipeline { + const lib = try mtl.Library.fromSource(device, text, mtl.CompileOptions.mlx()); + defer lib.deinit(); + return mtl.Pipeline.init(device, lib, "tf_lane", false); +} + +fn frozen(a: std.mem.Allocator, d: mtl.Device, l: lane.Layout) !mtl.Pipeline { + const current = try lane.source(a, l); + defer a.free(current); + const end = std.mem.indexOf(u8, current, "#include ") orelse return error.SourceHeader; + const old = try std.mem.concat(a, u8, &.{ current[0..end], original }); + defer a.free(old); + return pipe(d, old); +} + +fn chunks(q: mtl.Queue, p: lane.Projection, x: mtl.Buffer, w: mtl.Buffer, sb: mtl.Buffer, xs: ?lane.Sums, y: mtl.Buffer, rows: u32) !void { + const cb = q.commandBuffer(); + const e = cb.compute(.serial); + var first: u32 = 0; + while (first < rows) : (first += p.max_rows) { + const xr = lane.Ref{ .buf = x, .off = @as(usize, first) * p.k * 2 }; + const yr = lane.Ref{ .buf = y, .off = @as(usize, first) * p.y_stride * 2 }; + const n = @min(rows - first, p.max_rows); + if (xs) |s| { + try p.encodeSums(e, xr, .{ .buf = w }, .{ .buf = sb }, .{ .ref = .{ .buf = s.ref.buf, .off = first * 4 }, .stride = s.stride }, yr, n); + } else try p.encode(e, xr, .{ .buf = w }, .{ .buf = sb }, yr, n); + } + e.end(); + cb.commit(); + cb.wait(); + if (cb.failure() != null) return error.GpuFailure; +} + +fn clear(y: mtl.Buffer, rows: usize, n: usize, stride: usize, candidate: bool) void { + @memset(y.contents()[0..y.length()], 0x5a); + if (candidate) for (0..rows) |m| @memset(y.contents()[m * stride * 2 ..][0 .. n * 2], 0xa5); +} + +fn same(a: mtl.Buffer, b: mtl.Buffer) bool { + return std.mem.eql(u8, a.contents()[0..a.length()], b.contents()[0..b.length()]); +} + +fn controls(a: std.mem.Allocator, d: mtl.Device, q: mtl.Queue, l: lane.Layout, p: lane.Projection, x: mtl.Buffer, w: mtl.Buffer, sb: mtl.Buffer, sums: lane.Sums, want: mtl.Buffer, got: mtl.Buffer) !void { + const text = try lane.source(a, l); + defer a.free(text); + const from = [_][]const u8{ "int(tg.y) * 16", "XS += row0;", "XS[(g) * mdims[1] + (row)]" }; + const to = [_][]const u8{ "0", "", "XS[(g) * 16 + (row)]" }; + for (from, to, 0..) |needle, replacement, i| { + const bad = try std.mem.replaceOwned(u8, a, text, needle, replacement); + defer a.free(bad); + if (std.mem.eql(u8, bad, text)) return error.ControlNotPlanted; + const bad_pipe = try pipe(d, bad); + defer bad_pipe.deinit(); + var broken = p; + broken.pipe = bad_pipe; + clear(got, 17, l.n, l.y_stride.?, true); + try chunks(q, broken, x, w, sb, sums, got, 17); + if (same(want, got)) return error.ControlMissed; + std.debug.print("LANE LONG CONTROL {d} caught\n", .{i}); + } +} + +pub fn check(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, rng: std.Random) !usize { + var count: usize = 0; + const formats = [_]lane.Format{ + .{ .bits = 8, .group = 64 }, .{ .bits = 6, .group = 32 }, .{ .bits = 5, .group = 64 }, + .{ .bits = 3, .group = 32 }, .{ .bits = 2, .group = 64 }, .{ .bits = 8, .group = 128 }, + .{ .bits = 5, .group = 64 }, // the generic precomputed-sums layout, beside the normal inline-sums one + }; + for (formats, 0..) |format, fi| { + // Non-reg entries retain splitFor's loader split; only reg entries use promptSplit's smaller values. + for ([_][3]usize{ .{ 34816, 5120, 2 }, .{ 5120, 17408, 8 }, .{ 10240, 5120, 4 }, .{ 12288, 5120, 4 }, .{ 6240, 5120, 8 }, .{ 5120, 6144, 8 }, .{ 2048, 5120, 8 }, .{ 48, 512, 2 } }, 0..) |shape, si| { + const n = shape[0]; + const k = shape[1]; + const stride = n + @as(usize, if (si == 7) 32 else 0); + const most = lane.max_reg_rows; + const sum_stride = most + 16; + const ng = k / format.group; + // Group128 SK8 exceeds the unchanged admission bound; compare its supported SK4 variant on both sides. + const sk = if (format.group == 128 and shape[2] == 8) 4 else shape[2]; + const layout = lane.Layout{ .n = n, .k = k, .format = format, .sk = sk, .rows = most, .precompute_sums = fi == 0 or fi == 6, .y_stride = stride }; + const p_long = try lane.Projection.init(a, device, layout); + defer p_long.deinit(); + var short = layout; + short.rows = 16; + var p_short = try lane.Projection.init(a, device, short); + defer p_short.deinit(); + const old_pipe = try frozen(a, device, short); + p_short.pipe.deinit(); + p_short.pipe = old_pipe; + const raw = try a.alloc(u32, layout.sourceWords()); + defer a.free(raw); + for (raw) |*v| v.* = rng.int(u32); + const scales = try a.alloc(u16, n * ng); + defer a.free(scales); + const biases = try a.alloc(u16, n * ng); + defer a.free(biases); + for (scales, biases) |*s, *b| { + s.* = bf16(0.002 + rng.float(f32) * 0.02); + b.* = bf16(-0.01 - rng.float(f32) * 0.15); + } + const w = try device.buffer(layout.weightWords() * 4, opts); + defer w.deinit(); + const sb = try device.buffer(layout.metadataElements() * 2, opts); + defer sb.deinit(); + try lane.pack(layout, raw, scales, biases, w.slice(u32, layout.weightWords()), sb.slice(u16, layout.metadataElements())); + const x = try device.buffer(most * k * 2, opts); + defer x.deinit(); + const input = x.slice(u16, most * k); + for (input) |*v| v.* = bf16((rng.float(f32) - 0.5) * 4); + const given = try device.buffer(ng * sum_stride * 4, opts); + defer given.deinit(); + const sums = given.slice(f32, ng * sum_stride); + @memset(sums, -999); + for (0..ng) |g| for (0..most) |m| { + var acc: f32 = 0; + for (0..format.group) |j| acc += @as(f32, @bitCast(@as(u32, input[m * k + g * format.group + j]) << 16)); + sums[g * sum_stride + m] = acc; + }; + const want = try device.buffer(most * stride * 2, opts); + defer want.deinit(); + const got = try device.buffer(most * stride * 2, opts); + defer got.deinit(); + for ([_]u32{ 17, 1065, most }) |rows| for ([_]bool{ false, true }) |external| { + if (external and fi != 0) continue; + const xs: ?lane.Sums = if (external) .{ .ref = .{ .buf = given }, .stride = sum_stride } else null; + clear(want, rows, n, stride, false); + clear(got, rows, n, stride, true); + try chunks(queue, p_short, x, w, sb, xs, want, rows); + try chunks(queue, p_long, x, w, sb, xs, got, rows); + if (!same(want, got)) { + std.debug.print("LANE LONG DIFFER bits={d} group={d} n={d} k={d} sk={d} rows={d} external={}\n", .{ format.bits, format.group, n, k, sk, rows, external }); + return error.LongLaneBytesDiffer; + } + if (fi == 0 and si == 0 and rows == 17 and external) try controls(a, device, queue, layout, p_long, x, w, sb, xs.?, want, got); + count += 1; + }; + } + } + return count; +} diff --git a/zig/tests/core_lane_reg.zig b/zig/tests/core_lane_reg.zig index eb33896366f..b00151b9f71 100644 --- a/zig/tests/core_lane_reg.zig +++ b/zig/tests/core_lane_reg.zig @@ -1,4 +1,4 @@ -//! The reg kernel, every output byte: 1-16 rows against the cooperative one, 129-2048 rows against 128-row dispatches. +//! Each lane width's output bytes: reg against cooperative, whole prompts against short dispatches. const std = @import("std"); const mtl = @import("metal"); const lane = @import("core_lane"); @@ -117,6 +117,15 @@ pub fn main(init: std.process.Init) !void { if (!device.tensorUnits()) return error.TensorUnitsRequired; const queue = try device.queue(); defer queue.deinit(); + const args = try init.minimal.args.toSlice(a); + const only_long = args.len == 2 and std.mem.eql(u8, args[1], "--lane-long-only"); + if (args.len > 1 and !only_long) return error.BadArguments; + if (only_long) { + var random = std.Random.DefaultPrng.init(0x7e9); + const checked = try @import("core_lane_long.zig").check(a, device, queue, random.random()); + std.debug.print("LANE LONG CHECK PASS cases={d}: every byte equals main's frozen 16-row kernel\n", .{checked}); + return; + } const order = try lane.RegOrder.init(device); defer order.deinit(); var prng = std.Random.DefaultPrng.init(0x7e9); @@ -187,4 +196,6 @@ pub fn main(init: std.process.Init) !void { std.debug.print("REG CHECK PASS cases={d}: every output byte equals the cooperative kernel's\n", .{cases}); const long = try longCheck(a, device, queue, order, rng); std.debug.print("LONG CHECK PASS cases={d}: one dispatch of up to {d} rows equals {d}-row dispatches byte for byte\n", .{ long, lane.max_reg_rows, lane.reg_rows }); + const wide = try @import("core_lane_long.zig").check(a, device, queue, rng); + std.debug.print("LANE LONG CHECK PASS cases={d}: every byte equals main's frozen 16-row kernel\n", .{wide}); } diff --git a/zig/tests/fixtures/lane_prompt/README.md b/zig/tests/fixtures/lane_prompt/README.md new file mode 100644 index 00000000000..0ce2210fe66 --- /dev/null +++ b/zig/tests/fixtures/lane_prompt/README.md @@ -0,0 +1,7 @@ +The frozen Metal file is TensorFold's own lane kernel from main +fb8fdf3d8501255058a3221452164c56cb9a57ca. It is copied without modification from +zig/kernels/metal/core/lane_projection.metal and retains the repository's Apache-2.0 license. + +core_lane_long.zig compiles this kernel with each tested shape's defines and uses +16-row dispatches as the byte reference for the current whole-prompt kernel. +The fixture must stay frozen when the production kernel changes. diff --git a/zig/tests/fixtures/lane_prompt/main_fb8fdf3_lane.metal b/zig/tests/fixtures/lane_prompt/main_fb8fdf3_lane.metal new file mode 100644 index 00000000000..2fedb0c13c2 --- /dev/null +++ b/zig/tests/fixtures/lane_prompt/main_fb8fdf3_lane.metal @@ -0,0 +1,356 @@ +#include +#include +using namespace metal; +using namespace mpp::tensor_ops; + +// 1-16 bf16 rows share each group's weight read; every row's sums and split-K order hold at every width. + +#ifndef TF_SIMD_FRAGS +#define TF_SIMD_FRAGS 0 +#endif + +#ifndef TF_PRECOMPUTE_SUMS +#define TF_PRECOMPUTE_SUMS 0 +#endif +#ifndef TF_YS +#define TF_YS TF_N +#endif +// Output rows lie TF_YS elements apart, so a member of a fused stack writes its columns of the fused rows. +#define TF_ODD_WORDS (TF_GROUP * TF_BITS / 32 % 2) + +constant constexpr int NT = 32, GS = TF_GROUP, KG = TF_K / GS, NF = NT / 16; +constant constexpr int WP = GS * TF_BITS / 32; + +#if TF_ODD_WORDS +struct tf_raw { uint w[WP]; uint4 s[NF]; }; // 3- and 5-bit groups of 32 are an odd word count, loaded a word at a time +#else +struct tf_raw { uint2 w[WP / 2]; uint4 s[NF]; }; // a lane's packed code words and its scale/bias quads for one group +#endif + +inline tf_raw tf_lane_fetch(const device uint* W, const device uint4* SB, int tile, int g, int g_end, int n0, int lane, + short fn) { + tf_raw raw; + const bool ok = g < g_end; + const device uint2* src = (const device uint2*)(W + (((size_t)tile * KG + (ok ? g : 0)) * NT + lane) * WP); +#if TF_ODD_WORDS + for (int i = 0; i < WP; i++) raw.w[i] = ok && n0 + lane < TF_N ? ((const device uint*)src)[i] : 0u; +#elif TF_BITS != 4 && TF_BITS != 8 + for (int i = 0; i < WP / 2; i++) raw.w[i] = ok && n0 + lane < TF_N ? src[i] : uint2(0); +#else + for (int i = 0; i < WP / 2; i++) raw.w[i] = uint2(0); +#endif +#if TF_SIMD_FRAGS + const device uint2* sb2 = (const device uint2*)SB; + for (int f = 0; f < NF; f++) { + const int col = n0 + f * 16 + fn; + raw.s[f] = uint4(ok && col < TF_N ? sb2[((size_t)g * TF_N + col) / 2] : uint2(0), + ok && col + 8 < TF_N ? sb2[((size_t)g * TF_N + col + 8) / 2] : uint2(0)); + } +#else + // Four lanes load the metadata quads and broadcast them; all 32 lanes fetch, even for inactive rows. + const ushort leader = ushort((fn & 8) | ((fn >> 2) & 1)); + for (int f = 0; f < NF; f++) { + const uint4 sb = lane == int(leader) && ok && n0 + f * 16 + fn < TF_N + ? SB[((size_t)(ok ? g : 0) * TF_N + n0 + f * 16 + fn) / 4] : uint4(0); + raw.s[f] = simd_shuffle(sb, leader); + } +#endif + return raw; +} + +inline float tf_lane_xsum(const device bfloat* X, int m, int g, int M) { + float acc = 0.0f; + if (m < M) for (int i = 0; i < GS; i++) acc += float(X[(size_t)m * TF_K + g * GS + i]); + return acc; +} + +// Codes: 4/8-bit device views; 2/3-bit widened to nibbles, 5/6-bit to bytes. +#if TF_BITS == 4 +#define TF_STAGE_CODES(raw) +#define TF_GROUP_PRODUCT(g) \ + tensor, tensor_inline> b( \ + (device uchar*)W + ((size_t)tile * KG + (g)) * (NT * GS / 2), dextents(GS, NT)); \ + auto a = tA.slice((g) * GS, 0); \ + auto P = op.template get_destination_cooperative_tensor(); \ + op.run(a, b, P); +#elif TF_BITS == 8 +#define TF_STAGE_CODES(raw) +#define TF_GROUP_PRODUCT(g) \ + tensor, tensor_inline> b( \ + (device uint8_t*)W + ((size_t)tile * KG + (g)) * (NT * GS), dextents(GS, NT)); \ + auto a = tA.slice((g) * GS, 0); \ + auto P = op.template get_destination_cooperative_tensor(); \ + op.run(a, b, P); +#else +#if TF_BITS < 4 +inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, int lane) { + uint words[WP + 1]; +#if TF_ODD_WORDS + for (int i = 0; i < WP; i++) words[i] = raw.w[i]; +#else + for (int i = 0; i < WP / 2; i++) { words[2*i] = raw.w[i].x; words[2*i+1] = raw.w[i].y; } +#endif + words[WP] = 0u; +#if TF_BITS == 3 + for (int c = 0; c < GS / 8; c++) { + const int bit = 24 * c, i = bit >> 5, sh = bit & 31; + uint pack = words[i] >> sh; + if (sh > 8) pack |= words[i + 1] << (32 - sh); + uint nib = 0u; + for (int j = 0; j < 8; j++) nib |= ((pack >> (3 * j)) & 7u) << (4 * j); + stage[lane * (GS / 8) + c] = nib; + } +#else + for (int c = 0; c < GS / 8; c++) { // word c/2's half c%2: 8 codes of 2 bits to 8 nibbles + uint v = (words[c >> 1] >> (16 * (c & 1))) & 0xFFFFu; + v = (v | (v << 8)) & 0x00FF00FFu; + v = (v | (v << 4)) & 0x0F0F0F0Fu; + v = (v | (v << 2)) & 0x33333333u; + stage[lane * (GS / 8) + c] = v; + } +#endif +} +#else +inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, int lane) { + uint words[WP + 1]; +#if TF_ODD_WORDS + for (int i = 0; i < WP; i++) words[i] = raw.w[i]; +#else + for (int i = 0; i < WP / 2; i++) { words[2*i] = raw.w[i].x; words[2*i+1] = raw.w[i].y; } +#endif + words[WP] = 0u; + for (int c = 0; c < GS / 4; c++) { + uint word = 0u; + for (int j = 0; j < 4; j++) { + const int bit = (4*c + j) * TF_BITS, i = bit >> 5, sh = bit & 31; + uint code = words[i] >> sh; + if (sh + TF_BITS > 32) code |= words[i+1] << (32-sh); + word |= (code & ((1u << TF_BITS)-1u)) << (8*j); + } + stage[lane * (GS / 4) + c] = word; + } +} +#endif +#define TF_STAGE_CODES(raw) tf_stage_codes(stage, raw, lane); +#define TF_GROUP_PRODUCT(g) \ + simdgroup_barrier(mem_flags::mem_threadgroup); \ + auto a = tA.slice((g) * GS, 0); \ + auto P = op.template get_destination_cooperative_tensor(); \ + op.run(a, b, P); \ + simdgroup_barrier(mem_flags::mem_threadgroup); +#endif + +#if TF_PRECOMPUTE_SUMS +#if TF_BITS == 8 +#define TF_SUM_VALUE(row, g) ((row) < M ? XS[(g) * mdims[1] + (row)] : 0.0f) +#else +#define TF_SUM_VALUE(row, g) XS[(g) * 16 + (row)] +#endif +#else +#define TF_SUM_VALUE(row, g) tf_lane_xsum(X, row, g, M) +#endif + +// Each group loads packed codes or stages widened codes, performs the tensor product, then applies the scale/bias FMAs. +#define TF_LANE_GROUP(g, raw) \ + { \ + TF_STAGE_CODES(raw) \ + float s[NF][4], bb[NF][4]; \ + for (int f = 0; f < NF; f++) { \ + const vec v = as_type>(raw.s[f]); \ + for (int j = 0; j < 4; j++) { s[f][j] = float(v[2 * j]); bb[f][j] = float(v[2 * j + 1]); } \ + } \ + raw = tf_lane_fetch(W, sbv, tile, (g) + TF_PF, g_end, n0, lane, fn); \ + TF_GROUP_PRODUCT(g) \ + const float xs0 = TF_SUM_VALUE(fm, g); \ + const float xs1 = TF_SUM_VALUE(fm + 8, g); \ + for (int f = 0; f < NF; f++) \ + for (int r = 0; r < 2; r++) \ + for (int j = 0; j < 4; j++) { \ + const int i = f * 8 + r * 4 + j; \ + C[i] = fma(s[f][j], P[TF_SIMD_FRAGS ? r * 8 + f * 4 + j : i], \ + fma(bb[f][j], r ? xs1 : xs0, C[i])); \ + } \ + } + +[[kernel]] void tf_lane(const device bfloat* X [[buffer(0)]], const device uint* W [[buffer(1)]], + const device bfloat* SBt [[buffer(2)]], const device int* mdims [[buffer(3)]], device TF_OUT* Y [[buffer(4)]], +#if TF_PRECOMPUTE_SUMS + const device float* XS [[buffer(5)]], +#endif + uint sgi [[simdgroup_index_in_threadgroup]], uint lanei [[thread_index_in_simdgroup]], + uint tgx [[threadgroup_position_in_grid]]) { + const int lane = int(lanei), sg = int(sgi); + const short qid = short(lane) >> 2; + const short fm = (qid & 4) | ((short(lane) >> 1) & 3); +#if TF_SIMD_FRAGS + const short fn = ((short(lane) & 8) >> 1) | ((short(lane) & 1) << 1); +#else + const short fn = ((qid & 2) | (short(lane) & 1)) * 4; +#endif + const int M = mdims[0]; + const int tile = tf_tile(int(tgx)); + const int n0 = tile * NT; + const int g_begin = TF_G0 + (sg * TF_GN) / TF_SK, g_end = TF_G0 + ((sg + 1) * TF_GN) / TF_SK; + constexpr auto desc = matmul2d_descriptor(16, NT, GS, false, true, false, matmul2d_descriptor::mode::multiply); + matmul2d op; + tensor, tensor_inline> tA((device bfloat*)X, dextents(TF_K, M)); +#if TF_BITS < 4 + threadgroup uint stage_all[TF_SK * NT * (GS / 8)]; + threadgroup uint* stage = stage_all + sg * NT * (GS / 8); + tensor, tensor_inline> b((threadgroup uchar*)stage, dextents(GS, NT)); +#elif TF_BITS != 4 && TF_BITS != 8 + threadgroup uint stage_all[TF_SK * NT * (GS / 4)]; + threadgroup uint* stage = stage_all + sg * NT * (GS / 4); + tensor, tensor_inline> b((threadgroup uint8_t*)stage, dextents(GS, NT)); +#endif + const device uint4* sbv = (const device uint4*)SBt; + float C[NF * 8]; + for (int i = 0; i < NF * 8; i++) C[i] = 0.0f; + tf_raw r0 = tf_lane_fetch(W, sbv, tile, g_begin, g_end, n0, lane, fn); +#if TF_PF == 2 + tf_raw r1 = tf_lane_fetch(W, sbv, tile, g_begin + 1, g_end, n0, lane, fn); + for (int g = g_begin; g < g_end; g += 2) { + TF_LANE_GROUP(g, r0) + if (g + 1 < g_end) TF_LANE_GROUP(g + 1, r1) + } +#else + for (int g = g_begin; g < g_end; g++) TF_LANE_GROUP(g, r0) +#endif + threadgroup float part[(TF_SK > 1 ? TF_SK - 1 : 1) * NF * 8 * 32]; + if (TF_SK > 1) { + if (sg > 0) for (int i = 0; i < NF * 8; i++) part[((sg - 1) * NF * 8 + i) * 32 + lane] = C[i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) + for (int s2 = 1; s2 < TF_SK; s2++) for (int i = 0; i < NF * 8; i++) C[i] += part[((s2 - 1) * NF * 8 + i) * 32 + lane]; + } + if (sg == 0) + for (int f = 0; f < NF; f++) + for (int r = 0; r < 2; r++) { + const int m = fm + 8 * r; + const int nn = n0 + f * 16 + fn; + if (m < M && nn < TF_N) + for (int j = 0; j < 4; j++) { + const int col = nn + (TF_SIMD_FRAGS ? (j & 1) + 8 * (j >> 1) : j); + if (col < TF_N) Y[(size_t)m * TF_YS + col] = static_cast(C[f * 8 + r * 4 + j]); + } + } +} + +#if TF_PRECOMPUTE_SUMS +kernel void tf_lane_sums(const device bfloat* X [[buffer(0)]],constant int& rows [[buffer(1)]],device float* XS [[buffer(2)]],constant int& stride [[buffer(3)]],uint index [[thread_position_in_grid]]) { + const int row=int(index)%stride,group=int(index)/stride; + if(group>=KG)return; + XS[index]=tf_lane_xsum(X,row,group,rows); +} +#endif + +#if TF_BITS == 4 +kernel void tf_lane_pair(const device bfloat* X [[buffer(0)]],const device uint* W [[buffer(1)]],const device bfloat* SB [[buffer(2)]],const device int* mdims [[buffer(3)]],device TF_OUT* Y [[buffer(4)]], +#if TF_PRECOMPUTE_SUMS +const device float* XS [[buffer(5)]], +#endif +uint group_id [[simdgroup_index_in_threadgroup]],uint lane [[thread_index_in_simdgroup]],uint thread_id [[thread_position_in_threadgroup]],uint tile_id [[threadgroup_position_in_grid]]) { + const uint sg=group_id/2,tip=thread_id-sg*64; + const int M=mdims[0],tile=tf_tile(int(tile_id)),n0=tile*NT; + const int first=TF_G0+(int(sg)*TF_GN)/TF_SK,last=TF_G0+((int(sg)+1)*TF_GN)/TF_SK; + constexpr auto descriptor=matmul2d_descriptor(16,NT,GS,false,true,false,matmul2d_descriptor::mode::multiply); + matmul2d> op; + tensor,tensor_inline> a((device bfloat*)X,dextents(TF_K,M)); + tensor,tensor_inline> b0((device uchar*)W,dextents(GS,NT)); + auto a0=a.slice(0,0); + auto p=op.template get_destination_cooperative_tensor(); + short column[8],row[8];float value[8]; + for(int i=0;i<8;i++){auto index=p.get_multidimensional_index(i);column[i]=index[0];row[i]=index[1];value[i]=0.0f;} + for(int group=first;group,tensor_inline> b((device uchar*)W+((size_t)tile*KG+group)*(NT*GS/2),dextents(GS,NT)); + auto x=a.slice(group*GS,0);op.run(x,b,p); + for(int i=0;i<8;i++) { + const int n=n0+column[i]; + const float scale=float(SB[(size_t(group)*TF_N+n)*2]),bias=float(SB[(size_t(group)*TF_N+n)*2+1]); +#if TF_PRECOMPUTE_SUMS + const float sum=XS[group*16+row[i]]; +#else + const float sum=tf_lane_xsum(X,row[i],group,M); +#endif + value[i]=fma(scale,p[i],fma(bias,sum,value[i])); + } + } + threadgroup float partial[(TF_SK>1?TF_SK-1:1)*8*64]; + if(TF_SK>1) { + if(sg>0) for(int i=0;i<8;i++) partial[((sg-1)*8+i)*64+tip]=value[i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + if(sg==0) for(int part=1;part> 3) & 1), mb = ((lane >> 1) & 3) + 4 * ((lane >> 4) & 1); + const int first = TF_G0 + (ks * TF_GN) / TF_SK, last = TF_G0 + ((ks + 1) * TF_GN) / TF_SK; + constexpr auto desc = matmul2d_descriptor(16, NT, GS, false, true, false, matmul2d_descriptor::mode::multiply); + matmul2d op; + tensor, tensor_inline> a((device bfloat*)X, dextents(TF_K, M)); + auto a0 = a.slice(0, 0); + auto cb = op.template get_right_input_cooperative_tensor(); + auto P = op.template get_destination_cooperative_tensor, metal::remove_addrspace_t, float>(); + thread half2* b2 = (thread half2*)&cb[0]; + const device uint4* wl = W + (size_t)tile * KG * 64 + lane; + const device uint4* sbv = (const device uint4*)SB; + const half2 bias = half2(1024.0h); + float C[16]; + for (int i = 0; i < 16; i++) C[i] = 0.0f; + tf_reg_load next = {}; + if (first < last) next = tf_reg_fetch(wl, sbv, XS, first, n0 + nb, mb, stride); + for (int g = first; g < last; g++) { + const tf_reg_load cur = next; + for (int w = 0; w < 8; w++) { + // 1024 + code as fp16 bits, codes (0,4) (1,5) (2,6) (3,7) of a word; subtracting 1024 is exact + const uint word = w < 4 ? cur.w0[w] : cur.w1[w - 4], lo = word & 0x0F0F0F0Fu, hi = (word >> 4) & 0x0F0F0F0Fu; + const int e = 16 * (w & 1) + 2 * (w >> 1); + b2[e] = as_type((lo & 0x00FF00FFu) | 0x64006400u) - bias; + b2[e + 1] = as_type((hi & 0x00FF00FFu) | 0x64006400u) - bias; + b2[e + 8] = as_type(((lo >> 8) & 0x00FF00FFu) | 0x64006400u) - bias; + b2[e + 9] = as_type(((hi >> 8) & 0x00FF00FFu) | 0x64006400u) - bias; + } + if (g + 1 < last) next = tf_reg_fetch(wl, sbv, XS, g + 1, n0 + nb, mb, stride); + auto x = a.slice(g * GS, 0); + op.run(x, cb, P); + for (int j = 0; j < 4; j++) { + const float s0 = as_type(cur.q0[j] << 16), b0 = as_type(cur.q0[j] & 0xFFFF0000u); + const float s1 = as_type(cur.q1[j] << 16), b1 = as_type(cur.q1[j] & 0xFFFF0000u); + C[j] = fma(s0, P[j], fma(b0, cur.x0, C[j])); + C[4 + j] = fma(s0, P[4 + j], fma(b0, cur.x1, C[4 + j])); + C[8 + j] = fma(s1, P[8 + j], fma(b1, cur.x0, C[8 + j])); + C[12 + j] = fma(s1, P[12 + j], fma(b1, cur.x1, C[12 + j])); + } + } + threadgroup float part[(TF_SK > 1 ? TF_SK - 1 : 1) * 16 * 32]; + if (TF_SK > 1) { + if (ks > 0) for (int i = 0; i < 16; i++) part[((ks - 1) * 16 + i) * 32 + lane] = C[i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (ks == 0) for (int s2 = 1; s2 < TF_SK; s2++) for (int i = 0; i < 16; i++) C[i] += part[((s2 - 1) * 16 + i) * 32 + lane]; + } + if (ks == 0) + for (int i = 0; i < 16; i++) { + const int m = mb + 8 * ((i >> 2) & 1), n = n0 + nb + 16 * (i >> 3) + (i & 3); + if (m < M && n < TF_N) Y[(size_t)m * TF_YS + n] = static_cast(C[i]); + } +} +#endif From 461437572a33b852d8893bee84edc8e457f66982 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Wed, 11 Nov 2026 04:08:34 +0000 Subject: [PATCH 14/78] metal: record measured offline limits for standalone byte-lane sums The full lane contract compiles all four tf_lane_sums sources for M1 and M2 at a 1024-thread pipeline limit. Record those measured values and keep the actual 256-thread launch. Tensor-op sources remain explicitly unmeasured on these chips. Pin the distinction in a host regression; no shader changes. Refs #2818 --- zig/src/core/contract/lane_q8_records.zig | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/zig/src/core/contract/lane_q8_records.zig b/zig/src/core/contract/lane_q8_records.zig index 16bdd8bd2c1..0e734861227 100644 --- a/zig/src/core/contract/lane_q8_records.zig +++ b/zig/src/core/contract/lane_q8_records.zig @@ -1,4 +1,4 @@ -//! M5 byte-lane sources retain their launch envelopes before and after removing code staging. +//! Byte-lane launch envelopes; the standalone sums compiled for M1/M2 even when tensor ops cannot. const M = @import("limits.zig").MetalLaunch; pub const Bad = enum { none, group, stride }; pub const Source = struct { name: []const u8, n: usize, k: usize, sk: usize, group: usize = 64, sums: bool = false, pf: usize = 1, f32_out: bool = false, bad: Bad = .none }; @@ -35,7 +35,24 @@ pub const records = blk: { const groups: usize = @max(@as(usize, s.sk) -| 1, @as(usize, 1)); const threads: usize = if (sum) 256 else 32 * s.sk; const memory: usize = if (sum) 0 else groups * 16 * 32 * 4; - out[i] = .{ .name = s.name, .threads = @intCast(threads), .tg_memory = @intCast(memory), .offline = false }; + out[i] = .{ .name = s.name, .threads = @intCast(threads), .tg_memory = @intCast(memory), .m1 = if (sum) 1024 else 0, .m2 = if (sum) 1024 else 0, .offline = sum }; } break :blk out; }; + +test "byte lane records: standalone sums have measured offline limits; tensor ops stay unmeasured" { + const std = @import("std"); + var count: usize = 0; + for (records) |r| { + const sum = std.mem.endsWith(u8, r.name, "tf_lane_sums"); + try std.testing.expectEqual(sum, r.offline); + try std.testing.expectEqual(@as(u32, if (sum) 1024 else 0), r.m1); + try std.testing.expectEqual(r.m1, r.m2); + if (sum) { + count += 1; + try std.testing.expectEqual(@as(u32, 256), r.threads); + try std.testing.expectEqual(@as(u32, 0), r.tg_memory); + } + } + try std.testing.expectEqual(@as(usize, 4), count); +} From 01e6919ca1c189174c776f342d68f5e3284458fe Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Wed, 11 Nov 2026 04:08:34 +0000 Subject: [PATCH 15/78] metal: keep short byte-lane projections on a scalar row grid Only whole-prompt layouts compile the independent row-block header. Short layouts retain a scalar group coordinate and their original row extent; guarded sums remain bounds-safe. Keep older-chip sources on that short variant as well. Compare both short and long programs against frozen main at every maintained byte-check shape and row count. Refs #2818 --- zig/kernels/metal/core/lane_projection.metal | 17 ++++++++++++++++- zig/qwen27_quant_test.zig | 5 +++-- zig/src/core/lane_projection.zig | 5 ++++- zig/tests/core_lane_long.zig | 13 +++++++++---- 4 files changed, 32 insertions(+), 8 deletions(-) diff --git a/zig/kernels/metal/core/lane_projection.metal b/zig/kernels/metal/core/lane_projection.metal index 05b55977cd1..7e65ddb08b6 100644 --- a/zig/kernels/metal/core/lane_projection.metal +++ b/zig/kernels/metal/core/lane_projection.metal @@ -12,6 +12,9 @@ using namespace mpp::tensor_ops; #ifndef TF_PRECOMPUTE_SUMS #define TF_PRECOMPUTE_SUMS 0 #endif +#ifndef TF_ROW_BLOCKS +#define TF_ROW_BLOCKS 0 +#endif #ifndef TF_YS #define TF_YS TF_N #endif @@ -174,13 +177,20 @@ inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, in const device float* XS [[buffer(5)]], #endif uint sgi [[simdgroup_index_in_threadgroup]], uint lanei [[thread_index_in_simdgroup]], - uint2 tg [[threadgroup_position_in_grid]]) { +#if TF_ROW_BLOCKS + uint2 tg [[threadgroup_position_in_grid]] +#else + uint tgx [[threadgroup_position_in_grid]] +#endif +) { +#if TF_ROW_BLOCKS const int row0 = int(tg.y) * 16, rows = mdims[0] - row0; if (rows <= 0) return; X += (size_t)row0 * TF_K; Y += (size_t)row0 * TF_YS; #if TF_PRECOMPUTE_SUMS XS += row0; +#endif #endif const int lane = int(lanei), sg = int(sgi); const short qid = short(lane) >> 2; @@ -190,8 +200,13 @@ inline void tf_stage_codes(threadgroup uint* stage, thread const tf_raw& raw, in #else const short fn = ((qid & 2) | (short(lane) & 1)) * 4; #endif +#if TF_ROW_BLOCKS const int M = min(rows, 16); const int tile = tf_tile(int(tg.x)); +#else + const int M = mdims[0]; + const int tile = tf_tile(int(tgx)); +#endif const int n0 = tile * NT; const int g_begin = TF_G0 + (sg * TF_GN) / TF_SK, g_end = TF_G0 + ((sg + 1) * TF_GN) / TF_SK; constexpr auto desc = matmul2d_descriptor(16, NT, GS, false, true, false, matmul2d_descriptor::mode::multiply); diff --git a/zig/qwen27_quant_test.zig b/zig/qwen27_quant_test.zig index 1c151d8a718..783ec7c2eef 100644 --- a/zig/qwen27_quant_test.zig +++ b/zig/qwen27_quant_test.zig @@ -138,7 +138,7 @@ const wide_formats = [_]lane_projection.Format{ .{ .bits = 3, .group = 32 }, .{ .bits = 2, .group = 64 }, .{ .bits = 8, .group = 128 }, }; -test "qwen27 quant: non-reg layouts admit whole prompt chunks without changing source" { +test "qwen27 quant: non-reg layouts admit whole prompt chunks with row-block source variants" { for (wide_formats) |format| { const long = lane_projection.Layout{ .n = 64, .k = 5120, .format = format, .rows = 2048 }; try long.validate(); @@ -148,7 +148,8 @@ test "qwen27 quant: non-reg layouts admit whole prompt chunks without changing s defer a.free(text); const old = try lane_projection.source(a, short); defer a.free(old); - try t.expectEqualStrings(old, text); + try t.expect(std.mem.indexOf(u8, old, "#define TF_ROW_BLOCKS 0\n") != null); + try t.expect(std.mem.indexOf(u8, text, "#define TF_ROW_BLOCKS 1\n") != null); for ([_]usize{ 0, 15, 17, 2064 }) |rows| { short.rows = rows; try t.expectError(error.InvalidSchedule, short.validate()); diff --git a/zig/src/core/lane_projection.zig b/zig/src/core/lane_projection.zig index a27141818d9..ca7adc4121c 100644 --- a/zig/src/core/lane_projection.zig +++ b/zig/src/core/lane_projection.zig @@ -120,6 +120,7 @@ pub fn source(a: std.mem.Allocator, l: Layout) ![]u8 { } try text.appendSlice(a, if (l.ranges.len == 0) " return j;\n}\n" else " return 0;\n}\n"); try text.print(a, "#define TF_PRECOMPUTE_SUMS {d}\n", .{@intFromBool(l.precompute_sums)}); + try text.print(a, "#define TF_ROW_BLOCKS {d}\n", .{@intFromBool(!l.reg and !l.cooperative and l.rows > 16)}); if (l.reg) try text.appendSlice(a, "#define TF_REG 1\n"); if (l.y_stride) |s| try text.print(a, "#define TF_YS {d}\n", .{s}); try text.appendSlice(a, ks.core_lane_projection); @@ -249,7 +250,9 @@ pub const Projection = struct { const units = device.tensorUnits(); if (!units and (l.format.bits != 6 or l.format.group != 32 or l.cooperative or l.reg or l.precompute_sums)) return error.TensorUnitsRequired; const capacity: usize = if (l.cooperative or !units) 16 else l.rows; - const text = try source(a, l); + var kernel_layout = l; + if (!units and !l.cooperative and !l.reg) kernel_layout.rows = capacity; + const text = try source(a, kernel_layout); defer a.free(text); const adapted = try frags.source(device, a, text); defer a.free(adapted); diff --git a/zig/tests/core_lane_long.zig b/zig/tests/core_lane_long.zig index 502aa57f3da..6add4d36ffe 100644 --- a/zig/tests/core_lane_long.zig +++ b/zig/tests/core_lane_long.zig @@ -95,11 +95,12 @@ pub fn check(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, rng: st defer p_long.deinit(); var short = layout; short.rows = 16; - var p_short = try lane.Projection.init(a, device, short); + const p_short = try lane.Projection.init(a, device, short); defer p_short.deinit(); const old_pipe = try frozen(a, device, short); - p_short.pipe.deinit(); - p_short.pipe = old_pipe; + defer old_pipe.deinit(); + var p_old = p_short; + p_old.pipe = old_pipe; const raw = try a.alloc(u32, layout.sourceWords()); defer a.free(raw); for (raw) |*v| v.* = rng.int(u32); @@ -138,7 +139,10 @@ pub fn check(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, rng: st const xs: ?lane.Sums = if (external) .{ .ref = .{ .buf = given }, .stride = sum_stride } else null; clear(want, rows, n, stride, false); clear(got, rows, n, stride, true); - try chunks(queue, p_short, x, w, sb, xs, want, rows); + try chunks(queue, p_old, x, w, sb, xs, want, rows); + try chunks(queue, p_short, x, w, sb, xs, got, rows); + if (!same(want, got)) return error.ShortLaneBytesDiffer; + clear(got, rows, n, stride, true); try chunks(queue, p_long, x, w, sb, xs, got, rows); if (!same(want, got)) { std.debug.print("LANE LONG DIFFER bits={d} group={d} n={d} k={d} sk={d} rows={d} external={}\n", .{ format.bits, format.group, n, k, sk, rows, external }); @@ -149,5 +153,6 @@ pub fn check(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, rng: st }; } } + std.debug.print("LANE SHORT CHECK PASS cases={d}: current short variant equals frozen main\n", .{count}); return count; } From 92ef2bfdb80de70988a314ad708bc735e5b5c8d7 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Tue, 10 Nov 2026 23:13:35 -0800 Subject: [PATCH 16/78] metal: opt into whole-prompt row blocks explicitly Leave other affine lane callers on their scalar 16-row source and capacity. Only the Qwen prompt key and byte fixture request independent row blocks. Pin the shared shader bytes, describe the row layout accurately and cover the 16/17-row boundary. Refs #2818 --- zig/gemma4_shared_affine_test.zig | 4 +-- zig/qwen27_quant_test.zig | 39 ++++++++++++++++++++++++++- zig/src/core/lane_projection.zig | 10 ++++--- zig/src/families/qwen27/quant_gpu.zig | 1 + zig/tests/core_lane_long.zig | 6 ++--- 5 files changed, 50 insertions(+), 10 deletions(-) diff --git a/zig/gemma4_shared_affine_test.zig b/zig/gemma4_shared_affine_test.zig index 33e99b67ad7..4ab73b90968 100644 --- a/zig/gemma4_shared_affine_test.zig +++ b/zig/gemma4_shared_affine_test.zig @@ -109,8 +109,8 @@ test "landed Metal source bytes are pinned without shader compilation" { { var hash: [32]u8 = undefined; std.crypto.hash.sha2.Sha256.hash(ks.core_lane_projection, &hash, .{}); - // 66155d5ee's bytes (8-bit lane codes and producer sums, Refs #2502, under the metal_lane_q8 contract) - try std.testing.expectEqualStrings("159f39224b5bb60fd2ced301496d91662d2958056bd291d13d8a5db5186dd41a", &std.fmt.bytesToHex(hash, .lower)); + // Packed affine lanes with producer sums and independent prompt row blocks. + try std.testing.expectEqualStrings("7befe2028dae5036541f311f230a02d5652e634450ef16dc66dce6e15e935277", &std.fmt.bytesToHex(hash, .lower)); } } test "released lane parameters and conservative scratch size" { diff --git a/zig/qwen27_quant_test.zig b/zig/qwen27_quant_test.zig index 783ec7c2eef..8fcfa9a5bcf 100644 --- a/zig/qwen27_quant_test.zig +++ b/zig/qwen27_quant_test.zig @@ -140,7 +140,7 @@ const wide_formats = [_]lane_projection.Format{ test "qwen27 quant: non-reg layouts admit whole prompt chunks with row-block source variants" { for (wide_formats) |format| { - const long = lane_projection.Layout{ .n = 64, .k = 5120, .format = format, .rows = 2048 }; + const long = lane_projection.Layout{ .n = 64, .k = 5120, .format = format, .rows = 2048, .row_blocks = true }; try long.validate(); var short = long; short.rows = 16; @@ -219,3 +219,40 @@ test "qwen27 quant: cooperative and older-chip projections keep their 16-row lim const over_budget = lane_projection.Layout{ .n = 2048, .k = 5120, .format = .{ .bits = 8, .group = 128 }, .sk = 8 }; try t.expectError(error.ThreadgroupMemoryLimit, over_budget.validate()); } + +test "qwen27 quant: default non-reg callers retain scalar source and 16-row capacity" { + for (wide_formats) |format| { + var r: mtl.Record = .{}; + const layout = lane_projection.Layout{ .n = 64, .k = 5120, .format = format }; + const text = try lane_projection.source(a, layout); + defer a.free(text); + try t.expect(std.mem.indexOf(u8, text, "#define TF_ROW_BLOCKS 0\n") != null); + const p = try lane_projection.Projection.init(a, r.device(), layout); + defer p.deinit(); + try t.expectEqual(@as(u32, 16), p.max_rows); + const ref = lane_projection.Ref{ .buf = r.virtual() }; + try p.encode(r.encoder(), ref, ref, ref, ref, 16); + try t.expectEqual(@as(usize, 1), r.launches[0].grid.height); + try t.expectError(error.InvalidRows, p.encode(r.encoder(), ref, ref, ref, ref, 17)); + } +} + +test "qwen27 quant: non-reg prompt key opts in only above 16 rows" { + for (wide_formats) |format| for ([_]u32{ 16, 17 }) |rows| { + var r: mtl.Record = .{}; + var k = quant.Kernels{ .allocator = a, .device = r.device(), .long_prompt = true, .prompt = true }; + defer k.deinit(); + var p = matrix(&r, 64); + p.bits = format.bits; + p.group = @intCast(format.group); + const ref = quant.Ref{ .buffer = r.virtual() }; + const sums: ?quant.Sums = if (format.bits == 8 and format.group == 64) .{ .ref = ref, .stride = 32 } else null; + try k.quant(r.encoder(), p, ref, sums, ref, ref, rows); + var entries = k.lanes.iterator(); + const entry = entries.next().?; + try t.expectEqual(rows > 16, entry.key_ptr.long); + try t.expectEqual(@as(u32, if (rows > 16) 2048 else 16), entry.value_ptr.max_rows); + try t.expectEqual(@as(usize, 1), r.count); + try t.expectEqual(@as(usize, if (rows > 16) 2 else 1), r.launches[0].grid.height); + }; +} diff --git a/zig/src/core/lane_projection.zig b/zig/src/core/lane_projection.zig index ca7adc4121c..14e6029d370 100644 --- a/zig/src/core/lane_projection.zig +++ b/zig/src/core/lane_projection.zig @@ -36,8 +36,10 @@ pub const Layout = struct { cooperative: bool = false, /// 4-bit g64 codes in reg_layout, unpacked into tensor-op registers; every word equals the cooperative kernel's. reg: bool = false, - /// Rows one encode may take: it sizes scratch, never the kernel source; cooperative encodes still take 16. + /// Row capacity for reg layouts and opted-in row blocks; cooperative and scalar layouts take 16. rows: usize = reg_rows, + /// Non-reg prompt callers opt into independent 16-row blocks on the grid's y. + row_blocks: bool = false, /// Elements between output rows (default N): a member of a fused stack writes its columns of the fused rows. y_stride: ?usize = null, @@ -120,7 +122,7 @@ pub fn source(a: std.mem.Allocator, l: Layout) ![]u8 { } try text.appendSlice(a, if (l.ranges.len == 0) " return j;\n}\n" else " return 0;\n}\n"); try text.print(a, "#define TF_PRECOMPUTE_SUMS {d}\n", .{@intFromBool(l.precompute_sums)}); - try text.print(a, "#define TF_ROW_BLOCKS {d}\n", .{@intFromBool(!l.reg and !l.cooperative and l.rows > 16)}); + try text.print(a, "#define TF_ROW_BLOCKS {d}\n", .{@intFromBool(l.row_blocks and !l.reg and !l.cooperative and l.rows > 16)}); if (l.reg) try text.appendSlice(a, "#define TF_REG 1\n"); if (l.y_stride) |s| try text.print(a, "#define TF_YS {d}\n", .{s}); try text.appendSlice(a, ks.core_lane_projection); @@ -242,14 +244,14 @@ pub const Projection = struct { threads_per_slice: usize = 32, reg: bool = false, cooperative: bool = false, - /// Rows one encode takes at most: the layout's rows on tensor units, 16 on cooperative or older-chip paths. + /// Capacity is the reg or opted-in row count on tensor units, otherwise 16. max_rows: u32 = reg_rows, pub fn init(a: std.mem.Allocator, device: mtl.Device, l: Layout) !Projection { try l.validate(); const units = device.tensorUnits(); if (!units and (l.format.bits != 6 or l.format.group != 32 or l.cooperative or l.reg or l.precompute_sums)) return error.TensorUnitsRequired; - const capacity: usize = if (l.cooperative or !units) 16 else l.rows; + const capacity: usize = if (l.cooperative or !units or (!l.reg and !l.row_blocks)) 16 else l.rows; var kernel_layout = l; if (!units and !l.cooperative and !l.reg) kernel_layout.rows = capacity; const text = try source(a, kernel_layout); diff --git a/zig/src/families/qwen27/quant_gpu.zig b/zig/src/families/qwen27/quant_gpu.zig index be2149254ae..19fcd480139 100644 --- a/zig/src/families/qwen27/quant_gpu.zig +++ b/zig/src/families/qwen27/quant_gpu.zig @@ -174,6 +174,7 @@ pub const Kernels = struct { .output = if (cut[1] > 0) .f32 else .bf16, .precompute_sums = reg or (p.bits == 8 and p.group == 64), .reg = reg, + .row_blocks = long and !reg, .rows = if (long) lane.max_reg_rows else if (reg) lane.reg_rows else 16, }); k.joined(entry.value_ptr); diff --git a/zig/tests/core_lane_long.zig b/zig/tests/core_lane_long.zig index 6add4d36ffe..c2b00f105ad 100644 --- a/zig/tests/core_lane_long.zig +++ b/zig/tests/core_lane_long.zig @@ -1,4 +1,4 @@ -//! Whole prompt lanes at each width against main's frozen 16-row kernel, including strided oQ member formats. +//! Whole prompt lanes at each width against a frozen 16-row kernel, including strided oQ member formats. const std = @import("std"); const mtl = @import("metal"); const lane = @import("core_lane"); @@ -90,7 +90,7 @@ pub fn check(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, rng: st const ng = k / format.group; // Group128 SK8 exceeds the unchanged admission bound; compare its supported SK4 variant on both sides. const sk = if (format.group == 128 and shape[2] == 8) 4 else shape[2]; - const layout = lane.Layout{ .n = n, .k = k, .format = format, .sk = sk, .rows = most, .precompute_sums = fi == 0 or fi == 6, .y_stride = stride }; + const layout = lane.Layout{ .n = n, .k = k, .format = format, .sk = sk, .rows = most, .row_blocks = true, .precompute_sums = fi == 0 or fi == 6, .y_stride = stride }; const p_long = try lane.Projection.init(a, device, layout); defer p_long.deinit(); var short = layout; @@ -134,7 +134,7 @@ pub fn check(a: std.mem.Allocator, device: mtl.Device, queue: mtl.Queue, rng: st defer want.deinit(); const got = try device.buffer(most * stride * 2, opts); defer got.deinit(); - for ([_]u32{ 17, 1065, most }) |rows| for ([_]bool{ false, true }) |external| { + for ([_]u32{ 16, 17, 1065, most }) |rows| for ([_]bool{ false, true }) |external| { if (external and fi != 0) continue; const xs: ?lane.Sums = if (external) .{ .ref = .{ .buf = given }, .stride = sum_stride } else null; clear(want, rows, n, stride, false); From 88dfb1c8ab586ff8f19ab28d74583cbdcaaf14bf Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 12:14:41 +0000 Subject: [PATCH 17/78] changelog: #2851 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8975e9de0af..535cd9c58e2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Metal on M5: prompt lane projections at every width (8-, 6-, 5-, 3- and 2-bit and the oQ mixes) take a chunk's 16-row blocks in one dispatch, as 4-bit already did; Qwen3.8-27B MLX 8-bit cold prompts are about 33% faster on an M5 Ultra (769 to 1,029 tok/s at 2K and 8K) with identical tokens, and decode is unchanged (#2851, for #2818). + - GLM-5.3-Flash: the DFlash2 drafter's startup line and docs give its licence as Apache-2.0, as its model card now says (#3084, for #1583). - CUDA kernels: a bf16 format for grouped experts runs a checkpoint's own bf16 expert weights in mma fragment blocks, with relu² up and fp32 or bf16 down, in decode and prompt forms; no served checkpoint selects it yet (Nemotron 3.5 Transcribe, #1631, will), and every other expert format's dispatch is unchanged (#1630). Thanks to @BobClawblaw. From aca24c73fa617bf965ec45f357786591d0557d78 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Tue, 10 Nov 2026 19:25:40 +0900 Subject: [PATCH 18/78] glm cuda: TF_GLM_CONTEXT_CEILING serves windows past 1,048,576 tokens, read at GLM's open GLM-5.3-Flash's attention keeps no rotary positions, so only its trained length bounds the window. TF_GLM_CONTEXT_CEILING, a .glm row, sets the window --context may ask for, up to 2^31 - 32, where the slots keep i32 positions. GLM reads it at its open once its .glm rows are checked; unset, the ceiling and the refusal are main's. A window past the trained length says so at startup. --- docs/environment.md | 1 + docs/recipes/glm-5.3-flash-cuda.md | 9 ++++++ zig/src/families/glm/cuda_env_test.zig | 33 +++++++++++++++++++++- zig/src/families/glm/cuda_native.zig | 17 +++++++++++- zig/src/native/cuda.zig | 38 ++++++++++++++++++++++++-- zig/src/native/cuda_window.zig | 21 ++++++++++++++ zig/src/platform/environment.zig | 2 ++ 7 files changed, 117 insertions(+), 4 deletions(-) diff --git a/docs/environment.md b/docs/environment.md index 0af40a6f12c..b1818515681 100644 --- a/docs/environment.md +++ b/docs/environment.md @@ -43,6 +43,7 @@ Diagnostics quote and escape the supplied text so it cannot create another termi | `TF_GLM_SHARED` | Exactly `0` or `1`; default on (GLM windows share a forward); read on rank 0 | | `TF_GLM_DFLASH2_LOG` | Whole unsigned 64-bit count of DFlash2 rounds whose drafts are logged; default 0 | | `TF_GLM_DRAFT_CONFIDENCE` | Finite number 0 through 1 where a DFlash2 chain's chance product stops it; default 0.7 | +| `TF_GLM_CONTEXT_CEILING` | GLM-5.3-Flash on CUDA: whole token count 1 through 2147483616, the window `--context` may ask for and `--context 0` serves, in place of config.json's 1048576; read at GLM's open. Past that length the model's quality is the operator's to check. Unset keeps config.json's window | | `TF_PROMPT_ROWS` | Native serving: a multiple of 16 from 128 through 2048; unset keeps the chip/family default | | `TF_ANE_COLUMNS` | Whole u32 count; zero disables the experimental share; default zero | | `TF_FLASHNEXT_SCALE_INDEX` | Flash Next on Metal before tensor units: exactly `0` or `1`; `1` has the 4-bit expert kernels read each group's scale and bias through a 16-bit index into its tensor's table of pairs, the same values with fewer weight bytes; default 0 | diff --git a/docs/recipes/glm-5.3-flash-cuda.md b/docs/recipes/glm-5.3-flash-cuda.md index 93f671bd8fe..e8ed3350863 100644 --- a/docs/recipes/glm-5.3-flash-cuda.md +++ b/docs/recipes/glm-5.3-flash-cuda.md @@ -250,6 +250,15 @@ Long prompts and long runs of replies heat a GB10. The runbook's [heat bands](.. 1,035,295-token prompt at TP=3), [CPU pinning](../../RUNBOOK.md#gb10-cpu-pinning-for-long-cuda-runs) and [GPU clock cap](../../RUNBOOK.md#gb10-gpu-clock-cap-for-long-cuda-runs) each carry a measured GLM-5.3-Flash run. +## Windows past 1,048,576 tokens + +GLM-5.3-Flash's attention keeps no rotary positions (`qk_rope_head_dim` is 0, and KDA keeps none), so only its trained +length, config.json's 1,048,576 tokens, bounds the window. `TF_GLM_CONTEXT_CEILING` sets the window `--context` may ask +for and `--context 0` serves. GLM reads it at its open, and the startup log says when the window passes the trained +length, past which the quality is the operator's to check. On three GB10s with `Bizuayeu/GLM-5.3-Flash-XSFP4`, +`TF_GLM_CONTEXT_CEILING=2097152`, `--context 2097152 --parallel 1` and no prompt cache, rank 0 took 4,097 KV pages, and +a 1,996,430-token prompt answered all three of its passphrases. + ## Shared KDA operator `cuda.kda` owns the fused and three-kernel chains and single-layer and batched saved-row replay. diff --git a/zig/src/families/glm/cuda_env_test.zig b/zig/src/families/glm/cuda_env_test.zig index f7a6e7de348..023b4cebde7 100644 --- a/zig/src/families/glm/cuda_env_test.zig +++ b/zig/src/families/glm/cuda_env_test.zig @@ -18,6 +18,8 @@ const switches = [_]Name{ }; const counts = [_][:0]const u8{"TF_GLM_DFLASH2_LOG"}; const reals = [_][:0]const u8{"TF_GLM_DRAFT_CONFIDENCE"}; +/// The window GLM may serve past config.json's: GLM's raisedCeiling reads it at the open, after these rows' check. +const windows = [_][:0]const u8{"TF_GLM_CONTEXT_CEILING"}; /// The row by name without `rule`, so a missing row fails an assertion instead of panicking. fn row(name: []const u8) !e.Rule { @@ -54,12 +56,13 @@ test "GLM strict env: every GLM reader's name is a glm row, and the glm scope ho for (switches) |s| try t.expectEqual(e.Kind.flag, (try row(s.name)).kind); for (counts) |name| try t.expectEqual(e.Kind.u64, (try row(name)).kind); for (reals) |name| try t.expectEqual(e.Kind.real, (try row(name)).kind); + for (windows) |name| try t.expectEqual(e.Kind.u64, (try row(name)).kind); var n: usize = 0; for (e.rules) |r| if (r.scope == .glm) { try t.expect(std.mem.startsWith(u8, r.name, "TF_GLM_")); n += 1; }; - try t.expectEqual(switches.len + counts.len + reals.len, n); + try t.expectEqual(switches.len + counts.len + reals.len + windows.len, n); } test "GLM strict env: switches keep unset, blank, 0 and 1, and refuse words and padded digits by name" { @@ -137,3 +140,31 @@ test "GLM strict env: the DFlash2 cut is a finite chance from 0 to 1, read as a try t.expectError(error.InvalidEnvironment, confidence()); try t.expectError(error.InvalidEnvironment, @import("cuda_env.zig").startup()); } + +test "GLM strict env: the context ceiling is a whole token count from 1 to 2147483616; unset or blank keeps config.json's" { + for (windows) |name| { + _ = try row(name); + for ([_]?[]const u8{ null, "", " \t" }) |raw| { + try t.expectEqual(@as(?u64, null), try e.value(?u64, name, raw, null)); + if (raw) |v| try check(name, v, true); + } + for ([_][]const u8{ "1", "2097152", "2147483616" }) |raw| { + try t.expectEqual(@as(?u64, try std.fmt.parseInt(u64, raw, 10)), try e.value(?u64, name, raw, null)); + try check(name, raw, true); + } + for ([_][]const u8{ "0", "2147483617", "+2097152", " 2097152", "2097152 ", "2M", "0x200000", "2097152.0" }) |raw| try check(name, raw, false); + } +} + +test "GLM strict env: the open's check refuses a malformed ceiling before ranks, NCCL or weights" { + const startup = @import("cuda_env.zig").startup; + const name = "TF_GLM_CONTEXT_CEILING"; + _ = try row(name); + _ = env.unsetenv(name); + defer _ = env.unsetenv(name); + try startup(); + _ = env.setenv(name, "2097152", 1); + try startup(); + _ = env.setenv(name, "2M", 1); + try t.expectError(error.InvalidEnvironment, startup()); +} diff --git a/zig/src/families/glm/cuda_native.zig b/zig/src/families/glm/cuda_native.zig index 904d7cb72aa..01323464187 100644 --- a/zig/src/families/glm/cuda_native.zig +++ b/zig/src/families/glm/cuda_native.zig @@ -22,6 +22,15 @@ pub const model_type = "glm5_next"; pub const formats: []const []const u8 = &.{ "modelopt-nvfp4", "modelopt-mixed", "exl3" }; // cc-defer: a fixed window; the Python family sizes it from free memory (capacity.admit), port that when serving needs pub const default_context: i64 = 8192; +/// TF_GLM_CONTEXT_CEILING: --context past config.json's window, as the attention has no rotary positions to stretch. +pub const ceiling_switch = "TF_GLM_CONTEXT_CEILING"; + +/// The window TF_GLM_CONTEXT_CEILING raises --context's ceiling to, read once the open has checked every .glm row. +pub fn raisedCeiling() !?u64 { + try @import("cuda_env.zig").startup(); + const environment = @import("sys").environment; + return try environment.value(?u64, ceiling_switch, environment.Process.get(.{}, ceiling_switch), null); +} pub const max_segments: u32 = 1; pub const prompt_rows: u32 = fw.prefill_rows; /// Chunk rows change prompt arithmetic: states keep and resume only where a cold pass cuts (cuda_prompt_cuts). @@ -581,11 +590,17 @@ test "capacity keeps a verify window past the context and a multiple of 64" { test "the plan's window, the context and a verify window, takes the pages of a stream's table at that capacity" { const l: kv.Layout = .{ .fixed_bytes = 0, .token_bytes = 1, .page_tokens = @import("cuda_state.zig").page_rows, .snap_bytes = 0 }; - for ([_]usize{ 8192, 8191, 245_760, 524_288 - 33, 524_288 - 32, 524_288 - 31, 1_048_576 }) |context| { + for ([_]usize{ 8192, 8191, 245_760, 524_288 - 33, 524_288 - 32, 524_288 - 31, 1_048_576, 2_097_152, 2_097_152 + 1 }) |context| { try std.testing.expectEqual(pages_mod.tableLen(capacity(context)), l.pages(context + window_rows)); } } +test "TF_GLM_CONTEXT_CEILING's largest window is the last whose slots, verify rows included, keep i32 positions" { + const most = @import("sys").environment.rule(ceiling_switch).max; + try std.testing.expect(capacity(@intCast(most)) <= @as(usize, std.math.maxInt(i32)) + 1); + try std.testing.expect(capacity(@intCast(most + 1)) > @as(usize, std.math.maxInt(i32)) + 1); +} + // Enough words for every offered frame pair and image, even with a small text context. pub fn linkCapacity(cap: usize, media: ?images_mod.Limits) !usize { const lim = media orelse return cap; diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 57c34a6bac1..c9b1be77b68 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -181,18 +181,26 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O return null; } // a family whose rope can stretch (Qwen's YaRN, from its config or --rope-yarn) names its own ceiling - const native = if (comptime @hasDecl(F, "contextCeiling")) F.contextCeiling(a, io, o.dir, o.rope_yarn) catch |e| { + const own = if (comptime @hasDecl(F, "contextCeiling")) F.contextCeiling(a, io, o.dir, o.rope_yarn) catch |e| { problem.* = try std.fmt.allocPrint(a, "{s} ({s}): YaRN takes rope_type yarn, factor, original_max_position_embeddings and their product as max_position_embeddings, and --rope-yarn a factor from 1 matching it", .{ if (o.rope_yarn != null) "--rope-yarn" else "the checkpoint's rope scaling", @errorName(e) }); return null; } else modelContext(a, io, o.dir); + // a family with no positions to stretch (GLM-5.3-Flash) raises it with a row of its own, read at its open + const raised: ?u64 = if (comptime @hasDecl(F, "raisedCeiling")) F.raisedCeiling() catch |e| { + problem.* = F.explain(null, e) orelse @errorName(e); + return null; + } else null; + const switch_name = if (comptime @hasDecl(F, "raisedCeiling")) F.ceiling_switch else ""; + const native: i64 = if (raised) |n| @intCast(n) else own; const window = contextWindow(o.context, native, F.default_context) catch |e| { problem.* = switch (e) { error.Negative => try std.fmt.allocPrint(a, "--context {d}: a token count, or 0 for the model's window", .{o.context.?}), error.NoNative => "--context 0 asks for the model's window, and its config.json names none: give a token count", - error.PastNative => try std.fmt.allocPrint(a, "--context {d} exceeds this model's {d}-token window", .{ o.context.?, native }), + error.PastNative => if (raised != null) try windows.pastCeiling(a, o.context.?, native, switch_name) else try std.fmt.allocPrint(a, "--context {d} exceeds this model's {d}-token window", .{ o.context.?, native }), }; return null; }; + if (raised != null and window > own) std.log.info("{s}", .{try windows.pastTrained(a, window, own, switch_name)}); if (o.slide) { // a family learns once its module declares `slide` (its learner), `learnModel` and `linked` if (!@hasDecl(F, "slide")) { @@ -554,6 +562,32 @@ test "Flash Next takes kept prompt states and drafts together on any rank count" test "Flash Next takes --vision at family admission" { try @import("cuda_flashnext_refusal_test.zig").vision(); } +test "context ceilings: GLM reads its ceiling at its open, unset keeps main's refusal, and a malformed one is GLM's refusal" { + const env = platform.env; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try tmp.dir.writeFile(std.testing.io, .{ .sub_path = "config.json", .data = "{\"max_position_embeddings\": 1048576}" }); + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const dir = try tmp.dir.realPathFileAlloc(std.testing.io, ".", a); + var named: usize = 0; + inline for (registry) |F| if (comptime @hasDecl(F, "raisedCeiling")) { + defer _ = env.unsetenv(F.ceiling_switch); + var why: []const u8 = ""; + _ = env.unsetenv(F.ceiling_switch); + try std.testing.expect((try open(a, std.testing.allocator, std.testing.io, .{ .dir = dir, .model_type = F.model_type, .context = 2097152 }, &why)) == null); + try std.testing.expectEqualStrings("--context 2097152 exceeds this model's 1048576-token window", why); + _ = env.setenv(F.ceiling_switch, "2M", 1); + try std.testing.expect((try open(a, std.testing.allocator, std.testing.io, .{ .dir = dir, .model_type = F.model_type, .context = 2097152 }, &why)) == null); + try std.testing.expectEqualStrings(F.explain(null, error.InvalidEnvironment).?, why); + _ = env.setenv(F.ceiling_switch, "2097152", 1); + try std.testing.expect((try open(a, std.testing.allocator, std.testing.io, .{ .dir = dir, .model_type = F.model_type, .context = 2097153 }, &why)) == null); + try std.testing.expectEqualStrings("--context 2097153 exceeds the 2097152-token window " ++ F.ceiling_switch ++ " sets", why); + named += 1; + }; + try std.testing.expect(named >= 1); // GLM-5.3-Flash's TF_GLM_CONTEXT_CEILING +} test { _ = @import("cuda_health_test.zig"); _ = @import("cuda_vision_offer.zig"); diff --git a/zig/src/native/cuda_window.zig b/zig/src/native/cuda_window.zig index 355dc4d1b4b..7bd01019232 100644 --- a/zig/src/native/cuda_window.zig +++ b/zig/src/native/cuda_window.zig @@ -13,6 +13,17 @@ pub fn modelContext(a: Allocator, io: std.Io, dir: []const u8) i64 { return if (limit == .integer and limit.integer > 0) limit.integer else 0; } +/// A --context past the window a family's switch raises its ceiling to: the switch named. +pub fn pastCeiling(a: Allocator, requested: i64, ceiling: i64, switch_name: []const u8) ![]const u8 { + return std.fmt.allocPrint(a, "--context {d} exceeds the {d}-token window {s} sets", .{ requested, ceiling, switch_name }); +} + +/// The startup line of a window past the checkpoint's trained one, which only the family's switch serves. +pub fn pastTrained(a: Allocator, window: i64, trained: i64, switch_name: []const u8) ![]const u8 { + const words = "{s}: a {d}-token window, past the checkpoint's trained {d} tokens; past them the quality is the operator's to check"; + return std.fmt.allocPrint(a, words, .{ switch_name, window, trained }); +} + /// Prompt plus reply tokens a request may use: --context (0: the model's window), else the family default within it. pub fn contextWindow(requested: ?i64, native: i64, default: i64) error{ Negative, NoNative, PastNative }!i64 { const r = requested orelse return if (native > 0) @min(default, native) else default; @@ -31,3 +42,13 @@ test "context windows: the family default inside the model's, 0 for the model's, try std.testing.expectError(error.Negative, contextWindow(-5, 262144, 16384)); try std.testing.expectError(error.NoNative, contextWindow(0, 0, 16384)); } + +test "context windows: past a family's raised ceiling the refusal names its switch, and the startup line the trained window" { + const a = std.testing.allocator; + const past = try pastCeiling(a, 2097153, 2097152, "TF_GLM_CONTEXT_CEILING"); + defer a.free(past); + try std.testing.expectEqualStrings("--context 2097153 exceeds the 2097152-token window TF_GLM_CONTEXT_CEILING sets", past); + const line = try pastTrained(a, 2097152, 1048576, "TF_GLM_CONTEXT_CEILING"); + defer a.free(line); + try std.testing.expectEqualStrings("TF_GLM_CONTEXT_CEILING: a 2097152-token window, past the checkpoint's trained 1048576 tokens; past them the quality is the operator's to check", line); +} diff --git a/zig/src/platform/environment.zig b/zig/src/platform/environment.zig index 3ad7ecb4941..c73ab2066f4 100644 --- a/zig/src/platform/environment.zig +++ b/zig/src/platform/environment.zig @@ -82,6 +82,8 @@ pub const rules = [_]Rule{ .{ .name = "TF_GLM_SHARED", .scope = .glm, .kind = .flag, .forms = "0 or 1" }, .{ .name = "TF_GLM_DFLASH2_LOG", .scope = .glm, .kind = .u64, .forms = "a whole unsigned 64-bit round count (0 logs no drafts)" }, .{ .name = "TF_GLM_DRAFT_CONFIDENCE", .scope = .glm, .kind = .real, .max = 1, .forms = "a finite chain-chance cut from 0 to 1 (default 0.7)" }, + // 2^31 - 32: GLM's slots, the window and its 32 verify rows (glm/cuda_native.zig's capacity), keep i32 positions + .{ .name = "TF_GLM_CONTEXT_CEILING", .scope = .glm, .kind = .u64, .first_digit = true, .min = 1, .max = 2147483616, .forms = "a whole token count from 1 to 2147483616 (the window GLM-5.3-Flash may serve in place of config.json's)" }, .{ .name = "TF_DSV41_MHC_DEFER", .scope = .dsv41, .kind = .toggle, .trim = " \t", .no_word = false, .forms = "0, 1, true, false, on, off or yes (any case)" }, .{ .name = "TF_DSV41_COEF_LATE", .scope = .dsv41, .kind = .toggle, .trim = " \t", .no_word = false, .forms = "0, 1, true, false, on, off or yes (any case)" }, .{ .name = "TF_DSV41_LOOKUP", .scope = .dsv41, .kind = .toggle, .trim = " \t", .forms = "0, 1, true, false, on, off, yes or no (any case); blank turns lookup drafts off" }, From b6295aaa5159e0174f7e741cd95f1677ce383c9e Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Wed, 11 Nov 2026 02:17:33 +0900 Subject: [PATCH 19/78] glm cuda: TF_GLM_CONTEXT_CEILING's reader joins the .glm check in cuda_env.zig cuda_native.zig had passed the 600-line limit. ceiling_switch and raisedCeiling move to cuda_env.zig, which owns the .glm check, and the i32 test to cuda_env_test.zig; cuda_native.zig keeps the two names the CUDA open looks up. --- zig/src/families/glm/cuda_env.zig | 9 +++++++++ zig/src/families/glm/cuda_env_test.zig | 7 +++++++ zig/src/families/glm/cuda_native.zig | 17 ++--------------- 3 files changed, 18 insertions(+), 15 deletions(-) diff --git a/zig/src/families/glm/cuda_env.zig b/zig/src/families/glm/cuda_env.zig index c929b466b70..c5f37582c11 100644 --- a/zig/src/families/glm/cuda_env.zig +++ b/zig/src/families/glm/cuda_env.zig @@ -15,6 +15,15 @@ pub fn requestStartsAlone(learn_conversations: bool) !void { if (on and try kp.sharing() > 0) return error.GlmRequestStartsWithSharedKv; } +/// TF_GLM_CONTEXT_CEILING: --context past config.json's window, as the attention has no rotary positions to stretch. +pub const ceiling_switch = "TF_GLM_CONTEXT_CEILING"; + +/// The window TF_GLM_CONTEXT_CEILING raises --context's ceiling to, read once the open has checked every .glm row. +pub fn raisedCeiling() !?u64 { + try startup(); + return try environment.value(?u64, ceiling_switch, environment.Process.get(.{}, ceiling_switch), null); +} + test { _ = @import("cuda_env_test.zig"); } diff --git a/zig/src/families/glm/cuda_env_test.zig b/zig/src/families/glm/cuda_env_test.zig index 023b4cebde7..6f678d94dcb 100644 --- a/zig/src/families/glm/cuda_env_test.zig +++ b/zig/src/families/glm/cuda_env_test.zig @@ -168,3 +168,10 @@ test "GLM strict env: the open's check refuses a malformed ceiling before ranks, _ = env.setenv(name, "2M", 1); try t.expectError(error.InvalidEnvironment, startup()); } + +test "TF_GLM_CONTEXT_CEILING's largest window is the last whose slots, verify rows included, keep i32 positions" { + const capacity = @import("cuda_native.zig").capacity; + const most = e.rule(@import("cuda_env.zig").ceiling_switch).max; + try t.expect(capacity(@intCast(most)) <= @as(usize, std.math.maxInt(i32)) + 1); + try t.expect(capacity(@intCast(most + 1)) > @as(usize, std.math.maxInt(i32)) + 1); +} diff --git a/zig/src/families/glm/cuda_native.zig b/zig/src/families/glm/cuda_native.zig index 01323464187..099e47af6b2 100644 --- a/zig/src/families/glm/cuda_native.zig +++ b/zig/src/families/glm/cuda_native.zig @@ -22,15 +22,8 @@ pub const model_type = "glm5_next"; pub const formats: []const []const u8 = &.{ "modelopt-nvfp4", "modelopt-mixed", "exl3" }; // cc-defer: a fixed window; the Python family sizes it from free memory (capacity.admit), port that when serving needs pub const default_context: i64 = 8192; -/// TF_GLM_CONTEXT_CEILING: --context past config.json's window, as the attention has no rotary positions to stretch. -pub const ceiling_switch = "TF_GLM_CONTEXT_CEILING"; - -/// The window TF_GLM_CONTEXT_CEILING raises --context's ceiling to, read once the open has checked every .glm row. -pub fn raisedCeiling() !?u64 { - try @import("cuda_env.zig").startup(); - const environment = @import("sys").environment; - return try environment.value(?u64, ceiling_switch, environment.Process.get(.{}, ceiling_switch), null); -} +pub const ceiling_switch = @import("cuda_env.zig").ceiling_switch; +pub const raisedCeiling = @import("cuda_env.zig").raisedCeiling; pub const max_segments: u32 = 1; pub const prompt_rows: u32 = fw.prefill_rows; /// Chunk rows change prompt arithmetic: states keep and resume only where a cold pass cuts (cuda_prompt_cuts). @@ -595,12 +588,6 @@ test "the plan's window, the context and a verify window, takes the pages of a s } } -test "TF_GLM_CONTEXT_CEILING's largest window is the last whose slots, verify rows included, keep i32 positions" { - const most = @import("sys").environment.rule(ceiling_switch).max; - try std.testing.expect(capacity(@intCast(most)) <= @as(usize, std.math.maxInt(i32)) + 1); - try std.testing.expect(capacity(@intCast(most + 1)) > @as(usize, std.math.maxInt(i32)) + 1); -} - // Enough words for every offered frame pair and image, even with a small text context. pub fn linkCapacity(cap: usize, media: ?images_mod.Limits) !usize { const lim = media orelse return cap; From 1f7d03447c00d04b448b0c0df197c5df5812a4a1 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Wed, 11 Nov 2026 07:00:52 +0900 Subject: [PATCH 20/78] native cuda: the ceiling test and lookup leave cuda.zig for its own test file and cuda_window.zig Main grew cuda.zig to 587 lines, so the branch's additions passed the 600-line limit again. The open-time test of raised context ceilings moves to cuda_ceiling_test.zig, imported beside cuda_health_test.zig, and the lookup of a family's raised ceiling and its switch to cuda_window.zig's familyCeiling; cuda.zig keeps one call. --- zig/src/native/cuda.zig | 39 ++++------------------------ zig/src/native/cuda_ceiling_test.zig | 30 +++++++++++++++++++++ zig/src/native/cuda_window.zig | 13 ++++++++++ 3 files changed, 48 insertions(+), 34 deletions(-) create mode 100644 zig/src/native/cuda_ceiling_test.zig diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index c9b1be77b68..8036a4b1c82 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -186,21 +186,17 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O return null; } else modelContext(a, io, o.dir); // a family with no positions to stretch (GLM-5.3-Flash) raises it with a row of its own, read at its open - const raised: ?u64 = if (comptime @hasDecl(F, "raisedCeiling")) F.raisedCeiling() catch |e| { - problem.* = F.explain(null, e) orelse @errorName(e); - return null; - } else null; - const switch_name = if (comptime @hasDecl(F, "raisedCeiling")) F.ceiling_switch else ""; - const native: i64 = if (raised) |n| @intCast(n) else own; + const raised = windows.familyCeiling(F, problem) catch return null; + const native: i64 = if (raised.window) |n| @intCast(n) else own; const window = contextWindow(o.context, native, F.default_context) catch |e| { problem.* = switch (e) { error.Negative => try std.fmt.allocPrint(a, "--context {d}: a token count, or 0 for the model's window", .{o.context.?}), error.NoNative => "--context 0 asks for the model's window, and its config.json names none: give a token count", - error.PastNative => if (raised != null) try windows.pastCeiling(a, o.context.?, native, switch_name) else try std.fmt.allocPrint(a, "--context {d} exceeds this model's {d}-token window", .{ o.context.?, native }), + error.PastNative => if (raised.window != null) try windows.pastCeiling(a, o.context.?, native, raised.switch_name) else try std.fmt.allocPrint(a, "--context {d} exceeds this model's {d}-token window", .{ o.context.?, native }), }; return null; }; - if (raised != null and window > own) std.log.info("{s}", .{try windows.pastTrained(a, window, own, switch_name)}); + if (raised.window != null and window > own) std.log.info("{s}", .{try windows.pastTrained(a, window, own, raised.switch_name)}); if (o.slide) { // a family learns once its module declares `slide` (its learner), `learnModel` and `linked` if (!@hasDecl(F, "slide")) { @@ -562,34 +558,9 @@ test "Flash Next takes kept prompt states and drafts together on any rank count" test "Flash Next takes --vision at family admission" { try @import("cuda_flashnext_refusal_test.zig").vision(); } -test "context ceilings: GLM reads its ceiling at its open, unset keeps main's refusal, and a malformed one is GLM's refusal" { - const env = platform.env; - var tmp = std.testing.tmpDir(.{}); - defer tmp.cleanup(); - try tmp.dir.writeFile(std.testing.io, .{ .sub_path = "config.json", .data = "{\"max_position_embeddings\": 1048576}" }); - var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); - defer arena.deinit(); - const a = arena.allocator(); - const dir = try tmp.dir.realPathFileAlloc(std.testing.io, ".", a); - var named: usize = 0; - inline for (registry) |F| if (comptime @hasDecl(F, "raisedCeiling")) { - defer _ = env.unsetenv(F.ceiling_switch); - var why: []const u8 = ""; - _ = env.unsetenv(F.ceiling_switch); - try std.testing.expect((try open(a, std.testing.allocator, std.testing.io, .{ .dir = dir, .model_type = F.model_type, .context = 2097152 }, &why)) == null); - try std.testing.expectEqualStrings("--context 2097152 exceeds this model's 1048576-token window", why); - _ = env.setenv(F.ceiling_switch, "2M", 1); - try std.testing.expect((try open(a, std.testing.allocator, std.testing.io, .{ .dir = dir, .model_type = F.model_type, .context = 2097152 }, &why)) == null); - try std.testing.expectEqualStrings(F.explain(null, error.InvalidEnvironment).?, why); - _ = env.setenv(F.ceiling_switch, "2097152", 1); - try std.testing.expect((try open(a, std.testing.allocator, std.testing.io, .{ .dir = dir, .model_type = F.model_type, .context = 2097153 }, &why)) == null); - try std.testing.expectEqualStrings("--context 2097153 exceeds the 2097152-token window " ++ F.ceiling_switch ++ " sets", why); - named += 1; - }; - try std.testing.expect(named >= 1); // GLM-5.3-Flash's TF_GLM_CONTEXT_CEILING -} test { _ = @import("cuda_health_test.zig"); + _ = @import("cuda_ceiling_test.zig"); _ = @import("cuda_vision_offer.zig"); _ = @import("cuda_ledger.zig"); _ = kv_pages; diff --git a/zig/src/native/cuda_ceiling_test.zig b/zig/src/native/cuda_ceiling_test.zig new file mode 100644 index 00000000000..f9d1eacc0cd --- /dev/null +++ b/zig/src/native/cuda_ceiling_test.zig @@ -0,0 +1,30 @@ +//! A family's raised context ceiling (GLM's TF_GLM_CONTEXT_CEILING) through the real CUDA open, before any weights. +const std = @import("std"); +const env = @import("sys").env; +const cuda = @import("cuda.zig"); + +test "context ceilings: GLM reads its ceiling at its open, unset keeps main's refusal, and a malformed one is GLM's refusal" { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try tmp.dir.writeFile(std.testing.io, .{ .sub_path = "config.json", .data = "{\"max_position_embeddings\": 1048576}" }); + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const dir = try tmp.dir.realPathFileAlloc(std.testing.io, ".", a); + var named: usize = 0; + inline for (cuda.registry) |F| if (comptime @hasDecl(F, "raisedCeiling")) { + defer _ = env.unsetenv(F.ceiling_switch); + var why: []const u8 = ""; + _ = env.unsetenv(F.ceiling_switch); + try std.testing.expect((try cuda.open(a, std.testing.allocator, std.testing.io, .{ .dir = dir, .model_type = F.model_type, .context = 2097152 }, &why)) == null); + try std.testing.expectEqualStrings("--context 2097152 exceeds this model's 1048576-token window", why); + _ = env.setenv(F.ceiling_switch, "2M", 1); + try std.testing.expect((try cuda.open(a, std.testing.allocator, std.testing.io, .{ .dir = dir, .model_type = F.model_type, .context = 2097152 }, &why)) == null); + try std.testing.expectEqualStrings(F.explain(null, error.InvalidEnvironment).?, why); + _ = env.setenv(F.ceiling_switch, "2097152", 1); + try std.testing.expect((try cuda.open(a, std.testing.allocator, std.testing.io, .{ .dir = dir, .model_type = F.model_type, .context = 2097153 }, &why)) == null); + try std.testing.expectEqualStrings("--context 2097153 exceeds the 2097152-token window " ++ F.ceiling_switch ++ " sets", why); + named += 1; + }; + try std.testing.expect(named >= 1); // GLM-5.3-Flash's TF_GLM_CONTEXT_CEILING +} diff --git a/zig/src/native/cuda_window.zig b/zig/src/native/cuda_window.zig index 7bd01019232..2e94e48fc39 100644 --- a/zig/src/native/cuda_window.zig +++ b/zig/src/native/cuda_window.zig @@ -13,6 +13,19 @@ pub fn modelContext(a: Allocator, io: std.Io, dir: []const u8) i64 { return if (limit == .integer and limit.integer > 0) limit.integer else 0; } +/// The window a family's own row raises --context's ceiling to, and the switch its refusal names; none without one. +pub const Ceiling = struct { window: ?u64 = null, switch_name: []const u8 = "" }; + +/// A family's raised ceiling, read at its open (GLM's TF_GLM_CONTEXT_CEILING); a malformed row is the family's refusal. +pub fn familyCeiling(comptime F: type, problem: *[]const u8) error{Refused}!Ceiling { + if (comptime !@hasDecl(F, "raisedCeiling")) return .{}; + const window = F.raisedCeiling() catch |e| { + problem.* = F.explain(null, e) orelse @errorName(e); + return error.Refused; + }; + return .{ .window = window, .switch_name = F.ceiling_switch }; +} + /// A --context past the window a family's switch raises its ceiling to: the switch named. pub fn pastCeiling(a: Allocator, requested: i64, ceiling: i64, switch_name: []const u8) ![]const u8 { return std.fmt.allocPrint(a, "--context {d} exceeds the {d}-token window {s} sets", .{ requested, ceiling, switch_name }); From a69b066429f8c0f8f389a467f7e0061cf335ff32 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Thu, 12 Nov 2026 05:54:24 +0900 Subject: [PATCH 21/78] glm cuda: the capacity test leaves cuda_native.zig for cuda_pages_test.zig Main's eb29a68a5f added a line to cuda_native.zig, so merged with main the file comes to 601 lines, past the 600-line limit. The test that capacity keeps a verify window past the context, a multiple of 64, moves whole to cuda_pages_test.zig, which tests GLM's pages at that capacity: 594 lines in cuda_native.zig here, 595 merged with main. --- zig/src/families/glm/cuda_native.zig | 6 ------ zig/src/families/glm/cuda_pages_test.zig | 6 ++++++ 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/zig/src/families/glm/cuda_native.zig b/zig/src/families/glm/cuda_native.zig index 099e47af6b2..4a1f5cd1a5c 100644 --- a/zig/src/families/glm/cuda_native.zig +++ b/zig/src/families/glm/cuda_native.zig @@ -575,12 +575,6 @@ test "TF_KV_PAGES with --kv-cache fp8: auto and off keep windows at every rank c try std.testing.expect(explain(null, error.GlmKvHostNeedsPages) != null); } -test "capacity keeps a verify window past the context and a multiple of 64" { - try std.testing.expectEqual(@as(usize, 8256), capacity(8192)); - try std.testing.expectEqual(@as(usize, 330624), capacity(330529)); - try std.testing.expect(capacity(330529) % 64 == 0); -} - test "the plan's window, the context and a verify window, takes the pages of a stream's table at that capacity" { const l: kv.Layout = .{ .fixed_bytes = 0, .token_bytes = 1, .page_tokens = @import("cuda_state.zig").page_rows, .snap_bytes = 0 }; for ([_]usize{ 8192, 8191, 245_760, 524_288 - 33, 524_288 - 32, 524_288 - 31, 1_048_576, 2_097_152, 2_097_152 + 1 }) |context| { diff --git a/zig/src/families/glm/cuda_pages_test.zig b/zig/src/families/glm/cuda_pages_test.zig index 7fdfc60b657..76de1632cfa 100644 --- a/zig/src/families/glm/cuda_pages_test.zig +++ b/zig/src/families/glm/cuda_pages_test.zig @@ -530,3 +530,9 @@ test "every host allocation failure while making the pages leaves no device buff } }.run, .{}); } + +test "capacity keeps a verify window past the context and a multiple of 64" { + try expectEqual(@as(usize, 8256), native.capacity(8192)); + try expectEqual(@as(usize, 330624), native.capacity(330529)); + try expect(native.capacity(330529) % 64 == 0); +} From 36e9e0dcfd30e0948f6e1710e941041e3b315dfa Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 12:14:44 +0000 Subject: [PATCH 22/78] changelog: #2765 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 535cd9c58e2..144b294f96e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- GLM-5.3-Flash on CUDA: `TF_GLM_CONTEXT_CEILING` (opt-in) lets GLM open windows past 1,048,576 tokens, read once when the model opens; with it unset nothing changes, and cold prompts and decode on two and three GB10s read level with main (#2765, refs #2663). Thanks to @Bizuayeu. + - Metal on M5: prompt lane projections at every width (8-, 6-, 5-, 3- and 2-bit and the oQ mixes) take a chunk's 16-row blocks in one dispatch, as 4-bit already did; Qwen3.8-27B MLX 8-bit cold prompts are about 33% faster on an M5 Ultra (769 to 1,029 tok/s at 2K and 8K) with identical tokens, and decode is unchanged (#2851, for #2818). - GLM-5.3-Flash: the DFlash2 drafter's startup line and docs give its licence as Apache-2.0, as its model card now says (#3084, for #1583). From 976b408093bbea415f31868e3a3dfb7f07e0ee50 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Wed, 11 Nov 2026 08:39:08 +0000 Subject: [PATCH 23/78] fix(server): let the last --compact-at replace the earlier one A fraction after --compact-at auto left auto mode on, because the flag set compact_auto and a later fraction never cleared it, so the server kept the first value where every other flag keeps the last. Each occurrence now replaces both fields, which a config file's value under a command-line override depends on, and a filtered host step runs the test that fails without the change. --- zig/build/cuda_server_host.zig | 1 + zig/src/server/cli.zig | 4 +++- zig/src/server/compact_at_cli_test.zig | 23 +++++++++++++++++++++++ zig/src/server/root.zig | 1 + 4 files changed, 28 insertions(+), 1 deletion(-) create mode 100644 zig/src/server/compact_at_cli_test.zig diff --git a/zig/build/cuda_server_host.zig b/zig/build/cuda_server_host.zig index c77531a9b00..9da10918454 100644 --- a/zig/build/cuda_server_host.zig +++ b/zig/build/cuda_server_host.zig @@ -29,6 +29,7 @@ pub fn steps(b: *std.Build, host: std.Build.ResolvedTarget, cuda: *std.Build.Mod eos_policy.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"EOS policy:"} })).step); eos_policy.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = mods.tokenizer, .filters = &.{"EOS policy:"} })).step); eos_policy.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = kolibri1Module(b, host, .debug, cuda, mods.core, mods.lanes), .filters = &.{"EOS policy:"} })).step); + b.step("test-serve-flags-host", "Serve flag parsing a later flag replaces, without a GPU").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"--compact-at:"} })).step); b.step("test-dory-cli-host", "Dory CUDA loop option forwarding and invalid-value refusals").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"--device and --segments:"} })).step); b.step("test-dory-chat-host", "Dory chat template, EOS, reasoning split and qwen3_coder call on the CPU").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"Dory chat:"} })).step); b.step("test-embeddings-protocol-host", "Embedding request and response contracts without an admitted model or GPU").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests, .filters = &.{"embeddings protocol:"} })).step); diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig index 387dff1c8c6..82975c5f7e5 100644 --- a/zig/src/server/cli.zig +++ b/zig/src/server/cli.zig @@ -472,7 +472,9 @@ fn apply(a: Allocator, out: *Args, name: []const u8, value: ?[]const u8, u: *Usa out.decode_share = try float(u, a, name, v); if (!(out.decode_share.? >= 0 and out.decode_share.? < 1)) return fail(u, a, "argument --decode-share: expected a share from 0 to under 1: '{s}'", .{v}); } else if (is(name, "--prompt-seed")) out.prompt_seed = true else if (is(name, "--max-tokens")) out.max_tokens = try int(u, a, name, v) else if (is(name, "--temperature")) out.temperature = try float(u, a, name, v) else if (is(name, "--top-p")) out.top_p = try float(u, a, name, v) else if (is(name, "--top-k")) out.top_k = try int(u, a, name, v) else if (is(name, "--min-p")) out.min_p = try float(u, a, name, v) else if (is(name, "--instruct-sampling")) out.instruct_sampling = if (try @import("family_sampling.zig").pairsProblem(a, v)) |why| return fail(u, a, "argument {s}: {s}: '{s}'", .{ name, why, v }) else v else if (is(name, "--thinking")) out.thinking = true else if (is(name, "--no-thinking")) out.thinking = false else if (is(name, "--reasoning-effort")) out.reasoning_effort = v else if (is(name, "--thinking-budget")) out.thinking_budget = try int(u, a, name, v) else if (is(name, "--thinking-reserve")) out.thinking_reserve = try int(u, a, name, v) else if (is(name, "--loop-guard")) out.loop_guard = true else if (is(name, "--loop-guard-stop")) out.loop_guard_stop = true else if (is(name, "--drafter")) out.drafter = v else if (is(name, "--drafter-bits")) out.drafter_bits = @intCast(try int(u, a, name, v)) else if (is(name, "--record")) out.record = v else if (is(name, "--kv-cache")) out.kv_cache = v else if (is(name, "--prompt-precision")) out.prompt_precision = v else if (is(name, "--no-drafts")) out.no_drafts = true else if (is(name, "--keep-warm")) out.keep_warm = try int(u, a, name, v) else if (is(name, "--idle-unload")) out.idle_unload = try int(u, a, name, v) else if (is(name, "--compact-at")) { - if (std.mem.eql(u8, v, "auto")) out.compact_auto = true else { + out.compact_auto = std.mem.eql(u8, v, "auto"); // each occurrence replaces both fields, so the last one wins + out.compact_fraction = null; + if (!out.compact_auto) { const f = try float(u, a, name, v); if (!(f > 0 and f <= 1)) return fail(u, a, "argument --compact-at: expected auto or a fraction in (0, 1]: '{s}'", .{v}); out.compact_fraction = f; diff --git a/zig/src/server/compact_at_cli_test.zig b/zig/src/server/compact_at_cli_test.zig new file mode 100644 index 00000000000..31cd835c5af --- /dev/null +++ b/zig/src/server/compact_at_cli_test.zig @@ -0,0 +1,23 @@ +//! ``--compact-at`` given twice: each occurrence replaces both of its fields, so the last one wins. +const std = @import("std"); +const cli = @import("cli.zig"); + +test "--compact-at: a later occurrence replaces an earlier one in full" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var u: cli.Usage = .{}; + const fraction_last = try cli.parse(a, &.{ "m", "--compact-at", "auto", "--compact-at=0.5" }, &u); + try std.testing.expect(!fraction_last.compact_auto); + try std.testing.expectEqual(@as(?f64, 0.5), fraction_last.compact_fraction); + const auto_last = try cli.parse(a, &.{ "m", "--compact-at=0.5", "--compact-at", "auto" }, &u); + try std.testing.expect(auto_last.compact_auto); + try std.testing.expectEqual(@as(?f64, null), auto_last.compact_fraction); + const auto = try cli.parse(a, &.{ "m", "--compact-at", "auto" }, &u); + try std.testing.expect(auto.compact_auto and auto.compact_fraction == null); + const fraction = try cli.parse(a, &.{ "m", "--compact-at", "0.25" }, &u); + try std.testing.expect(!fraction.compact_auto); + try std.testing.expectEqual(@as(?f64, 0.25), fraction.compact_fraction); + try std.testing.expectError(error.Usage, cli.parse(a, &.{ "m", "--compact-at", "1.5" }, &u)); + try std.testing.expectEqualStrings("argument --compact-at: expected auto or a fraction in (0, 1]: '1.5'", u.message); +} diff --git a/zig/src/server/root.zig b/zig/src/server/root.zig index 3bfbc4ad96c..66d0da36073 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -76,6 +76,7 @@ test { _ = @import("cli.zig"); _ = @import("cli_test.zig"); _ = @import("cli_profile_test.zig"); + _ = @import("compact_at_cli_test.zig"); _ = @import("unix_socket_cli_test.zig"); _ = @import("effort_cli_audit_test.zig"); _ = @import("live.zig"); From 6482de24ae06de68a91b3d468adeb6aef036cd76 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 12:14:45 +0000 Subject: [PATCH 24/78] changelog: #2883 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 144b294f96e..a62f1873c78 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Serve: a later `--compact-at` replaces an earlier one, as for every other serve flag; `--compact-at auto --compact-at 0.5` used to keep automatic compaction because the fraction never cleared it (#2883). Thanks to @laudney. + - GLM-5.3-Flash on CUDA: `TF_GLM_CONTEXT_CEILING` (opt-in) lets GLM open windows past 1,048,576 tokens, read once when the model opens; with it unset nothing changes, and cold prompts and decode on two and three GB10s read level with main (#2765, refs #2663). Thanks to @Bizuayeu. - Metal on M5: prompt lane projections at every width (8-, 6-, 5-, 3- and 2-bit and the oQ mixes) take a chunk's 16-row blocks in one dispatch, as 4-bit already did; Qwen3.8-27B MLX 8-bit cold prompts are about 33% faster on an M5 Ultra (769 to 1,029 tok/s at 2K and 8K) with identical tokens, and decode is unchanged (#2851, for #2818). From 539a3fe5f2c1d95b9813a00e9cedbd33409dff16 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Wed, 11 Nov 2026 01:23:44 -0800 Subject: [PATCH 25/78] test: align Qwen27 reply history with newest-state retention Pin eviction of redundant own history before another conversation's newest state under pressure. Retain prompt-born provenance and prove a same-prompt regenerate restores history when all three states fit. Refs #2881 --- zig/src/native/qwen27_reply_cache_test.zig | 37 ++++++++++++++++++++++ zig/src/native/qwen27_served.zig | 24 ++------------ 2 files changed, 40 insertions(+), 21 deletions(-) diff --git a/zig/src/native/qwen27_reply_cache_test.zig b/zig/src/native/qwen27_reply_cache_test.zig index 91b8d078dbc..50d2be0d67c 100644 --- a/zig/src/native/qwen27_reply_cache_test.zig +++ b/zig/src/native/qwen27_reply_cache_test.zig @@ -9,6 +9,43 @@ const Metal = @import("qwen27_lanes.zig").Metal; const t = std.testing; const a = t.allocator; +/// Reply ends keep their request's provenance; spare budget also retains the pre-reply regenerate state. +pub fn replyBorn(keepEnd: anytype, snapshots: pc.Snapshots, live: *usize) !void { + const at: u32 = 5000; + for ([_]bool{ false, true }) |roomy| { + const rows: u64 = 3 * at + @as(u64, if (roomy) 107 else 100); + var store = pc.Store.init(a, snapshots, .{ .warm = .on }, rows << 16); + defer store.deinit(); + var other: [at + 7]u32 = undefined; + for (&other, 0..) |*word, i| word.* = @intCast(i); + try t.expect(store.keep(&other, at, null, &.{}, &.{})); + var tokens: [at + 107]u32 = undefined; + for (&tokens, 0..) |*word, i| word.* = @intCast(1_000_000 + i); + const prompt = tokens[0 .. at + 7]; + try t.expect(store.keep(prompt, at, null, &.{}, &.{})); + var s = try lanes.Stream.init(a, .{ .id = "reply", .prompt = prompt, .max_new = 100 }); + defer s.deinit(a); + try s.context.appendSlice(a, tokens[at + 7 ..]); + keepEnd(&store, a, &s, tokens.len); + const newest_other = store.find(&other, &.{}, &.{}) orelse return error.NewestConversationEvicted; + try t.expectEqual(at, newest_other.at); + const reply = store.entries.items[store.entries.items.len - 1]; + try t.expectEqual(@as(u32, tokens.len), reply.at); + try t.expectEqual(@as(u32, at + 7), reply.born); + try t.expectEqual(@as(usize, if (roomy) 3 else 2), store.entries.items.len); + try t.expectEqual(store.entries.items.len, live.*); + if (roomy) { + const history = store.find(prompt, &.{}, &.{}) orelse return error.HistoryEvicted; + try t.expectEqual(at, history.at); + const plan = try store.begin(a, prompt, at, &.{}, &.{}, null, &.{}); + defer a.free(plan.marks); + try t.expectEqual(at, plan.from); + try t.expectEqual(@as(u64, 1), store.counts.hits); + } else try t.expect(store.find(prompt, &.{}, &.{}) == null); + } + try t.expectEqual(@as(usize, 0), live.*); +} + const Family = struct { at: u32 = 0, sum: u64 = 0, diff --git a/zig/src/native/qwen27_served.zig b/zig/src/native/qwen27_served.zig index 05c40ddba7f..2e24131464e 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -265,28 +265,10 @@ test "--prompt-cache-gib 40: a later 192k-token slot takes its room back from ke try std.testing.expectEqual(@as(usize, 1), fake.live); } -test "a reply's end is born with its request's prompt: the history mark that request kept outlives the reply's keep" { - const gpa = std.testing.allocator; - const at: u32 = 5000; // a history mark at 5000 tokens, then a 7-token generation prompt and a 100-token reply +test "a reply's end keeps its prompt's born: newest conversations precede redundant history under pressure" { var fake: FakeStates = .{}; - var store = pc.Store.init(gpa, .{ .ptr = &fake, .vtable = &.{ .bytes = FakeStates.bytes, .save = FakeStates.save, .restore = FakeStates.restore, .drop = FakeStates.drop } }, .{ .warm = .on }, (3 * at + 100) << 16); - defer store.deinit(); - var other: [at + 7]u32 = undefined; // another conversation's prompt, its history kept first - for (&other, 0..) |*t, i| t.* = @intCast(i); - try std.testing.expect(store.keep(&other, at, null, &.{}, &.{})); - var tokens: [at + 107]u32 = undefined; - for (&tokens, 0..) |*t, i| t.* = @intCast(1_000_000 + i); - const prompt = tokens[0 .. at + 7]; - try std.testing.expect(store.keep(prompt, at, null, &.{}, &.{})); - var s = try lanes.Stream.init(gpa, .{ .id = "reply", .prompt = prompt, .max_new = 100 }); - defer s.deinit(gpa); - try s.context.appendSlice(gpa, tokens[at + 7 ..]); - keepEnd(&store, gpa, &s, tokens.len); // three states pass the budget: the least recently used goes - try std.testing.expect(store.find(&other, &.{}, &.{}) == null); - const history = store.find(prompt, &.{}, &.{}) orelse return error.HistoryEvicted; // a regenerate resumes it - try std.testing.expectEqual(at, history.at); - try std.testing.expectEqual(@as(u32, at + 7), store.entries.items[1].born); - try std.testing.expectEqual(@as(usize, 2), fake.live); + const snapshots: pc.Snapshots = .{ .ptr = &fake, .vtable = &.{ .bytes = FakeStates.bytes, .save = FakeStates.save, .restore = FakeStates.restore, .drop = FakeStates.drop } }; + try @import("qwen27_reply_cache_test.zig").replyBorn(keepEnd, snapshots, &fake.live); } /// Memory the slots past the ready ones take on first use: their K/V rows and taps rings. From 9597efdb6dca202986c5196766bf113526782a9c Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 12:14:47 +0000 Subject: [PATCH 26/78] changelog: #2889 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index a62f1873c78..7c6e70d4bfd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Tests: `qwen27_served`'s reply-history test follows #2699's rule (each conversation's newest state is kept first under pressure), and a new case checks that a regenerate restores at the history mark when the cache has room, so `zig build test-native` passes on main again (#2889, fixes #2881). + - Serve: a later `--compact-at` replaces an earlier one, as for every other serve flag; `--compact-at auto --compact-at 0.5` used to keep automatic compaction because the fraction never cleared it (#2883). Thanks to @laudney. - GLM-5.3-Flash on CUDA: `TF_GLM_CONTEXT_CEILING` (opt-in) lets GLM open windows past 1,048,576 tokens, read once when the model opens; with it unset nothing changes, and cold prompts and decode on two and three GB10s read level with main (#2765, refs #2663). Thanks to @Bizuayeu. From 833c83b80c6c925a9a16c0098c9bd219504349c5 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Wed, 11 Nov 2026 01:12:55 -0800 Subject: [PATCH 27/78] ci: run Metal engine host tests in the macOS nightly Give test-native a timed shard with the Metal toolchain. Keep it nightly until the CI runner establishes its wall time. Refs #2880 --- .github/workflows/no-gpu.yml | 20 +++++++++++++++++++- zig/build/host_ci.zig | 3 +++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/.github/workflows/no-gpu.yml b/.github/workflows/no-gpu.yml index 26b34ab3ef3..6960dc59073 100644 --- a/.github/workflows/no-gpu.yml +++ b/.github/workflows/no-gpu.yml @@ -335,7 +335,7 @@ jobs: strategy: fail-fast: false matrix: - suite: [fast, adapters, deepseek, lanes, deferred] + suite: [fast, adapters, native, deepseek, lanes, deferred] steps: - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0 with: @@ -344,11 +344,29 @@ jobs: with: version: 0.17.0 use-cache: false + - name: Metal toolchain for native engine host checks + if: matrix.suite == 'native' + run: xcrun -sdk macosx metal -v || { xcodebuild -downloadComponent MetalToolchain && xcrun -sdk macosx metal -v; } - name: Host checks and aggregate wall time shell: bash env: SUITE: ${{ matrix.suite }} run: | printf 'Runner class: %s %s\n' "$RUNNER_OS" "$RUNNER_ARCH" + if [ "$SUITE" = native ]; then + printf 'Native source: %s\n' "$GITHUB_SHA" + status=0 + for phase in initial warm; do + TIMEFORMAT="test-native $phase wall seconds: %R" + if { time zig build test-host-ci-macos-native --test-timeout 30m --summary all; }; then + printf 'test-native %s exit: 0\n' "$phase" + else + code=$? + printf 'test-native %s exit: %s\n' "$phase" "$code" + status=$code + fi + done + exit "$status" + fi TIMEFORMAT='Host checks wall seconds: %R' time zig build "test-host-ci-macos-$SUITE" --test-timeout 30m diff --git a/zig/build/host_ci.zig b/zig/build/host_ci.zig index 85206a8dd76..feaa56d56f0 100644 --- a/zig/build/host_ci.zig +++ b/zig/build/host_ci.zig @@ -14,6 +14,7 @@ pub fn steps(b: *std.Build) void { .macos => { wire(b, "test-host-ci-macos-fast", "Nightly host checks without a GPU or shader compiler", &macos_fast); wire(b, "test-host-ci-macos-adapters", "Nightly host checks without a GPU or shader compiler", &macos_adapters); + wire(b, "test-host-ci-macos-native", "Nightly Metal engine host checks with the Metal toolchain", &macos_native); wire(b, "test-host-ci-macos-deepseek", "Nightly host checks without a GPU or shader compiler", &macos_deepseek); wire(b, "test-host-ci-macos-lanes", "Nightly host checks without a GPU or shader compiler", &macos_lanes); wire(b, "test-host-ci-macos-deferred", "Nightly host checks without a GPU or shader compiler", &macos_deferred); @@ -289,6 +290,8 @@ const macos_adapters = [_][]const u8{ "test-verbs", }; +const macos_native = [_][]const u8{"test-native"}; + const macos_deepseek = [_][]const u8{ "test-dsv41-device-host", "test-dsv41-draft", From 9dbd140386d1e543db20637e0060878db769a769 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 12:14:49 +0000 Subject: [PATCH 28/78] changelog: #2891 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7c6e70d4bfd..4b01c602295 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- CI: the macOS nightly runs `zig build test-native`, the Metal engine host tests, installing the Metal toolchain when the runner lacks it and printing its cold and warm wall time (#2891, fixes #2880). + - Tests: `qwen27_served`'s reply-history test follows #2699's rule (each conversation's newest state is kept first under pressure), and a new case checks that a regenerate restores at the history mark when the cache has room, so `zig build test-native` passes on main again (#2889, fixes #2881). - Serve: a later `--compact-at` replaces an earlier one, as for every other serve flag; `--compact-at auto --compact-at 0.5` used to keep automatic compaction because the fraction never cleared it (#2883). Thanks to @laudney. From 5421def6f4a1d54701b39d078d7bfdb9b9275fe6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 11 Nov 2026 07:59:25 +0100 Subject: [PATCH 29/78] core: one block-FP8 reader and exact scale widening; Flash-Next MTP experts, Drama 3 and Kolibri-1 read through it --- zig/build/checkpoint_io_host.zig | 2 + zig/src/core/quant/fp8.zig | 54 +++++++ zig/src/core/quant/fp8_block_test.zig | 140 ++++++++++++++++++ zig/src/families/drama3/cuda/weights.zig | 106 +++++++++++-- zig/src/families/drama3/drama3.zig | 1 + .../flashnext/cuda_fp8_experts_load.zig | 14 +- zig/src/families/kolibri1/names.zig | 13 +- zig/src/families/kolibri1/pack.zig | 13 +- 8 files changed, 300 insertions(+), 43 deletions(-) create mode 100644 zig/src/core/quant/fp8_block_test.zig diff --git a/zig/build/checkpoint_io_host.zig b/zig/build/checkpoint_io_host.zig index 77767ffa710..583d42f9b1e 100644 --- a/zig/build/checkpoint_io_host.zig +++ b/zig/build/checkpoint_io_host.zig @@ -11,6 +11,8 @@ pub fn steps(b: *std.Build, target: std.Build.ResolvedTarget) void { step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = core, .filters = &.{"checkpoint admission:"} })).step); if (mode == .debug) b.step("test-modelopt-part-host", "The shared ModelOpt part reader (core/modelopt_part.zig) without a GPU") .dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = core, .filters = &.{"part:"} })).step); + if (mode == .debug) b.step("test-fp8-block-host", "The shared block-FP8 reader and exact scale widening (core/quant/fp8.zig)") + .dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = core, .filters = &.{"fp8 block:"} })).step); const qwen = b.createModule(.{ .root_source_file = b.path("zig/src/families/qwen3_5/loader.zig"), .target = target, .optimize = mode, .link_libc = true, .imports = &.{.{ .name = "core", .module = core }} }); step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = qwen, .filters = &.{"Qwen shard integrity:"} })).step); } diff --git a/zig/src/core/quant/fp8.zig b/zig/src/core/quant/fp8.zig index e95e38fbb43..3d400788b39 100644 --- a/zig/src/core/quant/fp8.zig +++ b/zig/src/core/quant/fp8.zig @@ -4,6 +4,7 @@ const types = @import("types.zig"); const convert = @import("convert.zig"); const Tensor = types.Tensor; +const Why = @import("../modelopt.zig").Why; /// Rows and columns one scale_inv entry covers. pub const block = 128; @@ -66,6 +67,59 @@ pub fn reference(h: Host, out: []f64) error{UnexpectedTensor}!void { }; } +/// The tensor beside `{module}.weight` that holds a block-FP8 projection's scales. +pub const scale_tail = "weight_scale_inv"; + +/// Scale blocks along a dimension of `n`: a last partial block takes a scale of its own. +pub fn blocks(n: usize) usize { + return (n + block - 1) / block; +} + +/// `{module}.weight`, e4m3 [n, k], from any store with `get(name) !Tensor`; a refusal named in `why`, else logged. +pub fn weight(store: anytype, module: []const u8, n: usize, k: usize, why: ?*Why) !Tensor { + return checked(store, module, "weight", .f8_e4m3, &.{ n, k }, why); +} + +/// `{module}.weight_scale_inv` of `dtype` [ceil(n/128), ceil(k/128)]; each family names the dtype it stores. +pub fn scales(store: anytype, module: []const u8, n: usize, k: usize, dtype: types.DType, why: ?*Why) !Tensor { + if (!convert.tableDtype(dtype)) return error.UnexpectedTensor; + return checked(store, module, scale_tail, dtype, &.{ blocks(n), blocks(k) }, why); +} + +/// A block-FP8 projection [n, k] as stored: views into the store's bytes, nothing copied. +pub fn read(store: anytype, module: []const u8, n: usize, k: usize, dtype: types.DType, why: ?*Why) !Host { + return .{ .weight = try weight(store, module, n, k, why), .scale_inv = try scales(store, module, n, k, dtype, why) }; +} + +fn checked(store: anytype, module: []const u8, tail: []const u8, dtype: types.DType, shape: []const usize, why: ?*Why) !Tensor { + var buf: [320]u8 = undefined; + const name = try std.fmt.bufPrint(&buf, "{s}.{s}", .{ module, tail }); + const t = try store.get(name); + const size = std.math.mul(usize, shape[0], shape[1]) catch return error.UnexpectedTensor; + if (t.is(dtype, shape) and t.bytes.len == size * dtype.size()) return t; + var own: Why = .{}; + const w = why orelse &own; + w.set("{s}: {t} {any}, expected {t} {any}", .{ name, t.dtype, t.shape[0..t.rank], dtype, shape }); + if (why == null) std.log.err("{s}", .{w.text()}); + return error.UnexpectedTensor; +} + +/// `scale_inv` as f32 into `out`, exactly (f32 copied bit for bit); `finite` refuses NaN and infinity. +pub fn widen(scale_inv: Tensor, out: []f32, finite: bool) error{ UnexpectedTensor, UnsupportedQuantization }!void { + const dtype = scale_inv.dtype; + if (!convert.tableDtype(dtype) or out.len != scale_inv.numel() or scale_inv.bytes.len != out.len * dtype.size()) + return error.UnexpectedTensor; + if (dtype == .f32) @memcpy(std.mem.sliceAsBytes(out), scale_inv.bytes); // NaN payloads too + for (out, 0..) |*v, i| { + if (dtype != .f32) v.* = convert.load(dtype, scale_inv.bytes, i); + if (finite and !std.math.isFinite(v.*)) return error.UnsupportedQuantization; + } +} + +test { + _ = @import("fp8_block_test.zig"); +} + test "e4m3 codes and their rounding" { try std.testing.expectEqual(@as(f64, 1), e4m3(0x38)); try std.testing.expectEqual(@as(f64, -2), e4m3(0xc0)); diff --git a/zig/src/core/quant/fp8_block_test.zig b/zig/src/core/quant/fp8_block_test.zig new file mode 100644 index 00000000000..6cd7ba82ff1 --- /dev/null +++ b/zig/src/core/quant/fp8_block_test.zig @@ -0,0 +1,140 @@ +//! fp8 block: the shared block-FP8 reader (fp8.read) and its exact scale widening (fp8.widen), on the host. +const std = @import("std"); +const fp8 = @import("fp8.zig"); +const types = @import("types.zig"); +const Why = @import("../modelopt.zig").Why; + +const Tensor = types.Tensor; +const testing = std.testing; + +fn tensor(dtype: types.DType, shape: []const usize, bytes: []const u8) Tensor { + var t: Tensor = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .bytes = bytes }; + @memcpy(t.shape[0..shape.len], shape); + return t; +} + +/// A store of named tensors, as safetensors maps, Drama 3's shards and the Flash-Next weight store answer `get`. +const Store = struct { + names: []const []const u8, + tensors: []const Tensor, + gets: usize = 0, + + pub fn get(s: *Store, name: []const u8) !Tensor { + s.gets += 1; + for (s.names, s.tensors) |n, t| if (std.mem.eql(u8, n, name)) return t; + return error.MissingTensor; + } +}; + +test "fp8 block: every bf16 bit pattern widens to its f32 value exactly, NaN and infinity kept without `finite`" { + var bits: [1 << 16]u16 = undefined; + for (&bits, 0..) |*b, i| b.* = @intCast(i); + var out: [1 << 16]f32 = undefined; + try fp8.widen(tensor(.bf16, &.{ 256, 256 }, std.mem.sliceAsBytes(&bits)), &out, false); + for (bits, out) |b, v| try testing.expectEqual(@as(u32, b) << 16, @as(u32, @bitCast(v))); +} + +test "fp8 block: f32 scales copy bit for bit, NaN payloads included" { + var prng = std.Random.DefaultPrng.init(0xf8b1); + var words: [4096]u32 = undefined; + prng.random().bytes(std.mem.sliceAsBytes(&words)); + words[0] = 0x7f800001; // a signalling NaN with a payload + words[1] = 0xffc12345; + var out: [4096]f32 = undefined; + try fp8.widen(tensor(.f32, &.{ 64, 64 }, std.mem.sliceAsBytes(&words)), &out, false); + try testing.expectEqualSlices(u32, &words, @ptrCast(&out)); +} + +test "fp8 block: every finite f16 bit pattern widens to its value exactly" { + var bits: [1 << 16]u16 = undefined; + for (&bits, 0..) |*b, i| b.* = @intCast(i); + var out: [1 << 16]f32 = undefined; + try fp8.widen(tensor(.f16, &.{1 << 16}, std.mem.sliceAsBytes(&bits)), &out, false); + for (bits, out) |b, v| { + const h: f16 = @bitCast(b); + if (std.math.isFinite(h)) try testing.expectEqual(@as(f64, @floatCast(h)), @as(f64, @floatCast(v))); + } +} + +test "fp8 block: `finite` refuses a NaN or infinite scale, and only then" { + for ([_]u16{ 0x7fc0, 0x7f80, 0xff80, 0x7f81 }) |bad| { + var bits = [_]u16{ 0x3f80, bad, 0x4000 }; + var out: [3]f32 = undefined; + try testing.expectError(error.UnsupportedQuantization, fp8.widen(tensor(.bf16, &.{ 1, 3 }, std.mem.sliceAsBytes(&bits)), &out, true)); + try fp8.widen(tensor(.bf16, &.{ 1, 3 }, std.mem.sliceAsBytes(&bits)), &out, false); + } + const ok = [_]f32{ 0.5, 1e-30, -2 }; + var out: [3]f32 = undefined; + try fp8.widen(tensor(.f32, &.{ 1, 3 }, std.mem.sliceAsBytes(&ok)), &out, true); + try testing.expectEqualSlices(f32, &ok, &out); +} + +test "fp8 block: widen refuses a wrong output length and a dtype the kernels do not read" { + const two = [_]f32{ 1, 2 }; + var out: [3]f32 = undefined; + try testing.expectError(error.UnexpectedTensor, fp8.widen(tensor(.f32, &.{ 1, 2 }, std.mem.sliceAsBytes(&two)), &out, false)); + const wide = [_]f64{1}; + try testing.expectError(error.UnexpectedTensor, fp8.widen(tensor(.f64, &.{ 1, 1 }, std.mem.sliceAsBytes(&wide)), out[0..1], false)); + try testing.expectError(error.UnexpectedTensor, fp8.widen(tensor(.f32, &.{ 1, 2 }, std.mem.sliceAsBytes(&two)[0..4]), out[0..2], false)); +} + +const codes: [130 * 260]u8 = @splat(0x38); +const inv: [2 * 3]f32 = .{ 1, 2, 3, 4, 5, 6 }; +const inv16: [2 * 3]u16 = .{ 1, 2, 3, 4, 5, 6 }; + +test "fp8 block: read returns the stored views of weight and weight_scale_inv, a partial last block counted" { + var s: Store = .{ + .names = &.{ "m.q.weight", "m.q.weight_scale_inv" }, + .tensors = &.{ tensor(.f8_e4m3, &.{ 130, 260 }, &codes), tensor(.f32, &.{ 2, 3 }, std.mem.sliceAsBytes(&inv)) }, + }; + var why: Why = .{}; + const h = try fp8.read(&s, "m.q", 130, 260, .f32, &why); + try testing.expectEqual(@intFromPtr(&codes), @intFromPtr(h.weight.bytes.ptr)); + try testing.expectEqual(@intFromPtr(&inv), @intFromPtr(h.scale_inv.bytes.ptr)); + try testing.expectEqual(@as(usize, 2), s.gets); + try testing.expectEqual(@as(usize, 0), why.len); + try testing.expectEqual(@as(usize, 2), fp8.blocks(130)); + try testing.expectEqual(@as(usize, 1), fp8.blocks(128)); +} + +test "fp8 block: read refuses by name a wrong dtype, a wrong shape and short bytes; a missing tensor is the store's" { + const Case = struct { w: Tensor, s: Tensor, n: usize, k: usize, dtype: types.DType, text: []const u8 }; + const w = tensor(.f8_e4m3, &.{ 130, 260 }, &codes); + const s = tensor(.f32, &.{ 2, 3 }, std.mem.sliceAsBytes(&inv)); + const cases = [_]Case{ + .{ .w = tensor(.u8, &.{ 130, 260 }, &codes), .s = s, .n = 130, .k = 260, .dtype = .f32, .text = "m.q.weight: u8 { 130, 260 }, expected f8_e4m3 { 130, 260 }" }, + .{ .w = w, .s = s, .n = 128, .k = 260, .dtype = .f32, .text = "m.q.weight: f8_e4m3 { 130, 260 }, expected f8_e4m3 { 128, 260 }" }, + .{ .w = w, .s = s, .n = 130, .k = 260, .dtype = .bf16, .text = "m.q.weight_scale_inv: f32 { 2, 3 }, expected bf16 { 2, 3 }" }, + .{ .w = w, .s = tensor(.f32, &.{ 3, 2 }, std.mem.sliceAsBytes(&inv)), .n = 130, .k = 260, .dtype = .f32, .text = "m.q.weight_scale_inv: f32 { 3, 2 }, expected f32 { 2, 3 }" }, + .{ .w = w, .s = tensor(.f32, &.{ 2, 3 }, std.mem.sliceAsBytes(&inv)[0..20]), .n = 130, .k = 260, .dtype = .f32, .text = "m.q.weight_scale_inv: f32 { 2, 3 }, expected f32 { 2, 3 }" }, + .{ .w = tensor(.f8_e4m3, &.{ 130, 260 }, codes[1..]), .s = s, .n = 130, .k = 260, .dtype = .f32, .text = "m.q.weight: f8_e4m3 { 130, 260 }, expected f8_e4m3 { 130, 260 }" }, + }; + for (cases) |c| { + var store: Store = .{ .names = &.{ "m.q.weight", "m.q.weight_scale_inv" }, .tensors = &.{ c.w, c.s } }; + var why: Why = .{}; + try testing.expectError(error.UnexpectedTensor, fp8.read(&store, "m.q", c.n, c.k, c.dtype, &why)); + try testing.expectEqualStrings(c.text, why.text()); + } + var none: Store = .{ .names = &.{"m.q.weight"}, .tensors = &.{w} }; + var why: Why = .{}; + try testing.expectError(error.MissingTensor, fp8.read(&none, "m.q", 130, 260, .f32, &why)); + // bf16 scales of the same blocks read as stored + var b: Store = .{ .names = &.{ "m.q.weight", "m.q.weight_scale_inv" }, .tensors = &.{ w, tensor(.bf16, &.{ 2, 3 }, std.mem.sliceAsBytes(&inv16)) } }; + _ = try fp8.read(&b, "m.q", 130, 260, .bf16, &why); +} + +test "fp8 block: read then widen gives reference()'s scales, block by block" { + var cs: [130 * 260]u8 = @splat(0x38); // 1.0 + cs[129 * 260 + 259] = 0x40; // 2.0 in the last, partial block + var s: Store = .{ .names = &.{ "e.weight", "e.weight_scale_inv" }, .tensors = &.{ tensor(.f8_e4m3, &.{ 130, 260 }, &cs), tensor(.f32, &.{ 2, 3 }, std.mem.sliceAsBytes(&inv)) } }; + var why: Why = .{}; + const h = try fp8.read(&s, "e", 130, 260, .f32, &why); + var widened: [6]f32 = undefined; + try fp8.widen(h.scale_inv, &widened, true); + var ref: [130 * 260]f64 = undefined; + try fp8.reference(h, &ref); + for (0..130) |r| for (0..260) |c| { + const want = fp8.e4m3(cs[r * 260 + c]) * widened[(r / fp8.block) * 3 + c / fp8.block]; + try testing.expectEqual(want, ref[r * 260 + c]); + }; +} diff --git a/zig/src/families/drama3/cuda/weights.zig b/zig/src/families/drama3/cuda/weights.zig index 7a0a2d941c3..0f7182bb991 100644 --- a/zig/src/families/drama3/cuda/weights.zig +++ b/zig/src/families/drama3/cuda/weights.zig @@ -7,6 +7,7 @@ const cfg = @import("../config.zig"); const shards = @import("../shards.zig"); const Shards = shards.Shards; +const fp8 = core.quant.fp8; pub const Gdn = struct { proj: u64, conv: u64, a_log: u64, dt_bias: u64, norm: u64, out: u64 }; pub const Attn = struct { qkv: u64, q_norm: u64, k_norm: u64, o: u64 }; @@ -114,6 +115,18 @@ fn fmt(buf: []u8, comptime f: []const u8, args: anytype) []const u8 { return std.fmt.bufPrint(buf, f, args) catch unreachable; } +/// An expert projection's e4m3 codes by the shared block-FP8 reader; a refused one stays Drama's BadExpert. +fn codes(sh: *const Shards, module: []const u8, n: usize, k: usize, why: ?*core.modelopt.Why) ![]const u8 { + const t = fp8.weight(sh, module, n, k, why) catch |e| return if (e == error.UnexpectedTensor) error.BadExpert else e; + return t.bytes; +} + +/// Its f32 weight_scale_inv blocks copied into `out` bit for bit; a refused one stays Drama's BadExpertScale. +fn scalesInto(sh: *const Shards, module: []const u8, n: usize, k: usize, out: []f32, why: ?*core.modelopt.Why) !void { + const t = fp8.scales(sh, module, n, k, .f32, why) catch |e| return if (e == error.UnexpectedTensor) error.BadExpertScale else e; + fp8.widen(t, out, false) catch return error.BadExpertScale; +} + const PackJob = struct { sh: *const Shards, layer: usize, @@ -135,28 +148,24 @@ const PackJob = struct { const gpa = std.heap.c_allocator; const n = cfg.expert_width; const k = cfg.hidden; - var nb: [160]u8 = undefined; var e = j.first; while (e < cfg.experts) : (e += j.stride) { const base = "model.language_model.layers.{d}.mlp.experts.{d}."; - const g = try j.sh.get(fmt(&nb, base ++ "gate_proj.weight", .{ j.layer, e })); - const u = try j.sh.get(fmt(&nb, base ++ "up_proj.weight", .{ j.layer, e })); - const dn = try j.sh.get(fmt(&nb, base ++ "down_proj.weight", .{ j.layer, e })); - if (!g.is(.f8_e4m3, &.{ n, k }) or !u.is(.f8_e4m3, &.{ n, k }) or !dn.is(.f8_e4m3, &.{ k, n })) return error.BadExpert; + var mb: [3][160]u8 = undefined; + const mods = [3][]const u8{ fmt(&mb[0], base ++ "gate_proj", .{ j.layer, e }), fmt(&mb[1], base ++ "up_proj", .{ j.layer, e }), fmt(&mb[2], base ++ "down_proj", .{ j.layer, e }) }; + const g = try codes(j.sh, mods[0], n, k, null); + const u = try codes(j.sh, mods[1], n, k, null); + const dn = try codes(j.sh, mods[2], k, n, null); const per_up = n * k * 2 / 4; - try cuda.experts.Fp8Experts.packGateUp(gpa, j.up[e * per_up ..][0..per_up], g.bytes, u.bytes, n, k); + try cuda.experts.Fp8Experts.packGateUp(gpa, j.up[e * per_up ..][0..per_up], g, u, n, k); const per_down = k * n / 4; - cuda.experts.Fp8Experts.packOne(j.down[e * per_down ..][0..per_down], dn.bytes, k, n); - const gs = try j.sh.get(fmt(&nb, base ++ "gate_proj.weight_scale_inv", .{ j.layer, e })); - const us = try j.sh.get(fmt(&nb, base ++ "up_proj.weight_scale_inv", .{ j.layer, e })); - const ds = try j.sh.get(fmt(&nb, base ++ "down_proj.weight_scale_inv", .{ j.layer, e })); - if (!gs.is(.f32, &.{ n / 128, k / 128 }) or !us.is(.f32, &.{ n / 128, k / 128 }) or !ds.is(.f32, &.{ k / 128, n / 128 })) return error.BadExpertScale; + cuda.experts.Fp8Experts.packOne(j.down[e * per_down ..][0..per_down], dn, k, n); const per_us = 2 * (n / 128) * (k / 128); - const usd = std.mem.sliceAsBytes(j.up_scale[e * per_us ..][0..per_us]); - @memcpy(usd[0..gs.bytes.len], gs.bytes); - @memcpy(usd[gs.bytes.len..][0..us.bytes.len], us.bytes); + const us = j.up_scale[e * per_us ..][0..per_us]; + try scalesInto(j.sh, mods[0], n, k, us[0 .. per_us / 2], null); + try scalesInto(j.sh, mods[1], n, k, us[per_us / 2 ..], null); const per_ds = (k / 128) * (n / 128); - @memcpy(std.mem.sliceAsBytes(j.down_scale[e * per_ds ..][0..per_ds]), ds.bytes); + try scalesInto(j.sh, mods[2], k, n, j.down_scale[e * per_ds ..][0..per_ds], null); for ([_][]const u8{ "gate_proj.weight", "up_proj.weight", "down_proj.weight" }) |suffix| { var nb2: [160]u8 = undefined; const name = std.fmt.bufPrint(&nb2, "model.language_model.layers.{d}.mlp.experts.{d}.{s}", .{ j.layer, e, suffix }) catch unreachable; @@ -303,3 +312,70 @@ pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, dir: []co if (opt.log) std.log.info("drama3: weights on the device: {d:.2} GB", .{@as(f64, @floatFromInt(w.bytes)) / 1e9}); return w; } + +test "drama experts: block-FP8 experts pack through the shared reader, f32 scales bit for bit, Drama's refusals kept" { + const a = std.testing.allocator; + const io = std.testing.io; + const n = cfg.expert_width; + const k = cfg.hidden; + const sb = (n / 128) * (k / 128); + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + var prng = std.Random.DefaultPrng.init(0xd3a3); + const pre = "model.language_model.layers.0.mlp.experts.1."; + const projs = [3][]const u8{ "gate_proj", "up_proj", "down_proj" }; + var codes_: [3][]u8 = undefined; + var inv: [3][sb]u32 = undefined; + var entries: [6]core.checkpoint_fixture.Entry = undefined; + var names: [6][96]u8 = undefined; + for (0..3) |p| { + codes_[p] = try a.alloc(u8, n * k); + prng.random().bytes(codes_[p]); + prng.random().bytes(std.mem.sliceAsBytes(&inv[p])); + const shape: []const usize = if (p == 2) &.{ k, n } else &.{ n, k }; + const blocks: []const usize = if (p == 2) &.{ k / 128, n / 128 } else &.{ n / 128, k / 128 }; + entries[2 * p] = .{ .name = fmt(&names[2 * p], pre ++ "{s}.weight", .{projs[p]}), .dtype = "F8_E4M3", .shape = shape, .bytes = codes_[p] }; + entries[2 * p + 1] = .{ .name = fmt(&names[2 * p + 1], pre ++ "{s}.weight_scale_inv", .{projs[p]}), .dtype = "F32", .shape = blocks, .bytes = std.mem.sliceAsBytes(&inv[p]) }; + } + defer for (codes_) |c| a.free(c); + inv[1][5] = 0x7f800123; // a NaN with a payload: Drama copies scales as stored, no finiteness rule + try core.checkpoint_fixture.write(a, io, tmp.dir, "model.safetensors", &entries); + var index: std.ArrayList(u8) = .empty; + defer index.deinit(a); + try index.appendSlice(a, "{\"weight_map\":{"); + for (entries, 0..) |e, i| try index.print(a, "{s}\"{s}\":\"model.safetensors\"", .{ if (i > 0) "," else "", e.name }); + try index.appendSlice(a, "}}"); + try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = index.items }); + const cwd = try std.process.currentPathAlloc(io, a); + defer a.free(cwd); + const dir = try std.fmt.allocPrint(a, "{s}/.zig-cache/tmp/{s}", .{ cwd, tmp.sub_path }); + defer a.free(dir); + var sh = try Shards.open(a, io, dir); + defer sh.close(io); + const per_up = n * k * 2 / 4; + const per_down = k * n / 4; + const up = try a.alloc(u32, 2 * per_up); + defer a.free(up); + const down = try a.alloc(u32, 2 * per_down); + defer a.free(down); + var up_scale: [2 * 2 * sb]f32 = undefined; + var down_scale: [2 * sb]f32 = undefined; + var job: PackJob = .{ .sh = &sh, .layer = 0, .first = 1, .stride = cfg.experts, .up = up, .down = down, .up_scale = &up_scale, .down_scale = &down_scale }; + try job.work(); + const want_up = try a.alloc(u32, per_up); + defer a.free(want_up); + try cuda.experts.Fp8Experts.packGateUp(a, want_up, codes_[0], codes_[1], n, k); + try std.testing.expectEqualSlices(u32, want_up, up[per_up..]); + const want_down = try a.alloc(u32, per_down); + defer a.free(want_down); + cuda.experts.Fp8Experts.packOne(want_down, codes_[2], k, n); + try std.testing.expectEqualSlices(u32, want_down, down[per_down..]); + try std.testing.expectEqualSlices(u32, &(inv[0] ++ inv[1]), @ptrCast(up_scale[2 * sb ..])); + try std.testing.expectEqualSlices(u32, &inv[2], @ptrCast(down_scale[sb..])); + // a weight or scale of the wrong shape keeps Drama's own error, the reason named by tensor + var why: core.modelopt.Why = .{}; + try std.testing.expectError(error.BadExpert, codes(&sh, pre ++ "down_proj", n, k, &why)); + try std.testing.expectEqualStrings(pre ++ "down_proj.weight: f8_e4m3 { 2048, 512 }, expected f8_e4m3 { 512, 2048 }", why.text()); + try std.testing.expectError(error.BadExpertScale, scalesInto(&sh, pre ++ "down_proj", n, k, down_scale[0..sb], &why)); + try std.testing.expectEqualStrings(pre ++ "down_proj.weight_scale_inv: f32 { 16, 4 }, expected f32 { 4, 16 }", why.text()); +} diff --git a/zig/src/families/drama3/drama3.zig b/zig/src/families/drama3/drama3.zig index 87a955bf15e..d929cb5b4af 100644 --- a/zig/src/families/drama3/drama3.zig +++ b/zig/src/families/drama3/drama3.zig @@ -43,4 +43,5 @@ test { test { _ = @import("cuda/owner_test.zig"); + _ = @import("cuda/weights.zig"); } diff --git a/zig/src/families/flashnext/cuda_fp8_experts_load.zig b/zig/src/families/flashnext/cuda_fp8_experts_load.zig index 3c79c3e52fe..0c4a0cf944c 100644 --- a/zig/src/families/flashnext/cuda_fp8_experts_load.zig +++ b/zig/src/families/flashnext/cuda_fp8_experts_load.zig @@ -1,8 +1,8 @@ //! The MTP head's block-FP8 routed experts: E4M3 weights, BF16 128x128 weight_scale_inv, packed for Fp8Experts (W8A16). const std = @import("std"); +const core = @import("core"); const cuda = @import("cuda"); const experts = @import("cuda_modelopt_experts.zig"); -const wt = @import("cuda_prompt_wt.zig"); const fx = cuda.experts.Fp8Experts; const blk = fx.block; @@ -78,15 +78,9 @@ pub fn packExpert(a: std.mem.Allocator, store: anytype, prefix: []const u8, i: u /// One projection's E4M3 weight [n, k]; its BF16 weight_scale_inv [n/128, k/128] widened (exactly) into `scale`. fn part(store: anytype, prefix: []const u8, i: usize, name: []const u8, n: usize, k: usize, scale: []f32) ![]const u8 { var buf: [512]u8 = undefined; - const w = try store.get(try std.fmt.bufPrint(&buf, "{s}.experts.{d}.{s}.weight", .{ prefix, i, name })); - if (!w.is(.f8_e4m3, &.{ n, k }) or w.bytes.len != n * k) return error.UnexpectedTensor; - const s = try store.get(try std.fmt.bufPrint(&buf, "{s}.experts.{d}.{s}.weight_scale_inv", .{ prefix, i, name })); - if (!s.is(.bf16, &.{ n / blk, k / blk }) or s.bytes.len != 2 * scale.len) return error.UnexpectedTensor; - for (scale, 0..) |*v, j| { - v.* = wt.promote(std.mem.readInt(u16, s.bytes[2 * j ..][0..2], .little)); - if (!std.math.isFinite(v.*)) return error.UnsupportedQuantization; - } - return w.bytes; + const h = try core.quant.fp8.read(store, try std.fmt.bufPrint(&buf, "{s}.experts.{d}.{s}", .{ prefix, i, name }), n, k, .bf16, null); + try core.quant.fp8.widen(h.scale_inv, scale, true); + return h.weight.bytes; } test { diff --git a/zig/src/families/kolibri1/names.zig b/zig/src/families/kolibri1/names.zig index 361864d86fb..2f179a07895 100644 --- a/zig/src/families/kolibri1/names.zig +++ b/zig/src/families/kolibri1/names.zig @@ -6,7 +6,7 @@ const Config = @import("config.zig").Config; const Tensor = core.checkpoint.Tensor; const Checkpoint = core.checkpoint.Checkpoint; -pub const block = 128; // the checkpoint's scale block, both ways +pub const block = core.quant.fp8.block; // the checkpoint's scale block, both ways /// A block-FP8 projection as stored: e4m3 codes [n, k] and fp32 scale_inv [ceil(n/128), ceil(k/128)]. pub const Fp8 = struct { w: Tensor, s: Tensor }; @@ -42,10 +42,6 @@ pub const Map = struct { } }; -fn blocks(n: usize) usize { - return (n + block - 1) / block; -} - const Names = struct { ck: *Checkpoint, buf: [160]u8 = undefined, @@ -59,9 +55,8 @@ const Names = struct { } fn fp8(self: *Names, n: usize, k: usize, comptime fmt: []const u8, args: anytype) !Fp8 { - const w = try self.ck.expect(try self.name(fmt ++ ".weight", args), .f8_e4m3, &.{ n, k }); - const s = try self.ck.expect(try self.name(fmt ++ ".weight_scale_inv", args), .f32, &.{ blocks(n), blocks(k) }); - return .{ .w = w, .s = s }; + const h = try core.quant.fp8.read(self.ck, try self.name(fmt, args), n, k, .f32, null); + return .{ .w = h.weight, .s = h.scale_inv }; } fn expert(self: *Names, c: Config, comptime pre: []const u8, args: anytype) !Expert { @@ -130,7 +125,7 @@ fn writeTiny(a: std.mem.Allocator, tmp: testing.TmpDir) !void { const Add = struct { fn fp8(list: anytype, al: std.mem.Allocator, pre: []const u8, n: usize, k: usize) !void { try list.append(al, .{ .name = try std.fmt.allocPrint(al, "{s}.weight", .{pre}), .dtype = "F8_E4M3", .shape = try al.dupe(usize, &.{ n, k }) }); - try list.append(al, .{ .name = try std.fmt.allocPrint(al, "{s}.weight_scale_inv", .{pre}), .dtype = "F32", .shape = try al.dupe(usize, &.{ blocks(n), blocks(k) }) }); + try list.append(al, .{ .name = try std.fmt.allocPrint(al, "{s}.weight_scale_inv", .{pre}), .dtype = "F32", .shape = try al.dupe(usize, &.{ core.quant.fp8.blocks(n), core.quant.fp8.blocks(k) }) }); } fn bf(list: anytype, al: std.mem.Allocator, name: []const u8, shape: []const usize) !void { try list.append(al, .{ .name = name, .dtype = "BF16", .shape = try al.dupe(usize, shape) }); diff --git a/zig/src/families/kolibri1/pack.zig b/zig/src/families/kolibri1/pack.zig index 8c93570cf65..a51c0975fc5 100644 --- a/zig/src/families/kolibri1/pack.zig +++ b/zig/src/families/kolibri1/pack.zig @@ -49,11 +49,6 @@ pub fn widen(gpa: std.mem.Allocator, t: Tensor) ![]f32 { return out; } -/// fp32 values of a tensor whose bytes may sit at any offset of the mapped file. -fn floats(dst: []f32, t: Tensor) void { - for (dst, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u32, t.bytes[4 * i ..][0..4], .little)); -} - /// Row-stacked block-FP8 parts as one projection: codes and scales concatenated by rows, then packed. fn linear(gpa: std.mem.Allocator, parts: []const names.Fp8) !Linear { const k = parts[0].w.dim(1); @@ -73,7 +68,7 @@ fn linear(gpa: std.mem.Allocator, parts: []const names.Fp8) !Linear { for (parts) |p| { @memcpy(codes[at..][0..p.w.bytes.len], p.w.bytes); at += p.w.bytes.len; - floats(inv[sat..][0..p.s.numel()], p.s); + try core.quant.fp8.widen(p.s, inv[sat..][0..p.s.numel()], false); sat += p.s.numel(); } const npad = fp8.padded(n); @@ -98,9 +93,9 @@ fn experts(gpa: std.mem.Allocator, l: *Layer, ex: []const names.Expert, c: Confi for (ex, 0..) |e, i| { try fx8.packGateUp(gpa, l.up[i * one / 2 ..][0 .. one / 2], e.gate.w.bytes, e.up.w.bytes, ni, d); fx8.packOne(l.down[i * one / 4 ..][0 .. one / 4], e.down.w.bytes, d, ni); - floats(l.up_scale[2 * i * sb ..][0..sb], e.gate.s); - floats(l.up_scale[(2 * i + 1) * sb ..][0..sb], e.up.s); - floats(l.down_scale[i * sb ..][0..sb], e.down.s); + try core.quant.fp8.widen(e.gate.s, l.up_scale[2 * i * sb ..][0..sb], false); + try core.quant.fp8.widen(e.up.s, l.up_scale[(2 * i + 1) * sb ..][0..sb], false); + try core.quant.fp8.widen(e.down.s, l.down_scale[i * sb ..][0..sb], false); } } From e2887ac3ab1e0c4dabe3287f65819c31a162a394 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 11 Nov 2026 12:33:51 +0100 Subject: [PATCH 30/78] core fp8: a byte-length refusal names the byte counts; widen's comment says an f16 NaN stays a NaN --- zig/src/core/quant/fp8.zig | 6 ++++-- zig/src/core/quant/fp8_block_test.zig | 4 ++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/zig/src/core/quant/fp8.zig b/zig/src/core/quant/fp8.zig index 3d400788b39..f28169ddae9 100644 --- a/zig/src/core/quant/fp8.zig +++ b/zig/src/core/quant/fp8.zig @@ -99,12 +99,14 @@ fn checked(store: anytype, module: []const u8, tail: []const u8, dtype: types.DT if (t.is(dtype, shape) and t.bytes.len == size * dtype.size()) return t; var own: Why = .{}; const w = why orelse &own; - w.set("{s}: {t} {any}, expected {t} {any}", .{ name, t.dtype, t.shape[0..t.rank], dtype, shape }); + if (t.is(dtype, shape)) { + w.set("{s}: {d} bytes, {t} {any} needs {d}", .{ name, t.bytes.len, dtype, shape, size * dtype.size() }); + } else w.set("{s}: {t} {any}, expected {t} {any}", .{ name, t.dtype, t.shape[0..t.rank], dtype, shape }); if (why == null) std.log.err("{s}", .{w.text()}); return error.UnexpectedTensor; } -/// `scale_inv` as f32 into `out`, exactly (f32 copied bit for bit); `finite` refuses NaN and infinity. +/// `scale_inv` as f32 into `out`: exact (f32 bit for bit), an f16 NaN stays a NaN; `finite` refuses NaN and infinity. pub fn widen(scale_inv: Tensor, out: []f32, finite: bool) error{ UnexpectedTensor, UnsupportedQuantization }!void { const dtype = scale_inv.dtype; if (!convert.tableDtype(dtype) or out.len != scale_inv.numel() or scale_inv.bytes.len != out.len * dtype.size()) diff --git a/zig/src/core/quant/fp8_block_test.zig b/zig/src/core/quant/fp8_block_test.zig index 6cd7ba82ff1..9ea6e30f906 100644 --- a/zig/src/core/quant/fp8_block_test.zig +++ b/zig/src/core/quant/fp8_block_test.zig @@ -106,8 +106,8 @@ test "fp8 block: read refuses by name a wrong dtype, a wrong shape and short byt .{ .w = w, .s = s, .n = 128, .k = 260, .dtype = .f32, .text = "m.q.weight: f8_e4m3 { 130, 260 }, expected f8_e4m3 { 128, 260 }" }, .{ .w = w, .s = s, .n = 130, .k = 260, .dtype = .bf16, .text = "m.q.weight_scale_inv: f32 { 2, 3 }, expected bf16 { 2, 3 }" }, .{ .w = w, .s = tensor(.f32, &.{ 3, 2 }, std.mem.sliceAsBytes(&inv)), .n = 130, .k = 260, .dtype = .f32, .text = "m.q.weight_scale_inv: f32 { 3, 2 }, expected f32 { 2, 3 }" }, - .{ .w = w, .s = tensor(.f32, &.{ 2, 3 }, std.mem.sliceAsBytes(&inv)[0..20]), .n = 130, .k = 260, .dtype = .f32, .text = "m.q.weight_scale_inv: f32 { 2, 3 }, expected f32 { 2, 3 }" }, - .{ .w = tensor(.f8_e4m3, &.{ 130, 260 }, codes[1..]), .s = s, .n = 130, .k = 260, .dtype = .f32, .text = "m.q.weight: f8_e4m3 { 130, 260 }, expected f8_e4m3 { 130, 260 }" }, + .{ .w = w, .s = tensor(.f32, &.{ 2, 3 }, std.mem.sliceAsBytes(&inv)[0..20]), .n = 130, .k = 260, .dtype = .f32, .text = "m.q.weight_scale_inv: 20 bytes, f32 { 2, 3 } needs 24" }, + .{ .w = tensor(.f8_e4m3, &.{ 130, 260 }, codes[1..]), .s = s, .n = 130, .k = 260, .dtype = .f32, .text = "m.q.weight: 33799 bytes, f8_e4m3 { 130, 260 } needs 33800" }, }; for (cases) |c| { var store: Store = .{ .names = &.{ "m.q.weight", "m.q.weight_scale_inv" }, .tensors = &.{ c.w, c.s } }; From 6351f260801367223f32b7a7ee4114a8748fdf1d Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 12:14:50 +0000 Subject: [PATCH 31/78] changelog: #2892 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4b01c602295..e55d1906beb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Internal: one block-FP8 reader in `core/quant/fp8.zig` loads Flash Next's, Kolibri-1's and Drama 3's FP8 weights and their scales, widening bf16 scales by the same bit shift and copying f32 ones as before; every packed buffer is byte-identical, and only Flash Next refuses non-finite scales, as before (#2892, #2460 item 7). Thanks to @jschmied. + - CI: the macOS nightly runs `zig build test-native`, the Metal engine host tests, installing the Metal toolchain when the runner lacks it and printing its cold and warm wall time (#2891, fixes #2880). - Tests: `qwen27_served`'s reply-history test follows #2699's rule (each conversation's newest state is kept first under pressure), and a new case checks that a regenerate restores at the history mark when the cache has room, so `zig build test-native` passes on main again (#2889, fixes #2881). From 95eb4f9639781ccb62453d0e8527ca96f058e5c5 Mon Sep 17 00:00:00 2001 From: Abraham <119352736+AbrahamCete@users.noreply.github.com> Date: Tue, 10 Nov 2026 22:26:07 +0100 Subject: [PATCH 32/78] perf(hip): the routed prompt tile takes K-parallel passes for an item of few rows On gfx1151 a routed prompt holds about 80 experts of 8 rows or fewer at any length. On the 64-row matrix tile each of their blocks reads 128 columns of codes a stage, a page of weights every four columns, and the cost of those reads follows the pages a block touches. wmma_item_tile now runs gemm_kp_wave for an item of at most the launch's cut: a lane a column's group, a column's whole row of codes read at once, no barrier. gemm_kp_wave is a wave's passes of gemm_kp_tile, factored out so both tiles run them; the K-parallel tiles keep their statement order and their register and scratch use. A pass of the item tile takes 2 rows and 2 columns a lane, which stays in registers, and keeps its partials in the block's code stages, so the tile's LDS is unchanged. The cut is a card's: `few_rows` in the tuning row (8 on gfx1151's item64 row, none elsewhere), carried to the kernel in Routing.few by prefillLaunch. A kernel test runs one plan through both paths at every width and compares bytes; the contract runs the tile with the cut and a control whose passes stop after a wave's first. The tile's record holds ROCm 10.1's hipcc. The build now tracks tiles/matrix_item.hpp. Refs #2795 --- zig/build/hip.zig | 18 ++--- zig/kernels/hip/contract_controls.hip | 13 ++++ zig/kernels/hip/quant/mlx_tiles.hpp | 11 ++- zig/kernels/hip/tiles/gemm_kp.hpp | 58 ++++++++++------ zig/kernels/hip/tiles/matrix_item.hpp | 30 +++++++- zig/kernels/hip/tiles/plan.hpp | 1 + zig/src/core/contract/hip_records_item.zig | 4 +- zig/src/core/registry.zig | 6 ++ zig/src/core/tuning.zig | 2 + zig/src/hip/kernel_tests.zig | 1 + zig/src/hip/kernel_tests/affine/cases.zig | 2 + zig/src/hip/kernel_tests/affine/item_few.zig | 72 ++++++++++++++++++++ zig/src/hip/kernel_tests/affine/product.zig | 9 ++- zig/src/hip/launches.zig | 1 + zig/src/hip/launches/affine.zig | 9 ++- zig/src/hip/launches/matrix_item.zig | 9 ++- zig/src/hip/tuning/gfx1151.zon | 4 +- zig/tests/hip/contract_main.zig | 10 ++- zig/tests/hip/contract_wmma_item.zig | 38 ++++++++--- 19 files changed, 245 insertions(+), 53 deletions(-) create mode 100644 zig/src/hip/kernel_tests/affine/item_few.zig diff --git a/zig/build/hip.zig b/zig/build/hip.zig index 70c3f26e360..88c0f6fb363 100644 --- a/zig/build/hip.zig +++ b/zig/build/hip.zig @@ -518,15 +518,15 @@ const groups = [_]Source{ }; /// What the group sources include, so an edit to one rebuilds them. const headers = [_][]const u8{ - "attention/attention_fa.hip", "attention/prefill.hpp", "decode/pages.hpp", "decode/plan.hpp", "ops/attention.hpp", - "ops/common.hpp", "ops/draw.hpp", "ops/tp.hpp", "ops/elementwise.hpp", "ops/linear.hpp", - "ops/moe.hpp", "ops/norms.hpp", "ops/rope.hpp", "common/arch.hpp", "common/dot2.hpp", - "common/vec.hpp", "common/wmma.hpp", "quant/act.hpp", "quant/mlx.hpp", "quant/mlx_decoder.hpp", - "quant/mlx_pieces.hpp", "quant/mlx_tiles.hpp", "tiles/dot2.hpp", "tiles/epilogue.hpp", "tiles/gemm.hpp", - "tiles/gemm_kp.hpp", "tiles/matrix_gemm.hpp", "tiles/plan.hpp", "tiles/stream.hpp", "flashnext/streams.hpp", - "flashnext/ngram.hpp", "flashnext/indexer.hpp", "flashnext/sparse.hpp", "ops/nemotron_moe.hpp", "quant/bf16_decoder.hpp", - "quant/dense_tiles.hpp", "tiles/wmma12_gemm.hpp", "tiles/wmma12_gemm_k32_ref.hpp", "ops/columns.hpp", "columns_controls.hpp", - "ops/vision.hpp", "ops/dense_text.hpp", "ops/vl_position.hpp", "flashnext/sparse_lanes.hpp", + "attention/attention_fa.hip", "attention/prefill.hpp", "decode/pages.hpp", "decode/plan.hpp", "ops/attention.hpp", + "ops/common.hpp", "ops/draw.hpp", "ops/tp.hpp", "ops/elementwise.hpp", "ops/linear.hpp", + "ops/moe.hpp", "ops/norms.hpp", "ops/rope.hpp", "common/arch.hpp", "common/dot2.hpp", + "common/vec.hpp", "common/wmma.hpp", "quant/act.hpp", "quant/mlx.hpp", "quant/mlx_decoder.hpp", + "quant/mlx_pieces.hpp", "quant/mlx_tiles.hpp", "tiles/dot2.hpp", "tiles/epilogue.hpp", "tiles/gemm.hpp", + "tiles/gemm_kp.hpp", "tiles/matrix_gemm.hpp", "tiles/matrix_item.hpp", "tiles/plan.hpp", "tiles/stream.hpp", + "flashnext/streams.hpp", "flashnext/ngram.hpp", "flashnext/indexer.hpp", "flashnext/sparse.hpp", "ops/nemotron_moe.hpp", + "quant/bf16_decoder.hpp", "quant/dense_tiles.hpp", "tiles/wmma12_gemm.hpp", "tiles/wmma12_gemm_k32_ref.hpp", "ops/columns.hpp", + "columns_controls.hpp", "ops/vision.hpp", "ops/dense_text.hpp", "ops/vl_position.hpp", "flashnext/sparse_lanes.hpp", }; /// The flags of the kernels' first build: no contraction, wave32 on RDNA, C++20. const flags = [_][]const u8{ "-D__HIP_PLATFORM_AMD__=1", "-DUSE_ROCM=1", "-DHIPBLAS_V2", "-fPIC", "-DCUDA_HAS_FP16=1", "-DHIP_ENABLE_WARP_SYNC_BUILTINS=1", "-std=c++20", "-fno-gpu-rdc", "-mno-wavefrontsize64", "-ffp-contract=off" }; diff --git a/zig/kernels/hip/contract_controls.hip b/zig/kernels/hip/contract_controls.hip index 231f769c61c..327da668856 100644 --- a/zig/kernels/hip/contract_controls.hip +++ b/zig/kernels/hip/contract_controls.hip @@ -7,6 +7,7 @@ #include "attention/prefill.hpp" #include "quant/dense_tiles.hpp" +#include "quant/mlx_tiles.hpp" #include "tiles/wmma12_gemm_k32_ref.hpp" __device__ __forceinline__ float tf_ctl_x(const uint16_t* x, int64_t i) { @@ -196,6 +197,18 @@ tf_ctl_fa_paged_unmasked_bf16(const float* q, const hip_bfloat16* k, const hip_b 0, 0, 0, 0, table, count); } +// Gate 3: the routed item tile's few-row passes stop after a wave's first, leaving the block's other columns unwritten. +struct TfCtlItemFewOnePass { + using Dot = tf::rocm::DotBF16; + using Epi = tf::rocm::GemmOut; + static constexpr bool every_pass = false; +}; + +extern "C" __global__ void __launch_bounds__(256) tf_ctl_item_few_one_pass(tf::rocm::Affine a) { + using namespace tf::rocm; + wmma_item_tile, BF16Act, WmmaBF16, GemmOut, 4, TfCtlItemFewOnePass>(a); +} + // Gate 3: the big block loading its operands without the stores' swizzle, and without its K loop's barriers. struct TfCtlUnswizzled { static constexpr bool swizzled_loads = false, ordered = true; diff --git a/zig/kernels/hip/quant/mlx_tiles.hpp b/zig/kernels/hip/quant/mlx_tiles.hpp index a3f6ca68940..70bdd72d1e5 100644 --- a/zig/kernels/hip/quant/mlx_tiles.hpp +++ b/zig/kernels/hip/quant/mlx_tiles.hpp @@ -45,11 +45,18 @@ __global__ void __launch_bounds__(256) affine_wmma_skew(Affine a) { wmma_gemm_tile, BF16Act, WmmaBF16, GemmOut, true>(a); } -// ---- the same on RM x 16 rows a block, for a routed item of few rows ---- +// ---- the same on RM x 16 rows a block, for a routed item of few rows; K-parallel passes for one of very few ---- + +// The item tile's few-row passes as served: BF16 dot2 into the fp32 sums, every pass of a wave's columns. +struct ItemFew { + using Dot = DotBF16; + using Epi = GemmOut; + static constexpr bool every_pass = true; +}; template __global__ void __launch_bounds__(256) affine_wmma_item(Affine a) { - wmma_item_tile, BF16Act, WmmaBF16, GemmOut, RM>(a); + wmma_item_tile, BF16Act, WmmaBF16, GemmOut, RM, ItemFew>(a); } // ---- the decode stream tile ---- diff --git a/zig/kernels/hip/tiles/gemm_kp.hpp b/zig/kernels/hip/tiles/gemm_kp.hpp index d759c7da59f..f4b7da96a48 100644 --- a/zig/kernels/hip/tiles/gemm_kp.hpp +++ b/zig/kernels/hip/tiles/gemm_kp.hpp @@ -12,26 +12,23 @@ namespace rocm { // The lanes a column takes, a group each: 32, or 16 or 8 when it has few groups (a wave then takes more columns). __host__ __device__ inline int kp_lanes(int groups) { return groups > 16 ? 32 : groups > 8 ? 16 : 8; } -// CB columns a lane, R rows a pass, WAVES waves a block; passes span <= RB adjacent row blocks to share weights. -template -__device__ __forceinline__ void gemm_kp_tile(typename Dec::Args& a) { - const int blocks = (a.m + R - 1) / R < RB ? (a.m + R - 1) / R : RB; // of the launch: a routed item has its own rows - if (!Dec::take_item(a, blockIdx.z)) return; +// A lane's row of each of a wave's LDS tables, what `lds(lane)` names: R * CB dots, R sums of x, CB scales, CB biases. +struct KpRows { + float* part; + float* spart; + float* scl; + float* bil; +}; + +// A wave's passes over its columns from `colw`, CB a lane and R rows a pass: rows from `first`, every `step` when LOOP. +template +__device__ __forceinline__ void gemm_kp_wave(typename Dec::Args& a, int colw, int first, int step, Lds lds) { constexpr int NO = R * CB; // outputs of a lane set - constexpr int PS = NO + 1; constexpr int OPL = (4 * NO + 31) / 32; // outputs a lane folds, at most const int groups = a.k / a.group; const int gp = kp_lanes(groups); // the lanes of a set; a wave has 32 / gp sets const int sets = 32 / gp; - const int colw = ((blockIdx.x / blocks) * WAVES + (threadIdx.x >> 5)) * CB * sets; - if (colw >= a.n) return; - __shared__ float part[WAVES][32][PS]; - __shared__ float spart[WAVES][32][R + 1]; - __shared__ float scl[WAVES][32][CB + 1]; - __shared__ float bil[WAVES][32][CB + 1]; - const int lane = threadIdx.x & 31; - const int wave = threadIdx.x >> 5; const int per = a.group / 32; const int rounds = (groups + gp - 1) / gp; const int gi = lane & (gp - 1); @@ -128,14 +125,14 @@ __device__ __forceinline__ void gemm_kp_tile(typename Dec::Args& a) { // this group's results to LDS, then the round's groups in order #pragma unroll for (int r = 0; r < R; ++r) { - spart[wave][lane][r] = sx[r]; + lds(lane).spart[r] = sx[r]; #pragma unroll - for (int c = 0; c < CB; ++c) part[wave][lane][r * CB + c] = d[r][c]; + for (int c = 0; c < CB; ++c) lds(lane).part[r * CB + c] = d[r][c]; } #pragma unroll for (int c = 0; c < CB; ++c) { - scl[wave][lane][c] = Dec::scale_value(a, term[c]); - bil[wave][lane][c] = Dec::bias_value(a, term[c]); + lds(lane).scl[c] = Dec::scale_value(a, term[c]); + lds(lane).bil[c] = Dec::bias_value(a, term[c]); } if (round + 1 < rounds) tables(g + gp); __syncwarp(); @@ -150,8 +147,8 @@ __device__ __forceinline__ void gemm_kp_tile(typename Dec::Args& a) { const int c = o % CB; float v = acc[t]; for (int gg = 0; gg < count; ++gg) { - v = Dec::fold_scale(v, part[wave][base + gg][o], scl[wave][base + gg][c]); - v = Dec::fold_bias(v, spart[wave][base + gg][r], bil[wave][base + gg][c]); + v = Dec::fold_scale(v, lds(base + gg).part[o], lds(base + gg).scl[c]); + v = Dec::fold_bias(v, lds(base + gg).spart[r], lds(base + gg).bil[c]); } acc[t] = v; } @@ -168,11 +165,28 @@ __device__ __forceinline__ void gemm_kp_tile(typename Dec::Args& a) { } }; if constexpr (LOOP) { - for (int row0 = (blockIdx.x % blocks) * R; row0 < a.m; row0 += blocks * R) pass(row0); + for (int row0 = first; row0 < a.m; row0 += step) pass(row0); } else { - pass((blockIdx.x % blocks) * R); + pass(first); } } +// CB columns a lane, R rows a pass, WAVES waves a block; passes span <= RB adjacent row blocks to share weights. +template +__device__ __forceinline__ void gemm_kp_tile(typename Dec::Args& a) { + const int blocks = (a.m + R - 1) / R < RB ? (a.m + R - 1) / R : RB; // of the launch: a routed item has its own rows + if (!Dec::take_item(a, blockIdx.z)) return; + const int sets = 32 / kp_lanes(a.k / a.group); + const int colw = ((blockIdx.x / blocks) * WAVES + (threadIdx.x >> 5)) * CB * sets; + if (colw >= a.n) return; + __shared__ float part[WAVES][32][R * CB + 1]; + __shared__ float spart[WAVES][32][R + 1]; + __shared__ float scl[WAVES][32][CB + 1]; + __shared__ float bil[WAVES][32][CB + 1]; + const int wave = threadIdx.x >> 5; + const auto lds = [wave](int lane) { return KpRows{part[wave][lane], spart[wave][lane], scl[wave][lane], bil[wave][lane]}; }; + gemm_kp_wave(a, colw, (blockIdx.x % blocks) * R, blocks * R, lds); +} + } // namespace rocm } // namespace tf diff --git a/zig/kernels/hip/tiles/matrix_item.hpp b/zig/kernels/hip/tiles/matrix_item.hpp index 06195bb373e..734ccf5a056 100644 --- a/zig/kernels/hip/tiles/matrix_item.hpp +++ b/zig/kernels/hip/tiles/matrix_item.hpp @@ -4,13 +4,19 @@ #include "common/wmma.hpp" #include "tiles/gemm.hpp" +#include "tiles/gemm_kp.hpp" #include "tiles/plan.hpp" namespace tf { namespace rocm { +// The most rows a launch's cut sends to K-parallel passes, the rows a pass takes and the columns a lane takes in it. +constexpr int kItemFewRows = 8; +constexpr int kItemFewPass = 2; +constexpr int kItemFewCols = 2; + // Same staging, stage order and per-element sums as wmma_gemm_tile, so each output equals that tile's byte for byte. -template +template __device__ __forceinline__ void wmma_item_tile(typename Dec::Args& a) { #if TF_DEVICE_WMMA_GFX11 if (!Dec::take_item(a, blockIdx.z)) return; @@ -27,6 +33,28 @@ __device__ __forceinline__ void wmma_item_tile(typename Dec::Args& a) { const int tid = threadIdx.x; const int n0 = bx * kGemmN; + + // An item within the launch's cut takes K-parallel passes with Few's Dot and epilogue: same bytes, no matrix core. + if (a.m <= (a.route.few < kItemFewRows ? a.route.few : kItemFewRows)) { + constexpr int CB = kItemFewCols; + constexpr int R = kItemFewPass; + // a pass's partials, a row a lane: each wave's own slice of the two code stages, which this path leaves unused + constexpr int kLane = (R * CB + 1) + (R + 1) + 2 * (CB + 1); + static_assert(8 * 32 * kLane <= 2 * kGemmN * kLdw, "the few-row passes' partials fit the tile's code stages"); + const int wave = __builtin_amdgcn_readfirstlane(tid >> 5); + const auto lds = [wave](int lane) { + float* const row = reinterpret_cast(&ws[0][0]) + (wave * 32 + lane) * kLane; + return KpRows{row, row + R * CB + 1, row + R * CB + 1 + R + 1, row + R * CB + 1 + R + 1 + CB + 1}; + }; + // The 8 waves walk the block's columns side by side, CB a lane set a pass, R rows of the item a pass. + const int span = CB * (32 / kp_lanes(a.k / a.group)); + const int end = n0 + kGemmN < a.n ? n0 + kGemmN : a.n; + for (int colw = n0 + wave * span; colw < end; colw += 8 * span) { + gemm_kp_wave(a, colw, 0, R, lds); + if constexpr (!Few::every_pass) break; // the contract's control stops after a wave's first + } + return; + } const int words_row = Dec::template row_words(a); const int groups = a.k / a.group; const int stages = a.k / kGemmK; diff --git a/zig/kernels/hip/tiles/plan.hpp b/zig/kernels/hip/tiles/plan.hpp index 6cd355ba754..260d14b3b1e 100644 --- a/zig/kernels/hip/tiles/plan.hpp +++ b/zig/kernels/hip/tiles/plan.hpp @@ -10,6 +10,7 @@ struct Routing { const int* members = nullptr; int x_div = 1; int first = 0; // set by the kernel from its item + int few = 0; // an item of at most this many rows takes its tile's K-parallel passes, where the tile has them }; // Row r of this block's x and out: the plain matrix's r, or the routed item's r-th pair. diff --git a/zig/src/core/contract/hip_records_item.zig b/zig/src/core/contract/hip_records_item.zig index e7d154bb843..265209c83c7 100644 --- a/zig/src/core/contract/hip_records_item.zig +++ b/zig/src/core/contract/hip_records_item.zig @@ -18,8 +18,8 @@ pub const records = [1]Record{ .launch = itemTile(4), .usage = .{ .{ .vgprs = 1, .sgprs = 1, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 176, .sgprs = 55, .lds = 34304, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 176, .sgprs = 55, .lds = 34304, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 190, .sgprs = 68, .lds = 34304, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 190, .sgprs = 68, .lds = 34304, .private = 0, .bound = 256, .wave = 32 }, .{ .vgprs = 1, .sgprs = 1, .lds = 0, .private = 0, .bound = 256, .wave = 32 }, }, }, diff --git a/zig/src/core/registry.zig b/zig/src/core/registry.zig index 7d9895e8bed..1f07408a611 100644 --- a/zig/src/core/registry.zig +++ b/zig/src/core/registry.zig @@ -110,6 +110,12 @@ pub fn Registry(comptime Launch: type) type { return r; } + /// The tuned row of one of the registry's entries, or null for an entry this GPU's table does not name. + pub fn rowOf(r: *const Self, e: *const E) ?tuning.Row { + const i = (@intFromPtr(e) - @intFromPtr(r.entries.ptr)) / @sizeOf(E); + return if (i < r.entries.len) r.tuned[i] else null; + } + /// The cost of entry `i` at `shape`, or null where its row has no room for the shape. fn cost(r: *const Self, i: usize, shape: Shape, env: Env) ?f32 { const e = &r.entries[i]; diff --git a/zig/src/core/tuning.zig b/zig/src/core/tuning.zig index a24e2a346a5..45a8a169483 100644 --- a/zig/src/core/tuning.zig +++ b/zig/src/core/tuning.zig @@ -12,6 +12,8 @@ pub const Row = struct { items: u32 = 0, /// The most pairs a routed plan holds per item: more than `pairs_per_item * items` is refused (0: any). pairs_per_item: f32 = 0, + /// A routed prompt item of at most this many rows takes its kernel's K-parallel passes (0: none does). + few_rows: u32 = 0, plain: bool = true, routed: bool = true, }; diff --git a/zig/src/hip/kernel_tests.zig b/zig/src/hip/kernel_tests.zig index f5e44b272e0..14415a76989 100644 --- a/zig/src/hip/kernel_tests.zig +++ b/zig/src/hip/kernel_tests.zig @@ -67,6 +67,7 @@ test { _ = @import("kernel_tests/affine/affine.zig"); _ = @import("kernel_tests/affine/rows.zig"); _ = @import("kernel_tests/affine/order.zig"); + _ = @import("kernel_tests/affine/item_few.zig"); _ = @import("kernel_tests/affine/round.zig"); _ = @import("kernel_tests/affine/layout.zig"); _ = @import("kernel_tests/affine/golden.zig"); diff --git a/zig/src/hip/kernel_tests/affine/cases.zig b/zig/src/hip/kernel_tests/affine/cases.zig index f367933e7fa..cf68a09ced8 100644 --- a/zig/src/hip/kernel_tests/affine/cases.zig +++ b/zig/src/hip/kernel_tests/affine/cases.zig @@ -19,6 +19,8 @@ pub const Case = struct { item_rows: usize = 8, /// Every expert takes the same number of pairs. even: bool = false, + /// The pairs each expert takes, in expert order (they add up to rows x slots); empty draws an expert a pair. + loads: []const u16 = &.{}, /// Each token's last slot is the shared expert (index `experts`), the others distinct routed experts. shared: bool = false, /// The plan serving passes (`hip/moe.zig`): items cut at its tile, padded with empty items to its capacity. diff --git a/zig/src/hip/kernel_tests/affine/item_few.zig b/zig/src/hip/kernel_tests/affine/item_few.zig new file mode 100644 index 00000000000..86cdba36f3d --- /dev/null +++ b/zig/src/hip/kernel_tests/affine/item_few.zig @@ -0,0 +1,72 @@ +//! The 64-row routed prompt tile's K-parallel passes write the matrix path's bytes for an item of few rows, at any cut. + +const std = @import("std"); +const base = @import("../rig.zig"); +const Buffer = base.Buffer; +const gpa = base.gpa; +const cases = @import("cases.zig"); +const prod = @import("product.zig"); +const Product = prod.Product; +const Rig = prod.Rig; +const Entry = @import("../../launches/affine.zig").Entry; +const item = @import("../../launches/matrix_item.zig"); + +/// Experts on both sides of every cut, of a row tile, of a block and of an item, then three of few rows at the end. +const loads = [_]u16{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 16, 17, 63, 64, 65, 128, 129, 200, 5, 8, 3 }; +const pairs = blk: { + var n: usize = 0; + for (loads) |l| n += l; + break :blk n; +}; + +/// The 35B's gate and up, its down, and a ragged product: a 72-column last block and a last round of one group. +const shapes = [_]cases.Case{ + .{ .name = "35b gate_up", .path = .prefill, .rows = pairs, .n = 1024, .k = 2048, .experts = loads.len, .slots = 1, .item_rows = 128, .loads = &loads }, + .{ .name = "35b down", .path = .prefill, .rows = pairs, .n = 2048, .k = 512, .experts = loads.len, .slots = 1, .item_rows = 128, .loads = &loads, .down = true }, + .{ .name = "ragged", .path = .prefill, .rows = pairs, .n = 200, .k = 2080, .group = 32, .experts = loads.len, .slots = 1, .item_rows = 128, .loads = &loads }, +}; + +/// What a run of outputs holds when no launch wrote it. +const unwritten: [16]u8 = @splat(0xa5); + +/// The tile's output on `p` with items of at most `few` rows on the K-parallel passes. +fn run(p: *const Product, e: *const Entry, out: Buffer, few: c_int) ![]u8 { + var a = p.arg(out.ptr, 0); + a.route.few = few; + try out.fill8(unwritten[0], null); + try e.launch(p.rig.kernels, .{ .r = &p.rig.base.r, .s = p.rig.base.stream.handle, .arg = a, .items = @intCast(p.items) }); + try p.rig.base.stream.synchronize(); + const bytes = try gpa.alloc(u8, p.outBytes(false)); + errdefer gpa.free(bytes); + try out.download(0, bytes); + return bytes; +} + +test "the routed item tile writes the matrix path's bytes through its K-parallel passes, at every cut and width" { + const b = try base.Rig.open(1 << 20); + defer b.close(); + var rig = Rig.open(b); + if (!rig.kernels.env().matrix) return error.SkipZigTest; + for (cases.widths) |bits| for (shapes) |shape| { + var c = shape; + c.bits = bits; + var p = try Product.init(&rig, c); + defer p.deinit(); + const e = for (rig.kernels.reg.entries) |*x| { + if (std.mem.endsWith(u8, x.id, ".routed.prefill.item64") and prod.takes(&p, x)) break x; + } else return error.NoItemTile; + var out = try Buffer.alloc(rig.base, p.outBytes(false)); + defer out.free(); + const matrix = try run(&p, e, out, 0); + defer gpa.free(matrix); + // every pair's output is written: no byte of the fill is left where an item of few rows sits + try base.expect(std.mem.indexOf(u8, matrix, &unwritten) == null, "{s} b{d}: the matrix path left outputs unwritten", .{ c.name, bits }); + // a cut below, at and past the rows a pass takes (the tile clamps the last) + for ([_]c_int{ 1, 3, item.few_rows, 64 }) |few| { + const passes = try run(&p, e, out, few); + defer gpa.free(passes); + const at = std.mem.indexOfDiff(u8, matrix, passes); + try base.expect(at == null, "{s} b{d} cut {d}: byte {?d} differs from the matrix path", .{ c.name, bits, few, at }); + } + }; +} diff --git a/zig/src/hip/kernel_tests/affine/product.zig b/zig/src/hip/kernel_tests/affine/product.zig index 24069c656be..2baeeaa27a3 100644 --- a/zig/src/hip/kernel_tests/affine/product.zig +++ b/zig/src/hip/kernel_tests/affine/product.zig @@ -101,7 +101,14 @@ pub const Product = struct { fn makePlan(p: *Product, pairs: usize) !void { const c = p.c; p.expert_of = try gpa.alloc(u32, pairs); - if (c.shared or c.served) { + if (c.loads.len != 0) { + var at: usize = 0; + for (c.loads, 0..) |load, e| { + @memset(p.expert_of[at..][0..load], @intCast(e)); + at += load; + } + std.debug.assert(at == pairs and c.loads.len <= c.experts); + } else if (c.shared or c.served) { drawServed(p.expert_of, c, &p.rig.rng); } else for (p.expert_of, 0..) |*e, i| { e.* = if (c.even) @intCast(i % c.experts) else @intCast(p.rig.rng.next() % c.experts); diff --git a/zig/src/hip/launches.zig b/zig/src/hip/launches.zig index 5d3c644de76..24b5aa308a2 100644 --- a/zig/src/hip/launches.zig +++ b/zig/src/hip/launches.zig @@ -10,6 +10,7 @@ const Stream = @import("stream.zig").Stream; const caps = @import("caps.zig"); const affine = @import("launches/affine.zig"); pub const dense = @import("launches/dense.zig"); +pub const item = @import("launches/matrix_item.zig"); const util = @import("launches/util.zig"); diff --git a/zig/src/hip/launches/affine.zig b/zig/src/hip/launches/affine.zig index 38d3d646731..a1b317a24b1 100644 --- a/zig/src/hip/launches/affine.zig +++ b/zig/src/hip/launches/affine.zig @@ -26,7 +26,7 @@ pub const Choices = struct { matrix: Matrix = .auto, stream: bool = true, gemm: /// quant/mlx.hpp's GroupTable, Routing and Affine as the kernels take them by value. pub const GroupTable = extern struct { p: u64, kind: c_int }; -pub const Routing = extern struct { items: u64 = 0, members: u64 = 0, x_div: c_int = 1, first: c_int = 0 }; +pub const Routing = extern struct { items: u64 = 0, members: u64 = 0, x_div: c_int = 1, first: c_int = 0, few: c_int = 0 }; pub const Arg = extern struct { x: u64, words: u64, @@ -49,7 +49,7 @@ pub const StreamSides = stream_layout.StreamSides; pub const Side = struct { words: u64, scale: u64, bias: u64, n: c_int, out: u64 }; comptime { - std.debug.assert(@sizeOf(GroupTable) == 16 and @sizeOf(Routing) == 24 and @sizeOf(Arg) == 112 and @sizeOf(StreamSides) == 184); + std.debug.assert(@sizeOf(GroupTable) == 16 and @sizeOf(Routing) == 32 and @sizeOf(Arg) == 120 and @sizeOf(StreamSides) == 184); } /// Everything a launch needs beside the kernels. @@ -457,7 +457,10 @@ pub const Kernels = struct { pub fn prefillLaunch(k: *const Kernels, r: *const runtime.Runtime, a: Arg, s: abi.Stream, items: c_int, pairs: c_int, expert_rows: c_int) Error!void { const op: registry.Op = if (a.route.items != 0) .routed else .project; const e = k.choose(op, .prefill, prefillShape(a, items, pairs, expert_rows)) orelse return refuse("prefill shape"); - try e.launch(k, .{ .r = r, .s = s, .arg = a, .items = items }); + var arg = a; + // this GPU's cut for a tile with K-parallel passes for items of few rows + if (k.reg.rowOf(e)) |row| arg.route.few = @intCast(row.few_rows); + try e.launch(k, .{ .r = r, .s = s, .arg = arg, .items = items }); } /// The product a prompt launch asks the registry about; a routed plan's pairs, and its average rows per expert. diff --git a/zig/src/hip/launches/matrix_item.zig b/zig/src/hip/launches/matrix_item.zig index b9be20377b4..3f9618d66a0 100644 --- a/zig/src/hip/launches/matrix_item.zig +++ b/zig/src/hip/launches/matrix_item.zig @@ -9,6 +9,8 @@ const util = @import("util.zig"); /// The 16-row tiles a block holds, by tier: 64 rows. pub const row_tiles = [_]usize{4}; +/// The most rows a tuning row's `few_rows` can cut at: what the tile sends to its K-parallel passes (kItemFewRows). +pub const few_rows = 8; pub const Functions = [affine.bit_widths.len][row_tiles.len]Function; /// The tier's functions at a width, as the module names them. @@ -28,9 +30,14 @@ pub fn rows(tier: usize) usize { /// Tier `tier` over every item of a plan: 128 columns by the tier's rows a block, all waves across the columns. pub fn go(k: *const affine.Kernels, c: affine.Call, tier: usize) util.Error!void { const b = std.mem.indexOfScalar(c_int, &affine.bit_widths, c.arg.bits) orelse return error.Invalid; + try goWith(k.item[b][tier], c, tier); +} + +/// `go` on a tile function of the tier's geometry: the tier's own, or a contract control in its place. +pub fn goWith(f: Function, c: affine.Call, tier: usize) util.Error!void { if (c.s == util.counting) return; var args: hl.Args = .{}; try args.add(c.arg); const grid = util.dim(util.cdiv(c.arg.n, 128), util.cdiv(c.arg.m, rows(tier)), c.items); - try hl.launch(k.item[b][tier], .{ .grid = grid, .block = .{ .x = 256 } }, .{ .r = c.r, .handle = c.s }, &args); + try hl.launch(f, .{ .grid = grid, .block = .{ .x = 256 } }, .{ .r = c.r, .handle = c.s }, &args); } diff --git a/zig/src/hip/tuning/gfx1151.zon b/zig/src/hip/tuning/gfx1151.zon index 175548d9394..cdecfe82346 100644 --- a/zig/src/hip/tuning/gfx1151.zon +++ b/zig/src/hip/tuning/gfx1151.zon @@ -7,9 +7,11 @@ // min_rows the fewest rows // items the most items of a routed plan (0: any) // pairs_per_item the most pairs a routed plan holds per item (0: any) +// few_rows a routed prompt item of at most this many rows takes the kernel's K-parallel passes (0: none) // plain, routed whether it takes plain products, routed plans (both by default) // gfx1100's rows, plus the 64-row routed matrix-core tile below 12 pairs an item (measured on gfx1151), // and the skewed-row matrix tile (three blocks a WGP) ahead of the padded one where it takes a product. +// The 64-row routed tile's items of 8 rows or fewer take its K-parallel passes (measured on gfx1151). .{ .gfx = "gfx1151", .entries = .{ @@ -43,7 +45,7 @@ .{ .id = "mlx.routed.prefill.kp3", .cost = 2, .items = 300, .plain = false }, .{ .id = "mlx.routed.prefill.kp1", .cost = 3, .rows = 4, .items = 300 }, .{ .id = "mlx.routed.prefill.kp2", .cost = 4, .rows = 32, .routed = false }, - .{ .id = "mlx.routed.prefill.item64", .cost = 9, .pairs_per_item = 12, .plain = false }, + .{ .id = "mlx.routed.prefill.item64", .cost = 9, .pairs_per_item = 12, .few_rows = 8, .plain = false }, .{ .id = "mlx.routed.prefill.skew", .cost = 9.5 }, .{ .id = "mlx.routed.prefill.matrix", .cost = 10 }, .{ .id = "mlx.routed.prefill.gemm", .cost = 11 }, diff --git a/zig/tests/hip/contract_main.zig b/zig/tests/hip/contract_main.zig index 2f656e6d6c7..59416a7b740 100644 --- a/zig/tests/hip/contract_main.zig +++ b/zig/tests/hip/contract_main.zig @@ -298,10 +298,14 @@ fn run(a: std.mem.Allocator, io: std.Io) !usize { for (&routers) |*x| try kernels.append(a, x.kernel()); // the few-row matrix-core tiles run on gfx11's matrix cores only const matrix = if (hip.Caps.of(arch)) |v| v.matrix == .wmma11 else false; - var items: [item_k.cases.len]item_k.Item = undefined; - for (&items, 0..) |*x, i| x.* = .init(a, &r, s, &l, i); + // the last one is the few-row passes' control, on its case's geometry + var items: [item_k.cases.len + 1]item_k.Item = undefined; + for (&items, 0..) |*x, i| x.* = .init(a, &r, s, &l, @min(i, item_k.control.index)); defer for (&items) |*x| x.deinit(); - if (matrix) for (&items) |*x| try kernels.append(a, x.kernel()); + if (matrix) { + items[item_k.cases.len].planted = try ctl.function(item_k.control.symbol); + for (&items) |*x| try kernels.append(a, x.kernel()); + } var embeds: [4]embed_k.Embed = undefined; var nembeds: usize = 0; defer for (embeds[0..nembeds]) |*x| x.deinit(); diff --git a/zig/tests/hip/contract_wmma_item.zig b/zig/tests/hip/contract_wmma_item.zig index f8e0b62bba1..33dbc3af8a7 100644 --- a/zig/tests/hip/contract_wmma_item.zig +++ b/zig/tests/hip/contract_wmma_item.zig @@ -13,18 +13,31 @@ const Kind = @typeInfo(@TypeOf(Kernels.blockGo)).@"fn".param_types[5].?; const Entry = @typeInfo(@FieldType(@FieldType(Kernels, "reg"), "entries")).pointer.child; const Call = @typeInfo(@typeInfo(@FieldType(Entry, "launch")).pointer.child).@"fn".param_types[1].?; -const Case = struct { tile: []const u8, rows: usize, group: usize, name: []const u8 }; +const Case = struct { tile: []const u8, rows: usize, group: usize, few: c_int, name: []const u8 }; const tiles = [_]struct { []const u8, usize }{.{ "item64", 64 }}; +/// The few-row cut a case's launches pass: none (the matrix path at every row count), then the tile's most. +const cuts = [_]c_int{ 0, hip.launches.item.few_rows }; pub const cases = blk: { - var out: [2 * tiles.len]Case = undefined; - for (tiles, 0..) |t, i| for ([_]usize{ 32, 64 }, 0..) |group, j| { - const name = std.fmt.comptimePrint("hip affine_wmma_item ({d} rows, 4-bit, group {d}, fp32 out)", .{ t[1], group }); - out[2 * i + j] = .{ .tile = t[0], .rows = t[1], .group = group, .name = name }; + var out: [2 * cuts.len * tiles.len]Case = undefined; + for (tiles, 0..) |t, i| for (cuts, 0..) |few, f| for ([_]usize{ 32, 64 }, 0..) |group, j| { + const name: []const u8 = if (few == 0) + std.fmt.comptimePrint("hip affine_wmma_item ({d} rows, 4-bit, group {d}, fp32 out)", .{ t[1], group }) + else + std.fmt.comptimePrint("hip affine_wmma_item ({d} rows, 4-bit, group {d}, fp32 out, K-parallel passes to {d} rows)", .{ t[1], group, few }); + out[2 * (cuts.len * i + f) + j] = .{ .tile = t[0], .rows = t[1], .group = group, .few = few, .name = name }; }; break :blk out; }; +/// The planted control on the last case's geometry: the few-row passes stop after a wave's first. Red at gate 3. +pub const control = struct { + pub const index = cases.len - 1; + pub const symbol = "tf_ctl_item_few_one_pass"; + const name = "control: routed item tile whose few-row passes stop after a wave's first"; + const gate: g.Gate = .shapes; +}; + /// The 128-row tile's chain: a group's products by k, its scale and bias, then the groups' sum. fn depth(s: c.Shape) usize { return s.k / 2 + s.k / 32 + 8; @@ -39,6 +52,8 @@ pub const Item = struct { f: ?contract.affine.Affine = null, arg: Arg = undefined, base: u64 = 0, + /// The planted control (its kernel from the controls' object, on the tile's geometry) in place of the served one. + planted: ?hip.Function = null, pub fn init(a: std.mem.Allocator, rt: *const hip.Runtime, s: hip.Stream, l: *const hip.Launcher, i: usize) Item { return .{ .rig = .{ .a = a, .r = rt, .s = s }, .l = l, .arena = .init(a), .index = i }; @@ -54,7 +69,7 @@ pub const Item = struct { const flash: c.Shape = .{ .rows = 22, .n = 1280, .k = 2560 }; return .{ .spec = .{ - .name = cfg.name, + .name = if (self.planted != null) control.name else cfg.name, .accum = .f32, .depth = depth, .out = .f32, @@ -62,10 +77,12 @@ pub const Item = struct { .nan = .{ .stated = "the 128-row tile's bytes" }, // a prompt's routed items: about 22 rows at 1,024 tokens a step, 88 at 4,096 .served = &.{ flash, .{ .rows = 88, .n = 1280, .k = 2560 }, .{ .rows = 5, .n = 1024, .k = 2048 } }, - .tiles = .{ .{ .size = cfg.rows }, .{ .size = 128 }, .{ .size = 256, .step = cfg.group, .min = cfg.group } }, + // rows: the block's edge, or the cut's when the launch passes one + .tiles = .{ .{ .size = if (cfg.few > 0) @intCast(cfg.few) else cfg.rows }, .{ .size = 128 }, .{ .size = 256, .step = cfg.group, .min = cfg.group } }, }, .ctx = self, .vt = &vt, + .control = if (self.planted != null) control.gate else null, }; } @@ -96,13 +113,18 @@ pub const Item = struct { .group = @intCast(cfg.group), .fp16 = 0, }; + self.arg.route.few = cfg.few; return true; } pub fn run(self: *Item) !void { const tile = cases[self.index].tile; + const call: Call = .{ .r = self.l.r, .s = self.rig.s.handle, .arg = self.arg, .items = 1 }; + if (self.planted) |f| { + try hip.launches.item.goWith(f, call, 0); + return self.rig.s.synchronize(); + } for (self.l.affine.reg.entries) |*e| if (std.mem.startsWith(u8, e.id, "mlx.routed.prefill.") and std.mem.endsWith(u8, e.id, tile)) { - const call: Call = .{ .r = self.l.r, .s = self.rig.s.handle, .arg = self.arg, .items = 1 }; try e.launch(&self.l.affine, call); try self.rig.s.synchronize(); return; From 02af5bdf7fe23b236bdb94c28f365971360d5f66 Mon Sep 17 00:00:00 2001 From: Abraham <119352736+AbrahamCete@users.noreply.github.com> Date: Wed, 11 Nov 2026 12:08:48 +0100 Subject: [PATCH 33/78] test(hip): pin the routed tile's cut from table to launch; rowOf checks bounds (#2878 review) The review's asks on the tile's cut for items of few rows: - registry: `rowOf` returns null for an entry before the registry's own instead of subtracting first, which trapped. A toy-registry test covers a named entry, an unnamed one and entries before and after it. - tuning: only `mlx.routed.prefill.item64` rows set `few_rows`, at no more than the tile's K-parallel passes take, and gfx1151's is set. - launches: `prefillLaunch` over a registry of toy entries carries its row's cut in `route.few`, whatever the caller's route held, and none for a row without one. It fails when the copy is lost. Refs #2795 --- zig/src/core/registry.zig | 24 ++++++++++++++-- zig/src/hip/launches/matrix_item.zig | 43 ++++++++++++++++++++++++++++ zig/src/hip/tuning/tuning.zig | 10 +++++++ 3 files changed, 75 insertions(+), 2 deletions(-) diff --git a/zig/src/core/registry.zig b/zig/src/core/registry.zig index 1f07408a611..e0ea439046e 100644 --- a/zig/src/core/registry.zig +++ b/zig/src/core/registry.zig @@ -110,9 +110,12 @@ pub fn Registry(comptime Launch: type) type { return r; } - /// The tuned row of one of the registry's entries, or null for an entry this GPU's table does not name. + /// The tuned row of entry `e`, or null for one this GPU's table does not name or the registry does not hold. pub fn rowOf(r: *const Self, e: *const E) ?tuning.Row { - const i = (@intFromPtr(e) - @intFromPtr(r.entries.ptr)) / @sizeOf(E); + const at = @intFromPtr(e); + const first = @intFromPtr(r.entries.ptr); + if (at < first) return null; + const i = (at - first) / @sizeOf(E); return if (i < r.entries.len) r.tuned[i] else null; } @@ -279,6 +282,23 @@ test "the cheapest entry that has room for the rows runs" { try std.testing.expect(std.mem.endsWith(u8, partial.select(toy_env, .mlx, .project, .decode, shape).?.id, "any")); } +test "an entry's tuned row is the table's, and null for an entry the table or the registry does not hold" { + const entries = [_]ToyEntry{ + toy("a.decode.before", .decode, .lanes, false), + toy("a.decode.few", .decode, .stream, true), + toy("a.decode.any", .decode, .lanes, false), + toy("a.decode.after", .decode, .lanes, false), + }; + const rows = [_]tuning.Row{.{ .id = "a.decode.few", .cost = 1, .few_rows = 8 }}; + // the registry holds the two entries in the middle, and the table names the first of them + const r = Toy.init(&.{ .gfx = "toy", .rows = &rows }, entries[1..3]); + try std.testing.expectEqual(@as(u32, 8), r.rowOf(&entries[1]).?.few_rows); + try std.testing.expect(r.rowOf(&entries[2]) == null); + // an entry before or after the registry's own has no row + try std.testing.expect(r.rowOf(&entries[0]) == null); + try std.testing.expect(r.rowOf(&entries[3]) == null); +} + test "a plan is chosen by its items, a plain product by its rows" { const entries = [_]ToyEntry{ toy("a.prefill.short", .prefill, .gemm, false), diff --git a/zig/src/hip/launches/matrix_item.zig b/zig/src/hip/launches/matrix_item.zig index 3f9618d66a0..0a6bd913830 100644 --- a/zig/src/hip/launches/matrix_item.zig +++ b/zig/src/hip/launches/matrix_item.zig @@ -6,6 +6,8 @@ const Module = @import("../module.zig").Module; const Function = @import("../module.zig").Function; const affine = @import("affine.zig"); const util = @import("util.zig"); +const registry = @import("core").registry; +const tuning = @import("../tuning/tuning.zig"); /// The 16-row tiles a block holds, by tier: 64 rows. pub const row_tiles = [_]usize{4}; @@ -41,3 +43,44 @@ pub fn goWith(f: Function, c: affine.Call, tier: usize) util.Error!void { const grid = util.dim(util.cdiv(c.arg.n, 128), util.cdiv(c.arg.m, rows(tier)), c.items); try hl.launch(f, .{ .grid = grid, .block = .{ .x = 256 } }, .{ .r = c.r, .handle = c.s }, &args); } + +// ---- the cut, from the tuning table to the launch + +var cut_launched: c_int = -1; + +fn keepCut(_: *const affine.Kernels, c: affine.Call) util.Error!void { + cut_launched = c.arg.route.few; +} + +fn anyEnv(_: registry.Env) bool { + return true; +} + +fn anyShape(_: registry.Shape) bool { + return true; +} + +fn toy(id: []const u8, op: registry.Op) affine.Entry { + return .{ .id = id, .format = .mlx, .op = op, .path = .prefill, .family = .gemm, .caps = &anyEnv, .policy = &anyEnv, .fits = &anyShape, .launch = &keepCut }; +} + +test "a prompt launch carries its entry's cut from the tuning table, and none for a row without one" { + const entries = [_]affine.Entry{ toy("toy.routed.prefill", .routed), toy("toy.project.prefill", .project) }; + const tuned = [_]tuning.Row{ .{ .id = "toy.routed.prefill", .cost = 1, .few_rows = few_rows }, .{ .id = "toy.project.prefill", .cost = 1 } }; + var k: affine.Kernels = undefined; + k.wmma = true; + k.matrix = true; + k.mode = .auto; + k.tile = .gemm; + k.stream_on = true; + k.reg = affine.Registry.init(&.{ .gfx = "toy", .rows = &tuned }, &entries); + // the caller's route holds another cut: the table's replaces it, for a plan and for a plain product + var a = std.mem.zeroes(affine.Arg); + a.m, a.n, a.k, a.bits, a.group = .{ 4, 128, 128, 4, 32 }; + a.route = .{ .items = 0x1000, .few = 3 }; + try k.prefillLaunch(undefined, a, util.counting, 2, 4, 2); + try std.testing.expectEqual(@as(c_int, few_rows), cut_launched); + a.route = .{ .few = 3 }; + try k.prefillLaunch(undefined, a, util.counting, 1, 0, 0); + try std.testing.expectEqual(@as(c_int, 0), cut_launched); +} diff --git a/zig/src/hip/tuning/tuning.zig b/zig/src/hip/tuning/tuning.zig index 4a636c328e6..b294625a2ba 100644 --- a/zig/src/hip/tuning/tuning.zig +++ b/zig/src/hip/tuning/tuning.zig @@ -19,6 +19,16 @@ test "only a routed prompt row bounds its pairs per item, by a finite number fro }; } +test "only the routed prompt's item tile cuts at few rows, at no more than its K-parallel passes take" { + const most = @import("../launches/matrix_item.zig").few_rows; + for ([_]*const Table{ &gfx1030, &gfx1100, &gfx1151, &gfx1201 }) |t| for (t.rows) |r| { + try std.testing.expect(r.few_rows <= most); + if (r.few_rows != 0) try std.testing.expectEqualStrings("mlx.routed.prefill.item64", r.id); + }; + // the cut is live where the receipts were taken: a table that lost it would pass the loop above + try std.testing.expect(find(&gfx1151, "mlx.routed.prefill.item64").?.few_rows != 0); +} + test "the tables name each entry once" { for ([_]*const Table{ &gfx1030, &gfx1100, &gfx1151, &gfx1201 }) |t| { for (t.rows, 0..) |r, i| { From aee8736615457bb6daeff6ab4f175965be6d8424 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 12:14:53 +0000 Subject: [PATCH 34/78] changelog: #2878 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index e55d1906beb..3ee3b8becd0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- HIP on gfx1151: the routed prompt tile runs an expert of at most 8 rows in K-parallel passes inside the same launch, with the matrix path's bytes; with Qwen3.6-35B-A3B on a Radeon 8060S, served cold prompts of 64 to 256 tokens are 11% to 22% faster (1.5% at 1,024), and logits and served tokens equal main's. Other cards' tables leave it off (#2878, fixes #2795). Thanks to @AbrahamCete. + - Internal: one block-FP8 reader in `core/quant/fp8.zig` loads Flash Next's, Kolibri-1's and Drama 3's FP8 weights and their scales, widening bf16 scales by the same bit shift and copying f32 ones as before; every packed buffer is byte-identical, and only Flash Next refuses non-finite scales, as before (#2892, #2460 item 7). Thanks to @jschmied. - CI: the macOS nightly runs `zig build test-native`, the Metal engine host tests, installing the Metal toolchain when the runner lacks it and printing its cold and warm wall time (#2891, fixes #2880). From 7d32f030557277e8f92b06f62d1bdab7d133bd00 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Wed, 11 Nov 2026 12:07:46 +0900 Subject: [PATCH 35/78] glm cuda: a split warm-up that NCCL fails refuses split by name before the rank stops Split's start-up warm-up sends and receives 4 MiB with every other rank, and runs before Engine.init's split probe, whose refusal names the fix. On a fabric where some ranks reach only their neighbours (a ring of four without a switch) the warm-up's group ends with an NCCL error first, and check stops every rank at once, so the probe's refusal never prints. From four ranks, where a ring has ranks that are not neighbours, the warm-up now prints that refusal, naming gather, when its group end fails under split, and hands check a name that starts with it, so a rank that stops on another rank's failure names the refusal as the cause too. Fewer ranks, ring and gather warm-ups, and a warm-up that finishes, print nothing new. --- zig/src/families/glm/cuda_warm_peers.zig | 17 ++++++++++++++++- zig/src/families/glm/cuda_warm_peers_test.zig | 15 +++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/zig/src/families/glm/cuda_warm_peers.zig b/zig/src/families/glm/cuda_warm_peers.zig index 0be9a6e49d0..d95f0054531 100644 --- a/zig/src/families/glm/cuda_warm_peers.zig +++ b/zig/src/families/glm/cuda_warm_peers.zig @@ -111,7 +111,22 @@ pub fn post(g: *const cuda.tp_group.Group, stream: cuda.Stream, buffer: u64, pee } break :posting {}; }; - const ended = g.lib.check(g.lib.api.ncclGroupEnd(), "ncclGroupEnd"); + const closed = g.lib.api.ncclGroupEnd(); + // split's warm-up is its first exchange with every rank; from four ranks a ring has some that are not neighbours + const refuse = peers == .every and g.world >= 4; + if (closed != 0 and refuse) refuseSplit(g.rank); // before check, which stops every rank and hands the others its name + const ended = g.lib.check(closed, if (refuse) "REFUSE TF_GLM_PREFILL_REDUCE=split at ncclGroupEnd" else "ncclGroupEnd"); try posted; try ended; } + +/// The latest rank that refused split, which a test build keeps instead of printing the refusal. +pub var refused_rank: ?usize = null; + +fn refuseSplit(rank: usize) void { + if (@import("builtin").is_test) { + refused_rank = rank; + return; + } + std.debug.print("REFUSE TF_GLM_PREFILL_REDUCE=split: rank {d} could not finish its start-up exchange with every other rank; split sends every rank's rows to every other, and a fabric where some ranks reach only their neighbours (a ring of four without a switch) needs TF_GLM_PREFILL_REDUCE=gather\n", .{rank}); +} diff --git a/zig/src/families/glm/cuda_warm_peers_test.zig b/zig/src/families/glm/cuda_warm_peers_test.zig index 6cea11161a9..50bd713ca99 100644 --- a/zig/src/families/glm/cuda_warm_peers_test.zig +++ b/zig/src/families/glm/cuda_warm_peers_test.zig @@ -139,6 +139,21 @@ test "GLM warm peers: failure at every posting closes only an opened group and s }; } +test "GLM warm peers: a split warm-up of four ranks or more that NCCL fails at its group's end refuses split by name first; fewer ranks, ring and gather do not" { + for (2..9) |world| for (0..world) |rank| for ([_]warm.Peers{ .every, .neighbours, .gather }) |peers| { + const g = Fake.group(rank, world); + Fake.reset(); + warm.refused_rank = null; + try Fake.run(&g, peers); + try t.expectEqual(@as(?usize, null), warm.refused_rank); + Fake.reset(); + if (peers == .gather) Fake.fail = 0 else Fake.end_fails = true; + try t.expectError(error.NcclFailed, Fake.run(&g, peers)); + try t.expectEqual(if (peers == .every and world >= 4) @as(?usize, rank) else null, warm.refused_rank); + }; + warm.refused_rank = null; +} + test "GLM warm peers: one rank does nothing and malformed placement refuses before NCCL" { Fake.reset(); var g = Fake.group(0, 1); From 70d0a30d06a6ed9cdce2ef4511784163067b3e8d Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Wed, 11 Nov 2026 16:10:36 +0900 Subject: [PATCH 36/78] changelog: #2868 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3ee3b8becd0..807b24df04f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -190,6 +190,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Internal: core gains one walk for copying kept prompt states, and GLM's copies on CUDA and Metal go through it, with the same copies in the same order; other families keep their own loops. Nothing a user sees changes (#2215, part of #1293). Thanks to @jschmied. +- GLM-5.3-Flash on CUDA at four ranks or more with the default split prefill: when the start-up warm-up's exchange with every rank fails, as on a switchless ring of four, the failing ranks print split's refusal naming `TF_GLM_PREFILL_REDUCE=gather` and the ranks that stop with them name it too; before, the ranks stopped without naming the setting (#2868). Thanks to @Bizuayeu. + - CUDA, opt-in build (`-Dnative-fp-decode`): qmmf's FP4 decode rows use the shared native e2m1 pair decode, `tf_common::e2m1_pairs`, which equals the existing decode on all 2^33 inputs; the default build's 113 fatbins are byte-identical and GB10 decode is level (#2721). Thanks to @jschmied. - Qwen3.8-27B on Metal: while a request runs, the keep-warm ticker now commits the residency sets' probe each interval, so the first command buffer after a wide prompt pass no longer waits in the kernel (likely for macOS to wire the model again; not traced). On an M5 Ultra with a 14 GiB prompt cache filling memory, five cold 24K- to 64K-token prompts in a row reach their first tokens 3.5% to 5.9% sooner in total; with the cache off, cold prompts are level. Replies are unchanged (#2725, fixes #2724). Thanks to @mrv777. From 685008cbb79ebbcf2705a51fb95bc00555696207 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 13:24:43 +0000 Subject: [PATCH 37/78] changelog: #2868's line from lines/2868.txt in place of its own --- CHANGELOG.md | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 807b24df04f..24ad3b85410 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- GLM-5.3-Flash on CUDA at four ranks or more: when the default split prefill's start-up exchange with every rank fails, as on a switchless ring of four, the failing rank now prints a refusal naming `TF_GLM_PREFILL_REDUCE=gather`, and rank 0 names it too when another rank fails first. Before, every rank stopped on a bare `ncclGroupEnd` error. Start-ups that succeed are unchanged (#2868). Thanks to @Bizuayeu. + - HIP on gfx1151: the routed prompt tile runs an expert of at most 8 rows in K-parallel passes inside the same launch, with the matrix path's bytes; with Qwen3.6-35B-A3B on a Radeon 8060S, served cold prompts of 64 to 256 tokens are 11% to 22% faster (1.5% at 1,024), and logits and served tokens equal main's. Other cards' tables leave it off (#2878, fixes #2795). Thanks to @AbrahamCete. - Internal: one block-FP8 reader in `core/quant/fp8.zig` loads Flash Next's, Kolibri-1's and Drama 3's FP8 weights and their scales, widening bf16 scales by the same bit shift and copying f32 ones as before; every packed buffer is byte-identical, and only Flash Next refuses non-finite scales, as before (#2892, #2460 item 7). Thanks to @jschmied. @@ -99,7 +101,6 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - CUDA: dots.mocr's text decoder serves text prompts on a GB10 from the publisher's BF16 weights, on a new shared engine for dense GQA decoders (`zig/src/cuda_dense/`); `--vision`, prompt caching and fp8 caches are refused by name, and served CUDA checkpoints are unchanged (#2673, part of #1927). Thanks to @laudney. - - The prompt cache can keep a bounded ladder of intermediate checkpoints under a long prompt's history mark, so a later, compacted prompt resumes inside its retained prefix instead of the shared cut. On one GB10 with `nvidia/Qwen3.8-Flash-Next-NVFP4`, a 104K-token request that kept its first 100K resumed 65,536 tokens instead of 2, cutting TTFT from 50.7 s to 19.9 s with the same tokens; a cold 165K prompt paid 0.9 s more, and the 8 GiB cache held two more states. A follow-up after that turn resumed 65,536 tokens too, because the compacted turn evicted the larger states. `TF_PROMPT_INTERMEDIATE=1` enables it for Flash Next on CUDA; it stays off by default (#1909). Its checkpoints go after a conversation's other redundant copies but before any conversation's newest state, and take only the room left after the endpoint, the rewind state and other requests' promised keeps. Thanks to @javierbits. - Metal: kernels nothing launches are removed: Qwen3.8-27B's `prompt_glue.metal` and `prompt_state.metal` (six `q27_prompt_*` kernels) and core's `tf_lane_rope` with `Attention.rotate`; saved Qwen3.8-27B prompt states on Metal are recomputed once after upgrading, since their identity covers the kernel sources (#3065, part of #2251). Thanks to @jschmied. @@ -190,8 +191,6 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Internal: core gains one walk for copying kept prompt states, and GLM's copies on CUDA and Metal go through it, with the same copies in the same order; other families keep their own loops. Nothing a user sees changes (#2215, part of #1293). Thanks to @jschmied. -- GLM-5.3-Flash on CUDA at four ranks or more with the default split prefill: when the start-up warm-up's exchange with every rank fails, as on a switchless ring of four, the failing ranks print split's refusal naming `TF_GLM_PREFILL_REDUCE=gather` and the ranks that stop with them name it too; before, the ranks stopped without naming the setting (#2868). Thanks to @Bizuayeu. - - CUDA, opt-in build (`-Dnative-fp-decode`): qmmf's FP4 decode rows use the shared native e2m1 pair decode, `tf_common::e2m1_pairs`, which equals the existing decode on all 2^33 inputs; the default build's 113 fatbins are byte-identical and GB10 decode is level (#2721). Thanks to @jschmied. - Qwen3.8-27B on Metal: while a request runs, the keep-warm ticker now commits the residency sets' probe each interval, so the first command buffer after a wide prompt pass no longer waits in the kernel (likely for macOS to wire the model again; not traced). On an M5 Ultra with a 14 GiB prompt cache filling memory, five cold 24K- to 64K-token prompts in a row reach their first tokens 3.5% to 5.9% sooner in total; with the cache off, cold prompts are level. Replies are unchanged (#2725, fixes #2724). Thanks to @mrv777. From 21cf81d131df6ed1f952585554c48a548e1acb3f Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 10 Nov 2026 18:53:09 +0000 Subject: [PATCH 38/78] refactor(radix): take a Query in the prompt store's begin and keep P2 of #2444 adds reuse identity (C6) beside the tree's words, so the store's begin and keep now take a Query instead of a bare prompt. The engine_api side carries every identity field with defaults and reads only the key for now; the core side carries the key and prompt length, which the CUDA adapter converts as it already does for Plan. Every caller passes the same words as before, so no store user changes behaviour. --- zig/src/core/kv_pool_store_test.zig | 4 +-- zig/src/core/kv_prompts.zig | 41 ++++++++++++++++-------- zig/src/core/prompt_radix.zig | 29 +++++++++++++++-- zig/src/core/prompt_radix_test.zig | 12 +++---- zig/src/families/dory/cuda_backend.zig | 4 +-- zig/src/families/gemma4/cuda_shared.zig | 4 +-- zig/src/families/glm/cuda_pages_test.zig | 16 ++++----- zig/src/families/glm/cuda_share.zig | 4 +-- zig/src/families/kolibri1/backend.zig | 4 +-- zig/src/families/qwen3_5/cuda_prefix.zig | 6 ++-- zig/src/families/spark25/backend.zig | 4 +-- zig/src/hip_serve/prefix.zig | 4 +-- zig/src/native/cuda_prompt_pages.zig | 17 ++++++---- 13 files changed, 96 insertions(+), 53 deletions(-) diff --git a/zig/src/core/kv_pool_store_test.zig b/zig/src/core/kv_pool_store_test.zig index 7260b1836a7..fa3f5b18123 100644 --- a/zig/src/core/kv_pool_store_test.zig +++ b/zig/src/core/kv_pool_store_test.zig @@ -79,7 +79,7 @@ const World = struct { fn admit(w: *World, a: Allocator, prompt: []const u32, history: u32, max_new: u32) !void { var adopt: std.ArrayList(u32) = .empty; defer adopt.deinit(w.gpa); - const plan = try w.store.begin(a, prompt, history, &.{}, null, &adopt); + const plan = try w.store.begin(a, .of(prompt), history, &.{}, null, &adopt); var l: Lease = .{}; errdefer w.pool.finish(&l); try w.pool.adopt(&l, adopt.items); @@ -89,7 +89,7 @@ const World = struct { for (plan.marks) |m| { var path: [64]u32 = undefined; const n = m / page; - const kept = w.store.keep(prompt, m, null, l.ids.items[0..n], path[0..n]); + const kept = w.store.keep(.of(prompt), m, null, l.ids.items[0..n], path[0..n]); for (path[0..kept.shared], 0..) |id, i| w.pool.share(&l, i, id); } try w.streams.append(w.gpa, l); diff --git a/zig/src/core/kv_prompts.zig b/zig/src/core/kv_prompts.zig index 19518d71b40..9f5f696336d 100644 --- a/zig/src/core/kv_prompts.zig +++ b/zig/src/core/kv_prompts.zig @@ -5,6 +5,19 @@ const kv = @import("kv_pool.zig"); /// Where a pass resumes and the marks it keeps states at (`marks` the store's allocation, freed by `free`). pub const Plan = struct { from: u32 = 0, marks: []const u32 = &.{} }; +/// What a pass asks the store about (prompt_radix.Query as a family on the core pool gives it). +pub const Query = struct { + /// The words the store compares: the request's cache key, one word a model row; for text, the prompt. + key: []const u32, + /// The prompt's length; a key of another length reuses nothing. + prompt_len: u32, + + /// A query of `key` alone, as long as its prompt. + pub fn of(key: []const u32) Query { + return .{ .key = key, .prompt_len = @intCast(key.len) }; + } +}; + /// The store as a family sees it: a prompt pass begins, keeps at its marks, and reports once its prompt is in. pub const Prompts = struct { ctx: *anyopaque, @@ -12,20 +25,20 @@ pub const Prompts = struct { pub const VTable = struct { /// The deepest state `prompt` resumes: its pages into `adopt` (held once more), the rest restored into `owner`. - begin: *const fn (ctx: *anyopaque, prompt: []const u32, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) anyerror!Plan, + begin: *const fn (ctx: *anyopaque, q: Query, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) anyerror!Plan, /// At `at` the store takes the pages of `mine` it lacks and the state; `path`'s first n replace the stream's. - keep: *const fn (ctx: *anyopaque, prompt: []const u32, at: u32, owner: *anyopaque, mine: []const u32, path: []u32) usize, + keep: *const fn (ctx: *anyopaque, q: Query, at: u32, owner: *anyopaque, mine: []const u32, path: []u32) usize, /// One log line after a prompt pass: where it resumed and how many states it kept. report: *const fn (ctx: *anyopaque, prompt: usize, from: u32, kept: u64) void, free: *const fn (ctx: *anyopaque, marks: []const u32) void, }; - pub fn begin(p: Prompts, prompt: []const u32, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) !Plan { - return p.vtable.begin(p.ctx, prompt, history_len, shared, owner, adopt); + pub fn begin(p: Prompts, q: Query, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) !Plan { + return p.vtable.begin(p.ctx, q, history_len, shared, owner, adopt); } - pub fn keep(p: Prompts, prompt: []const u32, at: u32, owner: *anyopaque, mine: []const u32, path: []u32) usize { - return p.vtable.keep(p.ctx, prompt, at, owner, mine, path); + pub fn keep(p: Prompts, q: Query, at: u32, owner: *anyopaque, mine: []const u32, path: []u32) usize { + return p.vtable.keep(p.ctx, q, at, owner, mine, path); } pub fn report(p: Prompts, prompt: usize, from: u32, kept: u64) void { @@ -37,8 +50,8 @@ pub const Prompts = struct { } /// `begin`, its marks in `gpa` for a lane to own; a failure gives `adopt`'s pages back. - pub fn plan(p: Prompts, gpa: std.mem.Allocator, pool: *kv.Pool, prompt: []const u32, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) !Plan { - const got = try p.begin(prompt, history_len, shared, owner, adopt); + pub fn plan(p: Prompts, gpa: std.mem.Allocator, pool: *kv.Pool, q: Query, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) !Plan { + const got = try p.begin(q, history_len, shared, owner, adopt); defer p.free(got.marks); const marks = gpa.dupe(u32, got.marks) catch |e| { for (adopt.items) |id| pool.release(id); @@ -48,13 +61,13 @@ pub const Prompts = struct { } /// At mark `at` the store takes the stream's pages and state; pages it held replace `l`'s, the whole table written. - pub fn keepPages(p: Prompts, gpa: std.mem.Allocator, pool: *kv.Pool, l: *kv.Lease, prompt: []const u32, at: usize, owner: *anyopaque, page_tokens: usize, table: kv.Table) !bool { + pub fn keepPages(p: Prompts, gpa: std.mem.Allocator, pool: *kv.Pool, l: *kv.Lease, q: Query, at: usize, owner: *anyopaque, page_tokens: usize, table: kv.Table) !bool { const n = at / page_tokens; const path = try gpa.alloc(u32, n); defer gpa.free(path); const mine = try gpa.dupe(u32, l.ids.items); defer gpa.free(mine); - const shared = p.keep(prompt, @intCast(at), owner, mine[0..n], path); + const shared = p.keep(q, @intCast(at), owner, mine[0..n], path); if (std.mem.eql(u32, mine[0..shared], path[0..shared])) return shared == n; @memcpy(mine[0..shared], path[0..shared]); // the same tokens' rows: a resume from any row is exact try table.write(table.ctx, mine, 0); @@ -163,10 +176,10 @@ const Stub = struct { fn of(ctx: *anyopaque) *Stub { return @ptrCast(@alignCast(ctx)); } - fn begin(_: *anyopaque, _: []const u32, _: u32, _: []const u32, _: *anyopaque, _: *std.ArrayList(u32)) anyerror!Plan { + fn begin(_: *anyopaque, _: Query, _: u32, _: []const u32, _: *anyopaque, _: *std.ArrayList(u32)) anyerror!Plan { return .{ .from = 256, .marks = try testing.allocator.dupe(u32, &.{512}) }; } - fn keep(ctx: *anyopaque, _: []const u32, _: u32, _: *anyopaque, mine: []const u32, path: []u32) usize { + fn keep(ctx: *anyopaque, _: Query, _: u32, _: *anyopaque, mine: []const u32, path: []u32) usize { @memcpy(path, Stub.of(ctx).held[0..mine.len]); return mine.len; } @@ -203,7 +216,7 @@ test "keeping swaps the store's pages in and writes the whole table: a table of const p: Prompts = .{ .ctx = &stub, .vtable = &Stub.vtable }; var w: Written = .{}; var owner: u8 = 0; - const whole = try p.keepPages(testing.allocator, &pool, &l, &.{ 1, 2, 3 }, 512, &owner, 256, .{ .ctx = &w, .write = Written.write }); + const whole = try p.keepPages(testing.allocator, &pool, &l, .of(&.{ 1, 2, 3 }), 512, &owner, 256, .{ .ctx = &w, .write = Written.write }); try testing.expect(whole); try testing.expectEqual(@as(usize, 0), w.from); try testing.expectEqual(@as(usize, 4), w.cap); @@ -221,7 +234,7 @@ test "a plan's marks are the lane's own; a blob's pages come from the pool and g var adopt: std.ArrayList(u32) = .empty; defer adopt.deinit(testing.allocator); var owner: u8 = 0; - const got = try p.plan(testing.allocator, &pool, &.{1}, 1, &.{}, &owner, &adopt); + const got = try p.plan(testing.allocator, &pool, .of(&.{1}), 1, &.{}, &owner, &adopt); defer testing.allocator.free(got.marks); try testing.expectEqual(@as(u32, 256), got.from); try testing.expectEqualSlices(u32, &.{512}, got.marks); diff --git a/zig/src/core/prompt_radix.zig b/zig/src/core/prompt_radix.zig index e1fd5586e16..54720dc4da8 100644 --- a/zig/src/core/prompt_radix.zig +++ b/zig/src/core/prompt_radix.zig @@ -41,6 +41,29 @@ pub fn pagesOf(comptime T: type, owner: *T, bytes: u64) Pages { } pub const Plan = pc.Plan; + +/// What a pass asks the store about: the words the tree compares, and what reuse must also match (enforced from C6). +pub const Query = struct { + /// The tree's words: the request's cache key, one word a model row; for text, the prompt. + key: []const u32, + /// The prompt's length; a key of another length reuses nothing. + prompt_len: u32, + /// The exact identities of the request's images. + images: []const @import("cache_images.zig").Identity = &.{}, + /// The request's image frame. + frame: u32 = 0, + /// What the rows depend on beyond key, images and frame. + namespace: u64 = 0, + /// The request's own chunk starts; empty: the store's grid. + starts: []const u32 = &.{}, + /// Rows this pass prefills; null: the prompt's length. + prefilled: ?u32 = null, + + /// A query of `key` alone, as long as its prompt. + pub fn of(key: []const u32) Query { + return .{ .key = key, .prompt_len = @intCast(key.len) }; + } +}; /// A paged store's `keep`: whether the mark holds a state, and how many leading pages the stream swaps for the store's. pub const Kept = struct { held: bool, shared: usize }; @@ -92,7 +115,8 @@ pub const Store = struct { } /// Restores the longest state `prompt` resumes into `owner`, its pages into `adopt` (held once more); plans marks. - pub fn begin(s: *Store, a: Allocator, prompt: []const u32, history_len: u32, shared: []const u32, owner: ?*anyopaque, adopt: *std.ArrayList(u32)) !Plan { + pub fn begin(s: *Store, a: Allocator, q: Query, history_len: u32, shared: []const u32, owner: ?*anyopaque, adopt: *std.ArrayList(u32)) !Plan { + const prompt = q.key; if (!s.on()) return .{ .marks = try a.alloc(u32, 0) }; const found = s.find(prompt); if (found == null) s.tree.counts.misses += 1; @@ -213,7 +237,8 @@ pub const Store = struct { } /// At `at` the tree takes the pages of `mine` it lacks and the snapshot, evicting to fit; `path` gets its pages. - pub fn keep(s: *Store, prompt: []const u32, at: u32, owner: ?*anyopaque, mine: []const u32, path: []u32) Kept { + pub fn keep(s: *Store, q: Query, at: u32, owner: ?*anyopaque, mine: []const u32, path: []u32) Kept { + const prompt = q.key; const t = &s.tree; if (!s.on() or at == 0 or at % s.step() != 0 or at + s.rules.lookahead > prompt.len or mine.len != at / s.rules.page or path.len != mine.len) return .{ .held = false, .shared = 0 }; t.clock += 1; diff --git a/zig/src/core/prompt_radix_test.zig b/zig/src/core/prompt_radix_test.zig index 55bb93bf3ab..bcefb3b4af6 100644 --- a/zig/src/core/prompt_radix_test.zig +++ b/zig/src/core/prompt_radix_test.zig @@ -84,7 +84,7 @@ fn seq(comptime n: usize, from: u32) [n]u32 { /// A pass over `prompt` as a backend runs it: begin, then a keep at each mark, the stream taking the tree's pages. fn pass(f: *Fake, s: *Store, a: Allocator, prompt: []const u32, history: u32, shared: []const u32, mine: *std.ArrayList(u32)) !Plan { mine.clearRetainingCapacity(); - const plan = try s.begin(a, prompt, history, shared, null, mine); + const plan = try s.begin(a, .of(prompt), history, shared, null, mine); var fresh: [32]u32 = undefined; const more = (prompt.len + 3) / 4 - mine.items.len; f.lane(fresh[0..more]); @@ -92,7 +92,7 @@ fn pass(f: *Fake, s: *Store, a: Allocator, prompt: []const u32, history: u32, sh for (plan.marks) |m| { var path: [32]u32 = undefined; const n = m / 4; - const kept = s.keep(prompt, m, null, mine.items[0..n], path[0..n]); + const kept = s.keep(.of(prompt), m, null, mine.items[0..n], path[0..n]); for (path[0..kept.shared], 0..) |id, i| if (mine.items[i] != id) { f.refs[id] += 1; f.refs[mine.items[i]] -= 1; @@ -237,7 +237,7 @@ test "a state past the limits is refused, and a failed copy keeps no state" { // the same state again takes no second copy var path: [5]u32 = undefined; const held = [_]u32{ 0, 1, 2, 3, 4 }; - _ = s.keep(&t, 20, null, &held, &path); + _ = s.keep(.of(&t), 20, null, &held, &path); try std.testing.expectEqual(@as(usize, 1), f.live); // a restore that fails drops the entry and the pass starts from the start f.fail_restore = true; @@ -412,7 +412,7 @@ fn sharedAfter(branch: usize, lookahead: u32) !usize { var own: [5]u32 = undefined; f.lane(&own); var path: [5]u32 = undefined; - const kept = s.keep(&second, 20, null, &own, &path); + const kept = s.keep(.of(&second), 20, null, &own, &path); f.leave(&own); try std.testing.expect(kept.held); try std.testing.expectEqual(s.tree.snaps, f.live); @@ -461,9 +461,9 @@ test "planned: a keep off the grid is refused, one on it kept" { defer f.leave(&own); var path: [4]u32 = undefined; const prompt = seq(24, 0); - try std.testing.expect(!s.keep(&prompt, 12, null, own[0..3], path[0..3]).held); + try std.testing.expect(!s.keep(.of(&prompt), 12, null, own[0..3], path[0..3]).held); try std.testing.expectEqual(@as(usize, 0), s.tree.snaps); - try std.testing.expect(s.keep(&prompt, 16, null, &own, &path).held); + try std.testing.expect(s.keep(.of(&prompt), 16, null, &own, &path).held); try std.testing.expectEqual(@as(usize, 1), s.tree.snaps); } diff --git a/zig/src/families/dory/cuda_backend.zig b/zig/src/families/dory/cuda_backend.zig index 023e886fb50..6e84aaba8f6 100644 --- a/zig/src/families/dory/cuda_backend.zig +++ b/zig/src/families/dory/cuda_backend.zig @@ -242,7 +242,7 @@ pub const Dory = struct { x.restore_slot = slot; defer x.restore_slot = null; x.restored = 0; - const got = try pr.plan(x.gpa, &x.pool.?, s.prompt(), s.history_len, s.shared_prefixes, @ptrCast(s), adopt); + const got = try pr.plan(x.gpa, &x.pool.?, .of(s.prompt()), s.history_len, s.shared_prefixes, @ptrCast(s), adopt); if (got.from == x.restored) return got; for (adopt.items) |id| x.pool.?.release(id); x.gpa.free(got.marks); @@ -309,7 +309,7 @@ pub const Dory = struct { /// TF_KV_SHARE: at mark `at` the store takes the lane's pages and rings; pages it held already replace the lane's. fn keepPages(x: *Dory, pr: kp.Prompts, s: *Stream, l: *Lane, at: usize) !void { var map: Map = .{ .x = x, .slot = l.slot }; - if (try pr.keepPages(x.gpa, &x.pool.?, &l.lease, s.prompt(), at, @ptrCast(s), pages.tokens, .{ .ctx = &map, .write = Map.write })) l.kept += 1; + if (try pr.keepPages(x.gpa, &x.pool.?, &l.lease, .of(s.prompt()), at, @ptrCast(s), pages.tokens, .{ .ctx = &map, .write = Map.write })) l.kept += 1; } /// One chunk from `l.fill` to the next grid line or mark; true once the prompt's last row is in. `_: u32`. diff --git a/zig/src/families/gemma4/cuda_shared.zig b/zig/src/families/gemma4/cuda_shared.zig index b43232c5152..3dd1dde0c38 100644 --- a/zig/src/families/gemma4/cuda_shared.zig +++ b/zig/src/families/gemma4/cuda_shared.zig @@ -108,7 +108,7 @@ pub fn opening(comptime Lane: type, x: anytype, s: *Stream) !*Lane { var adopted = false; errdefer if (!adopted) for (adopt.items) |id| x.pool.?.release(id); const pr = if (s.images.len == 0) x.prompts else null; - const plan: core.kv_prompts.Plan = if (pr) |p| try p.begin(prompt, s.history_len, s.shared_prefixes, s, &adopt) else .{}; + const plan: core.kv_prompts.Plan = if (pr) |p| try p.begin(.of(prompt), s.history_len, s.shared_prefixes, s, &adopt) else .{}; var l: Lane = .{ .slot = slot, .last = x.slot_bufs[2 * slot].ptr, .rule = x.slot_bufs[2 * slot + 1].ptr, .marks = plan.marks, .shared = pr != null }; errdefer { x.m.caps[slot] = 0; @@ -147,7 +147,7 @@ pub fn marked(x: anytype, s: *Stream, l: anytype, at: usize) !void { if (std.mem.indexOfScalar(u32, l.marks, @intCast(at)) != null) { const M = Map(@TypeOf(x)); var map: M = .{ .x = x, .slot = l.slot }; - if (try pr.keepPages(x.gpa, &x.pool.?, &l.lease, s.prompt(), at, s, fwd.pages.tokens, .{ .ctx = &map, .write = M.write })) l.kept += 1; + if (try pr.keepPages(x.gpa, &x.pool.?, &l.lease, .of(s.prompt()), at, s, fwd.pages.tokens, .{ .ctx = &map, .write = M.write })) l.kept += 1; } if (at == s.prompt().len) pr.report(at, s.cached, l.kept); } diff --git a/zig/src/families/glm/cuda_pages_test.zig b/zig/src/families/glm/cuda_pages_test.zig index 76de1632cfa..a84a958630a 100644 --- a/zig/src/families/glm/cuda_pages_test.zig +++ b/zig/src/families/glm/cuda_pages_test.zig @@ -424,13 +424,13 @@ const Store = struct { fn reclaim(ctx: *anyopaque, want: usize) bool { return of(ctx).store.reclaim(want); } - fn begin(ctx: *anyopaque, prompt: []const u32, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) anyerror!kp.Plan { + fn begin(ctx: *anyopaque, q: kp.Query, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) anyerror!kp.Plan { const x = of(ctx); - const plan = try x.store.begin(x.gpa, prompt, history_len, shared, owner, adopt); + const plan = try x.store.begin(x.gpa, .{ .key = q.key, .prompt_len = q.prompt_len }, history_len, shared, owner, adopt); return .{ .from = plan.from, .marks = plan.marks }; } - fn keep(ctx: *anyopaque, prompt: []const u32, at: u32, owner: *anyopaque, mine: []const u32, path: []u32) usize { - return of(ctx).store.keep(prompt, at, owner, mine, path).shared; + fn keep(ctx: *anyopaque, q: kp.Query, at: u32, owner: *anyopaque, mine: []const u32, path: []u32) usize { + return of(ctx).store.keep(.{ .key = q.key, .prompt_len = q.prompt_len }, at, owner, mine, path).shared; } fn report(_: *anyopaque, _: usize, _: u32, _: u64) void {} fn free(ctx: *anyopaque, marks: []const u32) void { @@ -460,7 +460,7 @@ test "TF_KV_SHARE: a pass's pages below its mark go to the shared store, and a l var a: [3000]u32 = undefined; for (&a, 1..) |*t, i| t.* = @intCast(i); // a cold pass: a mark on the grid below its end, where the tree takes the pages below it and the state - const cold = try prompts.begin(&a, 0, &.{}, &owner, &adopt); + const cold = try prompts.begin(.of(&a), 0, &.{}, &owner, &adopt); defer prompts.free(cold.marks); try std.testing.expectEqualSlices(u32, &.{2048}, cold.marks); try expect(cold.from == 0 and owner.pending == null and adopt.items.len == 0); @@ -474,7 +474,7 @@ test "TF_KV_SHARE: a pass's pages below its mark go to the shared store, and a l // the same 2,049 tokens, another tail: the pass resumes at 2,048 on the tree's pages, its state pending var b = a; for (b[2049..]) |*t| t.* += 10_000; - const warm = try prompts.begin(&b, 0, &.{}, &owner, &adopt); + const warm = try prompts.begin(.of(&b), 0, &.{}, &owner, &adopt); defer prompts.free(warm.marks); try expectEqual(@as(u32, 2048), warm.from); try expectEqual(@as(u32, 2048), owner.pending.?.at); @@ -487,14 +487,14 @@ test "TF_KV_SHARE: a pass's pages below its mark go to the shared store, and a l var c = a; for (c[2048..]) |*t| t.* += 10_000; owner = .{}; - const miss = try prompts.begin(&c, 0, &.{}, &owner, &adopt); + const miss = try prompts.begin(.of(&c), 0, &.{}, &owner, &adopt); defer prompts.free(miss.marks); try expect(miss.from == 0 and owner.pending == null and adopt.items.len == 0); // a stream that needs the tree's pages: its promise evicts the state, and the pages are the stream's _ = try p.openWith(0, c.len, 100, &.{}); // 7 pages: the 4 free and 3 of the tree's try expectEqual(@as(usize, 0), states.live); p.close(0); - const gone = try prompts.begin(&a, 0, &.{}, &owner, &adopt); + const gone = try prompts.begin(.of(&a), 0, &.{}, &owner, &adopt); defer prompts.free(gone.marks); try expect(gone.from == 0 and owner.pending == null and adopt.items.len == 0); try expect(p.pool.idle() and p.pool.census().free == 8); diff --git a/zig/src/families/glm/cuda_share.zig b/zig/src/families/glm/cuda_share.zig index 365a1f8b27e..75ee6b282a8 100644 --- a/zig/src/families/glm/cuda_share.zig +++ b/zig/src/families/glm/cuda_share.zig @@ -79,7 +79,7 @@ pub const Owner = struct { pending: ?*engine.Kept = null }; /// Before a prompt pass: the deepest kept state's pages in `adopt` (held for the stream), the state in `owner`. pub fn begin(e: *engine.Engine, s: *const Stream, owner: *Owner, adopt: *std.ArrayList(u32)) !kp.Plan { const prompts = of(e) orelse return error.NoPromptStore; - const plan = try prompts.begin(cacheKey(s) orelse return error.ImageCacheKeyRequired, s.history_len, s.shared_prefixes, owner, adopt); + const plan = try prompts.begin(.of(cacheKey(s) orelse return error.ImageCacheKeyRequired), s.history_len, s.shared_prefixes, owner, adopt); defer prompts.free(plan.marks); const mine = e.gpa.dupe(u32, plan.marks) catch |err| { // the lane's own marks, freed by its allocator e.pages.?.give(adopt.items); @@ -111,7 +111,7 @@ pub const Keeps = struct { const path = k.e.gpa.alloc(u32, n) catch return; defer k.e.gpa.free(path); var owner: Owner = .{}; - _ = prompts.keep(k.prompt, mark, &owner, ids[0..n], path); + _ = prompts.keep(.of(k.prompt), mark, &owner, ids[0..n], path); } }; diff --git a/zig/src/families/kolibri1/backend.zig b/zig/src/families/kolibri1/backend.zig index f6fafb94e29..078e9fa13f5 100644 --- a/zig/src/families/kolibri1/backend.zig +++ b/zig/src/families/kolibri1/backend.zig @@ -250,7 +250,7 @@ pub const Kolibri = struct { fn plan(x: *Kolibri, s: *Stream, adopt: *std.ArrayList(u32)) !kp.Plan { const pr = x.prompts orelse return .{}; x.pending = null; - return pr.plan(x.gpa, &x.pool.?, s.prompt(), s.history_len, s.shared_prefixes, @ptrCast(s), adopt); + return pr.plan(x.gpa, &x.pool.?, .of(s.prompt()), s.history_len, s.shared_prefixes, @ptrCast(s), adopt); } /// A new lane's table starts with the kept pages, held once more by its lease (rows below `from`, never written). @@ -328,7 +328,7 @@ pub const Kolibri = struct { /// TF_KV_SHARE: at mark `at` the store takes the lane's pages and rings; pages it held already replace the lane's. fn keepPages(x: *Kolibri, pr: kp.Prompts, s: *Stream, l: *Lane, at: usize) !void { var map: Map = .{ .x = x, .slot = l.slot }; - if (try pr.keepPages(x.gpa, &x.pool.?, &l.lease, s.prompt(), at, @ptrCast(s), fwd.page_tokens, .{ .ctx = &map, .write = Map.write })) l.kept += 1; + if (try pr.keepPages(x.gpa, &x.pool.?, &l.lease, .of(s.prompt()), at, @ptrCast(s), fwd.page_tokens, .{ .ctx = &map, .write = Map.write })) l.kept += 1; } /// One chunk from `l.fill` to the next grid line or mark; true once the prompt's last row is in. `_: u32`. diff --git a/zig/src/families/qwen3_5/cuda_prefix.zig b/zig/src/families/qwen3_5/cuda_prefix.zig index b33ae5dee87..0d3e40ec4df 100644 --- a/zig/src/families/qwen3_5/cuda_prefix.zig +++ b/zig/src/families/qwen3_5/cuda_prefix.zig @@ -97,7 +97,7 @@ pub const Prefix = struct { /// Before a pass: the deepest kept boundary (pages in `adopt`, state in `owner.pending`) and the marks to keep. pub fn begin(x: *Prefix, prompt: []const u32, history_len: u32, shared: []const u32, owner: *Owner, adopt: *std.ArrayList(u32)) !radix.Plan { - return x.store.begin(x.gpa, prompt, history_len, shared, owner, adopt); + return x.store.begin(x.gpa, .of(prompt), history_len, shared, owner, adopt); } /// The pending state copied into the stream's cache (its lane open and its pages adopted). @@ -121,7 +121,7 @@ pub const Prefix = struct { const mine = try x.gpa.dupe(u32, cache.paged.?.ids[0..n]); defer x.gpa.free(mine); var owner: Owner = .{ .cache = cache }; - const kept = x.store.keep(prompt, @intCast(at), &owner, mine, path); + const kept = x.store.keep(.of(prompt), @intCast(at), &owner, mine, path); try x.pages.swap(cache, path[0..kept.shared]); } @@ -130,7 +130,7 @@ pub const Prefix = struct { var owner: Owner = .{}; var adopt: std.ArrayList(u32) = .empty; defer adopt.deinit(x.gpa); - var plan = try x.store.begin(x.gpa, ids, history_len, shared, &owner, &adopt); + var plan = try x.store.begin(x.gpa, .of(ids), history_len, shared, &owner, &adopt); // the adopted pages are held for this call until openWith takes them plan.marks = x.onGrid(plan) catch |err| { for (adopt.items) |id| x.pages.pool.release(id); diff --git a/zig/src/families/spark25/backend.zig b/zig/src/families/spark25/backend.zig index 1da71fd557d..1b0c667461c 100644 --- a/zig/src/families/spark25/backend.zig +++ b/zig/src/families/spark25/backend.zig @@ -246,7 +246,7 @@ pub const Spark = struct { fn plan(x: *Spark, s: *Stream, adopt: *std.ArrayList(u32)) !kp.Plan { const pr = x.prompts orelse return .{}; x.pending = null; - return pr.plan(x.gpa, &x.pool.?, s.prompt(), s.history_len, s.shared_prefixes, @ptrCast(s), adopt); + return pr.plan(x.gpa, &x.pool.?, .of(s.prompt()), s.history_len, s.shared_prefixes, @ptrCast(s), adopt); } /// A new lane's table starts with the kept pages, held once more by its lease (rows below `from`, never written). @@ -325,7 +325,7 @@ pub const Spark = struct { /// TF_KV_SHARE: at mark `at` the store takes the lane's pages and rings; pages it held already replace the lane's. fn keepPages(x: *Spark, pr: kp.Prompts, s: *Stream, l: *Lane, at: usize) !void { var map: Map = .{ .x = x, .slot = l.slot }; - if (try pr.keepPages(x.gpa, &x.pool.?, &l.lease, s.prompt(), at, @ptrCast(s), fwd.page_tokens, .{ .ctx = &map, .write = Map.write })) l.kept += 1; + if (try pr.keepPages(x.gpa, &x.pool.?, &l.lease, .of(s.prompt()), at, @ptrCast(s), fwd.page_tokens, .{ .ctx = &map, .write = Map.write })) l.kept += 1; } /// One chunk from `l.fill` to the next grid line or mark. `_: u32` keeps its grid. diff --git a/zig/src/hip_serve/prefix.zig b/zig/src/hip_serve/prefix.zig index 092081b7f2f..26ce50605ef 100644 --- a/zig/src/hip_serve/prefix.zig +++ b/zig/src/hip_serve/prefix.zig @@ -68,7 +68,7 @@ pub fn Prefix(comptime F: type) type { o: *F.Owner, adopt: *std.ArrayList(u32), ) !radix.Plan { - return p.store.begin(p.gpa, prompt, history_len, shared, o, adopt); + return p.store.begin(p.gpa, .of(prompt), history_len, shared, o, adopt); } /// At `at` the tree takes the pages it lacks and the state; the stream swaps in the pages the tree had. @@ -78,7 +78,7 @@ pub fn Prefix(comptime F: type) type { const path = try p.gpa.alloc(u32, n); defer p.gpa.free(path); const table = F.table(o); - const kept = p.store.keep(prompt, @intCast(at), o, table[0..n], path); + const kept = p.store.keep(.of(prompt), @intCast(at), o, table[0..n], path); const own = path[0..kept.shared]; if (!std.mem.eql(u32, own, table[0..kept.shared])) { try F.adopt(p.f, o, own); diff --git a/zig/src/native/cuda_prompt_pages.zig b/zig/src/native/cuda_prompt_pages.zig index 3869d2a92fa..47b67828b7a 100644 --- a/zig/src/native/cuda_prompt_pages.zig +++ b/zig/src/native/cuda_prompt_pages.zig @@ -78,14 +78,19 @@ pub const Adapter = struct { return of(ctx).store.reclaim(want); } - fn begin(ctx: *anyopaque, prompt: []const u32, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) anyerror!hooks.Plan { + fn begin(ctx: *anyopaque, q: hooks.Query, history_len: u32, shared: []const u32, owner: *anyopaque, adopt: *std.ArrayList(u32)) anyerror!hooks.Plan { const x = of(ctx); - const p = try x.store.begin(x.gpa, prompt, history_len, shared, owner, adopt); + const p = try x.store.begin(x.gpa, query(q), history_len, shared, owner, adopt); return .{ .from = p.from, .marks = p.marks }; } - fn keep(ctx: *anyopaque, prompt: []const u32, at: u32, owner: *anyopaque, mine: []const u32, path: []u32) usize { - return of(ctx).store.keep(prompt, at, owner, mine, path).shared; + fn keep(ctx: *anyopaque, q: hooks.Query, at: u32, owner: *anyopaque, mine: []const u32, path: []u32) usize { + return of(ctx).store.keep(query(q), at, owner, mine, path).shared; + } + + /// The family's query as the store takes it. + fn query(q: hooks.Query) radix.Query { + return .{ .key = q.key, .prompt_len = q.prompt_len }; } fn report(ctx: *anyopaque, prompt: usize, from: u32, kept: u64) void { @@ -141,12 +146,12 @@ test "TF_KV_SHARE's adapter: a kept prompt's pages resume a later one, adopted; try f.pool.cover(&lease, 3, null); var path: [2]u32 = undefined; var owner: u8 = 0; - try testing.expectEqual(@as(usize, 2), p.keep(&prompt, 8, &owner, lease.ids.items[0..2], &path)); + try testing.expectEqual(@as(usize, 2), p.keep(.of(&prompt), 8, &owner, lease.ids.items[0..2], &path)); try testing.expectEqualSlices(u32, lease.ids.items[0..2], &path); f.pool.finish(&lease); var adopt: std.ArrayList(u32) = .empty; defer adopt.deinit(testing.allocator); - const plan = try p.begin(&prompt, 10, &.{}, &owner, &adopt); + const plan = try p.begin(.of(&prompt), 10, &.{}, &owner, &adopt); defer p.free(plan.marks); try testing.expectEqual(@as(u32, 8), plan.from); try testing.expectEqualSlices(u32, &path, adopt.items); From 38c7120f64f18256ce0be06478e1496e05599855 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 10 Nov 2026 18:59:32 +0000 Subject: [PATCH 39/78] refactor(radix): one store hold a page through a page registry The checkpoint manifest of #2444 (C3) lets several nodes and entries name one page, so the tree now holds its pages through Holds, a registry that takes one pool reference on a page's first name and gives it back on the last. The new prompt_radix_parts.zig also declares the manifest's types for the next steps. Node pages are distinct today, so the pool sees the same retains and releases in the same order and every store user is unchanged. --- zig/src/core/prompt_radix.zig | 9 ++- zig/src/core/prompt_radix_parts.zig | 96 ++++++++++++++++++++++++ zig/src/core/prompt_radix_parts_test.zig | 74 ++++++++++++++++++ zig/src/core/prompt_radix_test.zig | 6 +- zig/src/core/prompt_radix_tree.zig | 21 ++++-- zig/tests/qwen35/check_radix.zig | 4 +- 6 files changed, 194 insertions(+), 16 deletions(-) create mode 100644 zig/src/core/prompt_radix_parts.zig create mode 100644 zig/src/core/prompt_radix_parts_test.zig diff --git a/zig/src/core/prompt_radix.zig b/zig/src/core/prompt_radix.zig index 54720dc4da8..663b2cdce45 100644 --- a/zig/src/core/prompt_radix.zig +++ b/zig/src/core/prompt_radix.zig @@ -100,7 +100,7 @@ pub const Store = struct { pub fn limit(s: *Store, limits: Limits) void { s.limits = limits; while (s.tree.snaps > limits.snaps) s.tree.evict(s.victim(&.{}, null) orelse break); - while (s.tree.held_pages > limits.pages) if (!s.tree.evictLeaf()) break; + while (s.tree.heldPages() > limits.pages) if (!s.tree.evictLeaf()) break; } /// Whether anything can be kept. @@ -168,7 +168,7 @@ pub const Store = struct { const last = marks_[marks_.len - 1]; while (true) { const new = last / s.rules.page - s.tree.have(prompt[0 .. last + s.rules.lookahead]); - if (s.tree.held_pages + new <= s.limits.pages and s.tree.snaps + marks_.len <= s.limits.snaps) return; + if (s.tree.heldPages() + new <= s.limits.pages and s.tree.snaps + marks_.len <= s.limits.snaps) return; s.tree.evict(s.victim(prompt, from) orelse return); } } @@ -254,7 +254,7 @@ pub const Store = struct { t.touch(last.?); return .{ .held = true, .shared = mine.len }; } - if (t.held_pages + (mine.len - got) <= s.limits.pages and t.snaps + 1 <= s.limits.snaps) break; + if (t.heldPages() + (mine.len - got) <= s.limits.pages and t.snaps + 1 <= s.limits.snaps) break; const v = s.victim(prompt, null) orelse { t.counts.refused += 1; note("kept nothing at {d} tokens: no room", .{at}); @@ -332,7 +332,7 @@ pub const Store = struct { pub fn report(s: *const Store, prompt: usize, from: u32, kept: u64) void { if (@import("builtin").is_test) return; const t = &s.tree; - std.log.info("prompt cache: {d} tokens, resumed at {d}, kept {d}; {d} states and {d} pages ({d} MiB) of {d} states and {d} pages (hits {d}, misses {d}, evicted {d}, refused {d}, failed {d})", .{ prompt, from, kept, t.snaps, t.held_pages, t.held() >> 20, s.limits.snaps, s.limits.pages, t.counts.hits, t.counts.misses, t.counts.evicted, t.counts.refused, t.counts.failed }); + std.log.info("prompt cache: {d} tokens, resumed at {d}, kept {d}; {d} states and {d} pages ({d} MiB) of {d} states and {d} pages (hits {d}, misses {d}, evicted {d}, refused {d}, failed {d})", .{ prompt, from, kept, t.snaps, t.heldPages(), t.held() >> 20, s.limits.snaps, s.limits.pages, t.counts.hits, t.counts.misses, t.counts.evicted, t.counts.refused, t.counts.failed }); } }; @@ -344,4 +344,5 @@ fn note(comptime fmt: []const u8, args: anytype) void { test { _ = @import("prompt_radix_test.zig"); + _ = @import("prompt_radix_parts.zig"); } diff --git a/zig/src/core/prompt_radix_parts.zig b/zig/src/core/prompt_radix_parts.zig new file mode 100644 index 00000000000..8985549a2f2 --- /dev/null +++ b/zig/src/core/prompt_radix_parts.zig @@ -0,0 +1,96 @@ +//! The radix store's manifest (C3, #2444): a family's cache groups, a checkpoint's parts, and one hold a page. +const std = @import("std"); +const pc = @import("prompt_cache.zig"); +const Allocator = std.mem.Allocator; + +/// What a cache group holds; the core pool's kinds (kv_groups.zig), named here for the store. +pub const Kind = enum { full, window, state, draft }; + +/// How a checkpoint group's part is made: the stream's own read-only pages, or store-held pages its bytes go to. +pub const Make = enum { retained, copied }; + +/// A checkpoint group: its kind, how its part is made, and (retained) its rows a page. +pub const PartSpec = struct { kind: Kind, make: Make, rows: u32 = 0 }; + +/// A family's cache groups as its store sees them, declared once when the store opens. +pub const Groups = struct { + /// Rows a page of each prefix group; empty: one prefix group of the store's page. + prefix_rows: []const u32 = &.{}, + /// The checkpoint groups, in order. + parts: []const PartSpec = &.{}, + /// A legacy snapshot (the family's Snapshots) beside the parts. + legacy: bool = true, + + /// Prefix groups: page ids a logical page on the tree's path. + pub fn prefixCount(g: Groups) usize { + return @max(1, g.prefix_rows.len); + } + + /// No checkpoint at all (pure full attention): any matched page edge resumes. + pub fn none(g: Groups) bool { + return g.parts.len == 0 and !g.legacy; + } +}; + +/// One checkpoint group's pages in one entry; page `i` sits at logical index `base + i` in the group's table. +pub const Part = struct { base: u32, ids: []u32, bytes: u64 }; + +/// A keep's input for one retained group: its pages from logical index `base`, the last partial if rows fall short. +pub const Retained = struct { base: u32, ids: []const u32 }; + +/// A copy a stream makes from a held page into one of its own: rows [0, rows) of `src` to the group's page `index`. +pub const Copy = struct { group: u32, index: u32, src: u32, rows: u32 }; + +/// The store's references: one pool reference a page, however many nodes and entries name it. +pub const Holds = struct { + counts: std.AutoHashMapUnmanaged(u32, u32) = .empty, + + pub fn deinit(h: *Holds, gpa: Allocator) void { + h.counts.deinit(gpa); + } + + /// Room for `n` more pages, so the registrations after it cannot fail. + pub fn reserve(h: *Holds, gpa: Allocator, n: usize) error{OutOfMemory}!void { + try h.counts.ensureUnusedCapacity(gpa, @intCast(n)); + } + + /// The store names a page a stream holds: its first name takes one pool reference. + pub fn shareAssumeCapacity(h: *Holds, pages: pc.Pages, id: u32) void { + const slot = h.counts.getOrPutAssumeCapacity(id); + if (slot.found_existing) { + slot.value_ptr.* += 1; + return; + } + slot.value_ptr.* = 1; + pages.vtable.retain(pages.ptr, id); + } + + /// The store names a page it took itself (Pages.hold): that hold is its pool reference. + pub fn adoptHeldAssumeCapacity(h: *Holds, id: u32) void { + const slot = h.counts.getOrPutAssumeCapacity(id); + std.debug.assert(!slot.found_existing); // a page just held is new to the store + slot.value_ptr.* = 1; + } + + /// One name fewer; the last gives the pool reference back. + pub fn drop(h: *Holds, pages: pc.Pages, id: u32) void { + const slot = h.counts.getPtr(id) orelse std.debug.panic("prompt cache: page {d} dropped while not held", .{id}); + slot.* -= 1; + if (slot.* > 0) return; + _ = h.counts.remove(id); + pages.vtable.release(pages.ptr, id); + } + + pub fn contains(h: *const Holds, id: u32) bool { + return h.counts.contains(id); + } + + /// Distinct pages the store holds: what its page limit counts. + pub fn count(h: *const Holds) usize { + return h.counts.count(); + } +}; + +test { + _ = @import("prompt_radix_parts_test.zig"); +} diff --git a/zig/src/core/prompt_radix_parts_test.zig b/zig/src/core/prompt_radix_parts_test.zig new file mode 100644 index 00000000000..1341a6ad546 --- /dev/null +++ b/zig/src/core/prompt_radix_parts_test.zig @@ -0,0 +1,74 @@ +//! The manifest's registry alone: one pool reference a page, whatever names it, given back on the last drop. +const std = @import("std"); +const parts = @import("prompt_radix_parts.zig"); +const pc = @import("prompt_cache.zig"); +const testing = std.testing; + +/// Pages whose holders are counted, as a pool's are. +const Refs = struct { + refs: [8]u32 = @splat(0), + calls: u32 = 0, // retains and releases seen + + fn of(p: *anyopaque) *Refs { + return @ptrCast(@alignCast(p)); + } + fn retain(p: *anyopaque, id: u32) void { + of(p).refs[id] += 1; + of(p).calls += 1; + } + fn release(p: *anyopaque, id: u32) void { + of(p).refs[id] -= 1; + of(p).calls += 1; + } + fn holders(p: *anyopaque, id: u32) u32 { + return of(p).refs[id]; + } + fn available(_: *anyopaque) usize { + return 0; + } + fn pages(r: *Refs) pc.Pages { + return .{ .ptr = r, .bytes = 1, .vtable = &.{ .retain = retain, .release = release, .holders = holders, .available = available } }; + } +}; + +test "the store's registry: a page named twice holds one reference, given back on its last drop" { + var r: Refs = .{}; + r.refs[3] = 1; // a stream's page + var h: parts.Holds = .{}; + defer h.deinit(testing.allocator); + try h.reserve(testing.allocator, 2); + h.shareAssumeCapacity(r.pages(), 3); + h.shareAssumeCapacity(r.pages(), 3); // a second entry names it + try testing.expectEqual(@as(u32, 2), r.refs[3]); + try testing.expectEqual(@as(usize, 1), h.count()); + h.drop(r.pages(), 3); + try testing.expectEqual(@as(u32, 2), r.refs[3]); // still named once + h.drop(r.pages(), 3); + try testing.expectEqual(@as(u32, 1), r.refs[3]); + try testing.expectEqual(@as(usize, 0), h.count()); + try testing.expectEqual(@as(u32, 2), r.calls); // one retain, one release +} + +test "the store's registry: a page it held itself is its reference, released once" { + var r: Refs = .{}; + r.refs[5] = 1; // Pages.hold's reference + var h: parts.Holds = .{}; + defer h.deinit(testing.allocator); + try h.reserve(testing.allocator, 1); + h.adoptHeldAssumeCapacity(5); + try testing.expect(h.contains(5)); + try testing.expectEqual(@as(u32, 1), r.refs[5]); // no retain + h.drop(r.pages(), 5); + try testing.expectEqual(@as(u32, 0), r.refs[5]); +} + +test "a family's groups: one prefix group by default, a page-only store has no checkpoint" { + const d: parts.Groups = .{}; + try testing.expectEqual(@as(usize, 1), d.prefixCount()); + try testing.expect(!d.none()); + const two: parts.Groups = .{ .prefix_rows = &.{ 4, 4 }, .legacy = false }; + try testing.expectEqual(@as(usize, 2), two.prefixCount()); + try testing.expect(two.none()); + const state: parts.Groups = .{ .parts = &.{.{ .kind = .state, .make = .copied }}, .legacy = false }; + try testing.expect(!state.none()); +} diff --git a/zig/src/core/prompt_radix_test.zig b/zig/src/core/prompt_radix_test.zig index bcefb3b4af6..761bb229424 100644 --- a/zig/src/core/prompt_radix_test.zig +++ b/zig/src/core/prompt_radix_test.zig @@ -276,7 +276,7 @@ test "reclaim frees the pages of the least valuable leaf nobody holds, and spare try std.testing.expectEqual(@as(usize, 1), s.tree.snaps); for (y_pages) |id| f.refs[id] -= 1; s.limit(.{ .pages = 40, .snaps = 0 }); - try std.testing.expectEqual(@as(usize, 0), s.tree.held_pages); + try std.testing.expectEqual(@as(usize, 0), s.tree.heldPages()); try std.testing.expectEqual(@as(usize, 0), f.live); } @@ -333,7 +333,7 @@ fn goldenRun(a: Allocator, f: *Fake, s: *Store, out: *std.ArrayList(Row)) !void for (passes) |p| { const plan = try pass(f, s, a, p.prompt, p.history, p.shared, &mine); const c = s.tree.counts; - try out.append(a, .{ .from = plan.from, .marks = plan.marks, .snaps = s.tree.snaps, .pages = s.tree.held_pages, .bytes = s.tree.held(), .hits = c.hits, .misses = c.misses, .kept = c.kept, .evicted = c.evicted, .refused = c.refused }); + try out.append(a, .{ .from = plan.from, .marks = plan.marks, .snaps = s.tree.snaps, .pages = s.tree.heldPages(), .bytes = s.tree.held(), .hits = c.hits, .misses = c.misses, .kept = c.kept, .evicted = c.evicted, .refused = c.refused }); f.leave(mine.items); } } @@ -502,7 +502,7 @@ test "GLM's shape, a grid of four pages and lookahead 1: a chat resumes at its h try std.testing.expectEqualSlices(u32, &.{48}, plan.marks); f.leave(mine.items); try std.testing.expectEqual(@as(usize, 2), s.tree.snaps); - try std.testing.expectEqual(@as(usize, 12), s.tree.held_pages); + try std.testing.expectEqual(@as(usize, 12), s.tree.heldPages()); try std.testing.expectEqual(s.tree.snaps, f.live); } diff --git a/zig/src/core/prompt_radix_tree.zig b/zig/src/core/prompt_radix_tree.zig index 2c6770ee58f..3eff0969810 100644 --- a/zig/src/core/prompt_radix_tree.zig +++ b/zig/src/core/prompt_radix_tree.zig @@ -4,6 +4,7 @@ const std = @import("std"); const Allocator = std.mem.Allocator; const pc = @import("prompt_cache.zig"); +const parts = @import("prompt_radix_parts.zig"); pub const Saved = pc.Saved; pub const Snapshots = pc.Snapshots; @@ -50,8 +51,9 @@ pub const Tree = struct { ahead: usize, root: *Node, clock: u64 = 0, - /// Pages and snapshots the tree holds. - held_pages: usize = 0, + /// The pages the tree holds, one pool reference each however many nodes name them. + holds: parts.Holds = .{}, + /// Snapshots the tree holds. snaps: usize = 0, counts: Counts = .{}, @@ -63,6 +65,12 @@ pub const Tree = struct { pub fn deinit(t: *Tree) void { t.drop(t.root); + t.holds.deinit(t.gpa); + } + + /// Distinct pages the tree holds: what its page limit counts. + pub fn heldPages(t: *const Tree) usize { + return t.holds.count(); } fn drop(t: *Tree, n: *Node) void { @@ -72,7 +80,7 @@ pub const Tree = struct { t.family.vtable.drop(t.family.ptr, e.saved); t.gpa.free(e.last); } - for (n.pages) |id| t.pages.vtable.release(t.pages.ptr, id); + for (n.pages) |id| t.holds.drop(t.pages, id); t.gpa.free(n.tokens); t.gpa.free(n.pages); t.gpa.destroy(n); @@ -163,10 +171,10 @@ pub const Tree = struct { leaf.* = .{ .parent = node, .end = @intCast(np * t.page), .tokens = try t.gpa.dupe(u32, tokens[pi * t.page ..]), .pages = try t.gpa.dupe(u32, mine[pi..]) }; errdefer t.gpa.free(leaf.tokens); errdefer t.gpa.free(leaf.pages); + try t.holds.reserve(t.gpa, leaf.pages.len); // before the leaf is linked: what follows cannot fail try node.kids.append(t.gpa, leaf); - for (leaf.pages) |id| t.pages.vtable.retain(t.pages.ptr, id); + for (leaf.pages) |id| t.holds.shareAssumeCapacity(t.pages, id); @memcpy(path[pi..], leaf.pages); - t.held_pages += leaf.pages.len; t.touch(leaf); return leaf; }; @@ -224,7 +232,6 @@ pub const Tree = struct { if (n.kids.items.len > 0 or n.entry != null) return; const slot = std.mem.indexOfScalar(*Node, parent.kids.items, n).?; _ = parent.kids.swapRemove(slot); - t.held_pages -= n.pages.len; t.drop(n); n = parent; } @@ -260,7 +267,7 @@ pub const Tree = struct { /// Bytes the tree holds: its pages and its snapshots. pub fn held(t: *const Tree) u64 { - var n: u64 = t.held_pages * t.pages.bytes; + var n: u64 = t.heldPages() * t.pages.bytes; sum(t.root, &n); return n; } diff --git a/zig/tests/qwen35/check_radix.zig b/zig/tests/qwen35/check_radix.zig index f715078b268..a970125a9f2 100644 --- a/zig/tests/qwen35/check_radix.zig +++ b/zig/tests/qwen35/check_radix.zig @@ -161,11 +161,11 @@ fn memoryLine(c: *Ctx) void { const done = c.session.run(c.arena, .{ .prompts = set, .max_new = 1, .solo = true }) catch |err| return c.report.broke(name, err); const tree = &h.prefix.store.tree; const copied = copiedBytes(c, set, done.replies); - const unique = pagesBytes(c, tree.held_pages); + const unique = pagesBytes(c, tree.heldPages()); const states = tree.held() - unique; var cached: usize = 0; for (done.replies) |r| cached += r.cached; - c.report.pass(name, "{d} requests sharing {d} system tokens, {d} tokens each: the tree holds {d:.1} MiB ({d} pages {d:.1} MiB, {d} snapshots {d:.1} MiB); copies of the caches at each cut held {d:.1} MiB; {d} prompt tokens came from the tree", .{ set.len, system_len, set[0].ids.len, mib(tree.held()), tree.held_pages, mib(unique), tree.snaps, mib(states), mib(copied), cached }); + c.report.pass(name, "{d} requests sharing {d} system tokens, {d} tokens each: the tree holds {d:.1} MiB ({d} pages {d:.1} MiB, {d} snapshots {d:.1} MiB); copies of the caches at each cut held {d:.1} MiB; {d} prompt tokens came from the tree", .{ set.len, system_len, set[0].ids.len, mib(tree.held()), tree.heldPages(), mib(unique), tree.snaps, mib(states), mib(copied), cached }); } fn pagesBytes(c: *Ctx, n: usize) u64 { From 3153e25202e135ad3dd408826e88a32383436455 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 10 Nov 2026 19:08:46 +0000 Subject: [PATCH 40/78] feat(radix): page-major prefix groups on the store's path A family whose cache has several prefix groups (a target and a drafter's full attention on one pool, #2444 C3) needs one page id of each group for every logical page on the tree's path. The tree now keeps per ids a logical page, page-major, counts logical pages in its walks and ids in its limits, and initGroups declares the groups, refusing groups with other rows a page than the store's. With one group, today's case, every index and count is as before. --- zig/src/core/prompt_radix.zig | 37 +++++-- zig/src/core/prompt_radix_groups_test.zig | 125 ++++++++++++++++++++++ zig/src/core/prompt_radix_tree.zig | 39 ++++--- 3 files changed, 174 insertions(+), 27 deletions(-) create mode 100644 zig/src/core/prompt_radix_groups_test.zig diff --git a/zig/src/core/prompt_radix.zig b/zig/src/core/prompt_radix.zig index 663b2cdce45..3376ad0adaa 100644 --- a/zig/src/core/prompt_radix.zig +++ b/zig/src/core/prompt_radix.zig @@ -3,6 +3,7 @@ const std = @import("std"); const pc = @import("prompt_cache.zig"); const tree_mod = @import("prompt_radix_tree.zig"); +const parts = @import("prompt_radix_parts.zig"); const Allocator = std.mem.Allocator; pub const Saved = tree_mod.Saved; @@ -11,6 +12,7 @@ pub const Pages = tree_mod.Pages; pub const Counts = tree_mod.Counts; pub const Node = tree_mod.Node; pub const Tree = tree_mod.Tree; +pub const Groups = parts.Groups; /// The rules of prompt_cache.zig, paged: `page` set; a planned store cuts and resumes only on `grid`, a multiple of it. pub const Rules = pc.Rules; @@ -64,7 +66,7 @@ pub const Query = struct { return .{ .key = key, .prompt_len = @intCast(key.len) }; } }; -/// A paged store's `keep`: whether the mark holds a state, and how many leading pages the stream swaps for the store's. +/// A paged store's `keep`: whether the mark holds a state, and how many leading ids the stream swaps for the store's. pub const Kept = struct { held: bool, shared: usize }; pub const Store = struct { @@ -77,8 +79,19 @@ pub const Store = struct { const SHARED_KEYS = 64; pub fn init(gpa: Allocator, family: Snapshots, pages: Pages, rules: Rules) !Store { + return initGroups(gpa, family, pages, rules, .{}); + } + + /// `init` for a family that declares its cache groups (C3, #2444): `P` prefix groups page-major on the path. + pub fn initGroups(gpa: Allocator, family: Snapshots, pages: Pages, rules: Rules, groups: Groups) !Store { if (rules.page == 0 or (rules.planned and (rules.grid == 0 or rules.grid % rules.page != 0))) return error.BadPromptRadixRules; - return .{ .tree = try Tree.init(gpa, family, pages, rules.page, rules.lookahead), .rules = rules }; + for (groups.prefix_rows) |rows| if (rows != rules.page) return error.UnequalPrefixRows; // one key a page for all + return .{ .tree = try Tree.init(gpa, family, pages, rules.page, rules.lookahead, groups.prefixCount()), .rules = rules }; + } + + /// Page ids a logical page: one per prefix group. + fn per(s: *const Store) usize { + return s.tree.per; } /// Where a pass may cut and resume: a page edge, or a planned store's grid. @@ -167,7 +180,7 @@ pub const Store = struct { if (marks_.len == 0) return; const last = marks_[marks_.len - 1]; while (true) { - const new = last / s.rules.page - s.tree.have(prompt[0 .. last + s.rules.lookahead]); + const new = (last / s.rules.page - s.tree.have(prompt[0 .. last + s.rules.lookahead])) * s.per(); if (s.tree.heldPages() + new <= s.limits.pages and s.tree.snaps + marks_.len <= s.limits.snaps) return; s.tree.evict(s.victim(prompt, from) orelse return); } @@ -198,9 +211,9 @@ pub const Store = struct { var out: std.ArrayList(u32) = .empty; errdefer out.deinit(a); for (marks_) |m| { - if (m / s.rules.page <= s.limits.pages and s.limits.snaps > 0) try out.append(a, m) else { + if (m / s.rules.page * s.per() <= s.limits.pages and s.limits.snaps > 0) try out.append(a, m) else { s.tree.counts.refused += 1; - note("kept nothing at {d} tokens: {d} pages pass the limit of {d}", .{ m, m / s.rules.page, s.limits.pages }); + note("kept nothing at {d} tokens: {d} pages pass the limit of {d}", .{ m, m / s.rules.page * s.per(), s.limits.pages }); } } return out.toOwnedSlice(a); @@ -236,29 +249,30 @@ pub const Store = struct { return out.toOwnedSlice(a); } - /// At `at` the tree takes the pages of `mine` it lacks and the snapshot, evicting to fit; `path` gets its pages. + /// At `at` the tree takes the ids of `mine` (page-major, `per` a page) it lacks, and the snapshot; `path` its ids. pub fn keep(s: *Store, q: Query, at: u32, owner: ?*anyopaque, mine: []const u32, path: []u32) Kept { const prompt = q.key; const t = &s.tree; - if (!s.on() or at == 0 or at % s.step() != 0 or at + s.rules.lookahead > prompt.len or mine.len != at / s.rules.page or path.len != mine.len) return .{ .held = false, .shared = 0 }; + const np = at / s.rules.page; // logical pages below `at` + if (!s.on() or at == 0 or at % s.step() != 0 or at + s.rules.lookahead > prompt.len or mine.len != np * s.per() or path.len != mine.len) return .{ .held = false, .shared = 0 }; t.clock += 1; const tokens = prompt[0 .. at + s.rules.lookahead]; // a page past which the lookahead differs is another page const flagged = std.mem.indexOfScalar(u64, s.shared_keys.items, sharedKey(tokens)) != null; while (true) { var last: ?*Node = null; - const got = t.heldPrefix(tokens, mine.len, path, &last); - if (got == mine.len and last != null and last.?.end == at and last.?.entry != null) { + const got = t.heldPrefix(tokens, np, path, &last); + if (got == np and last != null and last.?.end == at and last.?.entry != null) { const e = &last.?.entry.?; e.used = t.clock; // the same state again: no copy e.shared = e.shared or flagged; t.touch(last.?); return .{ .held = true, .shared = mine.len }; } - if (t.heldPages() + (mine.len - got) <= s.limits.pages and t.snaps + 1 <= s.limits.snaps) break; + if (t.heldPages() + (np - got) * s.per() <= s.limits.pages and t.snaps + 1 <= s.limits.snaps) break; const v = s.victim(prompt, null) orelse { t.counts.refused += 1; note("kept nothing at {d} tokens: no room", .{at}); - return .{ .held = false, .shared = got }; + return .{ .held = false, .shared = got * s.per() }; }; t.evict(v); } @@ -345,4 +359,5 @@ fn note(comptime fmt: []const u8, args: anytype) void { test { _ = @import("prompt_radix_test.zig"); _ = @import("prompt_radix_parts.zig"); + _ = @import("prompt_radix_groups_test.zig"); } diff --git a/zig/src/core/prompt_radix_groups_test.zig b/zig/src/core/prompt_radix_groups_test.zig new file mode 100644 index 00000000000..827672bbb12 --- /dev/null +++ b/zig/src/core/prompt_radix_groups_test.zig @@ -0,0 +1,125 @@ +//! The store over two prefix groups (C3, #2444): ids page-major on the path, limits counted in ids. +const std = @import("std"); +const radix = @import("prompt_radix.zig"); +const testing = std.testing; + +/// Pages whose holders are counted, and states that are the position they were kept at; a page is 4 tokens. +const Fake = struct { + refs: [32]u32 = @splat(0), + states: [8]u32 = @splat(0), + nstates: usize = 0, + restored: ?u32 = null, + + fn of(p: *anyopaque) *Fake { + return @ptrCast(@alignCast(p)); + } + fn retain(p: *anyopaque, id: u32) void { + of(p).refs[id] += 1; + } + fn release(p: *anyopaque, id: u32) void { + of(p).refs[id] -= 1; + } + fn holders(p: *anyopaque, id: u32) u32 { + return of(p).refs[id]; + } + fn available(p: *anyopaque) usize { + var n: usize = 0; + for (of(p).refs) |r| n += @intFromBool(r == 0); + return n; + } + fn bytes(_: *anyopaque, _: u32) u64 { + return 1; + } + fn save(p: *anyopaque, _: ?*anyopaque, at: u32) anyerror!radix.Saved { + const f = of(p); + f.states[f.nstates] = at; + f.nstates += 1; + return @ptrCast(&f.states[f.nstates - 1]); + } + fn restore(p: *anyopaque, _: ?*anyopaque, saved: radix.Saved) anyerror!void { + const at: *u32 = @ptrCast(@alignCast(saved)); + of(p).restored = at.*; + } + fn dropState(_: *anyopaque, _: radix.Saved) void {} + + fn pages(f: *Fake) radix.Pages { + return .{ .ptr = f, .bytes = 1, .vtable = &.{ .retain = retain, .release = release, .holders = holders, .available = available } }; + } + fn family(f: *Fake) radix.Snapshots { + return .{ .ptr = f, .vtable = &.{ .bytes = bytes, .save = save, .restore = restore, .drop = dropState } }; + } + /// A stream's own ids: taken pages, one holder each. + fn take(f: *Fake, ids: []const u32) void { + for (ids) |id| f.refs[id] = 1; + } +}; + +const two: radix.Groups = .{ .prefix_rows = &.{ 4, 4 } }; + +fn open(f: *Fake, limit_ids: usize) !radix.Store { + var s = try radix.Store.initGroups(testing.allocator, f.family(), f.pages(), .{ .page = 4, .lookahead = 0, .min_gap = 1, .min_prompt = 0 }, two); + s.limit(.{ .pages = limit_ids, .snaps = 4 }); + return s; +} + +test "two prefix groups: a keep holds each logical page's two ids, and a resume adopts them page-major" { + var f: Fake = .{}; + var s = try open(&f, 16); + defer s.deinit(); + const prompt = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 }; + const mine = [_]u32{ 0, 1, 2, 3 }; // page 0: ids 0 (group 0) and 1 (group 1); page 1: 2 and 3 + f.take(&mine); + var path: [4]u32 = undefined; + const kept = s.keep(.of(&prompt), 8, null, &mine, &path); + try testing.expect(kept.held); + try testing.expectEqual(@as(usize, 4), kept.shared); // ids, not logical pages + try testing.expectEqualSlices(u32, &mine, &path); + try testing.expectEqual(@as(usize, 4), s.tree.heldPages()); + for (mine) |id| try testing.expectEqual(@as(u32, 2), f.refs[id]); // the stream and the store + for (mine) |id| f.refs[id] -= 1; // the stream ends + + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const plan = try s.begin(testing.allocator, .of(&prompt), 10, &.{}, null, &adopt); + defer testing.allocator.free(plan.marks); + try testing.expectEqual(@as(u32, 8), plan.from); + try testing.expectEqual(@as(?u32, 8), f.restored); + try testing.expectEqualSlices(u32, &mine, adopt.items); + for (adopt.items) |id| f.refs[id] -= 1; +} + +test "two prefix groups: a longer prompt shares the first page's ids and adds its own, split page-major" { + var f: Fake = .{}; + var s = try open(&f, 16); + defer s.deinit(); + const a = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 0 }; + const b = [_]u32{ 1, 2, 3, 4, 9, 9, 9, 9, 0 }; + const mine_a = [_]u32{ 0, 1, 2, 3 }; + const mine_b = [_]u32{ 4, 5, 6, 7 }; + f.take(&mine_a); + f.take(&mine_b); + var path: [4]u32 = undefined; + try testing.expect(s.keep(.of(&a), 8, null, &mine_a, &path).held); + const kept = s.keep(.of(&b), 8, null, &mine_b, &path); + try testing.expect(kept.held); + try testing.expectEqualSlices(u32, &.{ 0, 1, 6, 7 }, &path); // the shared first page is a's ids; b's second is its own + try testing.expectEqual(@as(usize, 6), s.tree.heldPages()); +} + +test "two prefix groups: the page limit counts ids, so a keep needing more ids than it holds evicts or refuses" { + var f: Fake = .{}; + var s = try open(&f, 4); // two logical pages of two groups + defer s.deinit(); + const long = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 0 }; + const mine = [_]u32{ 0, 1, 2, 3, 4, 5 }; + f.take(&mine); + var path: [6]u32 = undefined; + try testing.expect(!s.keep(.of(&long), 12, null, &mine, &path).held); // 6 ids pass the limit of 4 + try testing.expectEqual(@as(usize, 0), s.tree.heldPages()); + try testing.expect(s.keep(.of(&long), 8, null, mine[0..4], path[0..4]).held); // 4 ids fit +} + +test "prefix groups with other rows a page than the store's are refused when it opens" { + var f: Fake = .{}; + try testing.expectError(error.UnequalPrefixRows, radix.Store.initGroups(testing.allocator, f.family(), f.pages(), .{ .page = 4, .min_prompt = 0 }, .{ .prefix_rows = &.{ 4, 8 } })); +} diff --git a/zig/src/core/prompt_radix_tree.zig b/zig/src/core/prompt_radix_tree.zig index 3eff0969810..eee106ebf31 100644 --- a/zig/src/core/prompt_radix_tree.zig +++ b/zig/src/core/prompt_radix_tree.zig @@ -47,6 +47,8 @@ pub const Tree = struct { pages: Pages, /// Tokens a page holds. page: usize, + /// Prefix groups: page ids a logical page, page-major on each node (`pages[p * per + g]`). + per: usize = 1, /// Tokens past a page that its rows read (the family's lookahead): a page matches only when they match too. ahead: usize, root: *Node, @@ -57,10 +59,15 @@ pub const Tree = struct { snaps: usize = 0, counts: Counts = .{}, - pub fn init(gpa: Allocator, family: Snapshots, pages: Pages, page: usize, ahead: usize) !Tree { + pub fn init(gpa: Allocator, family: Snapshots, pages: Pages, page: usize, ahead: usize, per: usize) !Tree { const root = try gpa.create(Node); root.* = .{ .parent = null, .end = 0, .tokens = &.{}, .pages = &.{} }; - return .{ .gpa = gpa, .family = family, .pages = pages, .page = page, .ahead = ahead, .root = root }; + return .{ .gpa = gpa, .family = family, .pages = pages, .page = page, .per = per, .ahead = ahead, .root = root }; + } + + /// Logical pages of `n`: its ids over the prefix groups. + pub fn npages(t: *const Tree, n: *const Node) usize { + return n.pages.len / t.per; } pub fn deinit(t: *Tree) void { @@ -101,7 +108,7 @@ pub const Tree = struct { fn common(t: *const Tree, n: *const Node, tokens: []const u32, most: usize) usize { const w = t.page + t.ahead; var p: usize = 0; - while (p < n.pages.len and p < most and std.mem.eql(u32, n.tokens[p * t.page ..][0..w], tokens[p * t.page ..][0..w])) p += 1; + while (p < t.npages(n) and p < most and std.mem.eql(u32, n.tokens[p * t.page ..][0..w], tokens[p * t.page ..][0..w])) p += 1; return p; } @@ -113,7 +120,7 @@ pub const Tree = struct { while (at + t.page <= limit_at) { const c = t.child(node, prompt[at..]) orelse break; const p = t.common(c, prompt[at..], (limit_at - at) / t.page); - if (p < c.pages.len) break; + if (p < t.npages(c)) break; at += p * t.page; node = c; if (c.entry != null and c.end % step == 0) best = c; @@ -121,17 +128,17 @@ pub const Tree = struct { return best; } - /// Pages of `tokens` (`np` whole pages and `ahead` more) the tree holds, into `path` if any; `last` the end node. + /// Logical pages of `tokens` (`np` whole pages and `ahead` more) the tree holds, their ids into `path` if any. pub fn heldPrefix(t: *const Tree, tokens: []const u32, np: usize, path: ?[]u32, last: *?*Node) usize { var node = t.root; var pi: usize = 0; while (pi < np) { const c = t.child(node, tokens[pi * t.page ..]) orelse break; const p = t.common(c, tokens[pi * t.page ..], np - pi); - if (path) |out| @memcpy(out[pi..][0..p], c.pages[0..p]); + if (path) |out| @memcpy(out[pi * t.per ..][0 .. p * t.per], c.pages[0 .. p * t.per]); pi += p; node = c; - if (p < c.pages.len) break; + if (p < t.npages(c)) break; } last.* = if (node == t.root) null else node; return pi; @@ -149,7 +156,7 @@ pub const Tree = struct { defer path.deinit(t.gpa); var w: ?*Node = n; while (w) |x| : (w = x.parent) if (x.parent != null) try path.append(t.gpa, x); - try adopt.ensureUnusedCapacity(t.gpa, n.end / t.page); + try adopt.ensureUnusedCapacity(t.gpa, n.end / t.page * t.per); while (path.pop()) |x| { for (x.pages) |id| { t.pages.vtable.retain(t.pages.ptr, id); @@ -159,28 +166,28 @@ pub const Tree = struct { } } - /// Puts `tokens` (whole pages, `mine` the caller's, and `ahead` more) in the tree; the end node, `path` its pages. + /// Puts `tokens` (whole pages, `mine` the caller's ids page-major, `ahead` more) in the tree; `path` gets its ids. pub fn insert(t: *Tree, tokens: []const u32, mine: []const u32, path: []u32) !*Node { - const np = mine.len; + const np = mine.len / t.per; var node = t.root; var pi: usize = 0; while (pi < np) { const c = t.child(node, tokens[pi * t.page ..]) orelse { const leaf = try t.gpa.create(Node); errdefer t.gpa.destroy(leaf); - leaf.* = .{ .parent = node, .end = @intCast(np * t.page), .tokens = try t.gpa.dupe(u32, tokens[pi * t.page ..]), .pages = try t.gpa.dupe(u32, mine[pi..]) }; + leaf.* = .{ .parent = node, .end = @intCast(np * t.page), .tokens = try t.gpa.dupe(u32, tokens[pi * t.page ..]), .pages = try t.gpa.dupe(u32, mine[pi * t.per ..]) }; errdefer t.gpa.free(leaf.tokens); errdefer t.gpa.free(leaf.pages); try t.holds.reserve(t.gpa, leaf.pages.len); // before the leaf is linked: what follows cannot fail try node.kids.append(t.gpa, leaf); for (leaf.pages) |id| t.holds.shareAssumeCapacity(t.pages, id); - @memcpy(path[pi..], leaf.pages); + @memcpy(path[pi * t.per ..], leaf.pages); t.touch(leaf); return leaf; }; const p = t.common(c, tokens[pi * t.page ..], np - pi); - const up = if (p < c.pages.len) try t.split(c, p) else c; - @memcpy(path[pi..][0..p], up.pages); + const up = if (p < t.npages(c)) try t.split(c, p) else c; + @memcpy(path[pi * t.per ..][0 .. p * t.per], up.pages); t.touch(up); node = up; pi += p; @@ -194,11 +201,11 @@ pub const Tree = struct { errdefer t.gpa.destroy(up); const head_tokens = try t.gpa.dupe(u32, c.tokens[0 .. p * t.page + t.ahead]); // its pages, the tail's first `ahead` errdefer t.gpa.free(head_tokens); - const head_pages = try t.gpa.dupe(u32, c.pages[0..p]); + const head_pages = try t.gpa.dupe(u32, c.pages[0 .. p * t.per]); errdefer t.gpa.free(head_pages); const tail_tokens = try t.gpa.dupe(u32, c.tokens[p * t.page ..]); errdefer t.gpa.free(tail_tokens); - const tail_pages = try t.gpa.dupe(u32, c.pages[p..]); + const tail_pages = try t.gpa.dupe(u32, c.pages[p * t.per ..]); errdefer t.gpa.free(tail_pages); const parent = c.parent.?; up.* = .{ .parent = parent, .end = parent.end + @as(u32, @intCast(p * t.page)), .tokens = head_tokens, .pages = head_pages, .used = c.used }; From 0a686d18e7c440af8553771a31ca8ecbbbf27822 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 10 Nov 2026 19:35:35 +0000 Subject: [PATCH 41/78] feat(radix): checkpoint parts made, restored and reclaimed C3 of #2444: a family may declare checkpoint groups whose pages an entry holds beside its legacy snapshot. A copied part takes store-held pool pages its live bytes are copied into; a retained part names the stream's own pages and copies a partial last page. Entries are published only whole, with every copy fenced before a page goes back; a node is pinned while a keep builds it or a resume restores it, so reclaim never frees its own source; reclaim first frees idle parts. Also fixes a leak in insert when its second allocation failed. Stores without parts are unchanged. --- zig/src/core/kv_pool_checks.zig | 1 + zig/src/core/kv_pool_parts_test.zig | 393 ++++++++++++++++++++++ zig/src/core/prompt_cache.zig | 2 + zig/src/core/prompt_radix.zig | 137 +++++++- zig/src/core/prompt_radix_groups_test.zig | 4 +- zig/src/core/prompt_radix_make.zig | 156 +++++++++ zig/src/core/prompt_radix_parts.zig | 19 ++ zig/src/core/prompt_radix_tree.zig | 49 ++- 8 files changed, 732 insertions(+), 29 deletions(-) create mode 100644 zig/src/core/kv_pool_parts_test.zig create mode 100644 zig/src/core/prompt_radix_make.zig diff --git a/zig/src/core/kv_pool_checks.zig b/zig/src/core/kv_pool_checks.zig index fb0678a4a73..a938dcc392d 100644 --- a/zig/src/core/kv_pool_checks.zig +++ b/zig/src/core/kv_pool_checks.zig @@ -1,6 +1,7 @@ //! The page pool's checks beside engine_api: under the paged prompt store and a lane host (kv_pool is its own module). test { _ = @import("kv_pool_store_test.zig"); + _ = @import("kv_pool_parts_test.zig"); _ = @import("kv_pool_host_test.zig"); _ = @import("kv_pool_ledger_test.zig"); } diff --git a/zig/src/core/kv_pool_parts_test.zig b/zig/src/core/kv_pool_parts_test.zig new file mode 100644 index 00000000000..eebc961d2a7 --- /dev/null +++ b/zig/src/core/kv_pool_parts_test.zig @@ -0,0 +1,393 @@ +//! The paged store's checkpoint parts (C3, #2444) over a real pool: made all or nothing, held once, reclaimed, pinned. +const std = @import("std"); +const kv = @import("kv_pool"); +const radix = @import("engine_api").prompt_radix; +const Pool = kv.Pool; +const Lease = kv.Lease; +const testing = std.testing; +const expect = testing.expect; +const expectEqual = testing.expectEqual; +const expectEqualSlices = testing.expectEqualSlices; + +const page = 4; // tokens a page of the prefix group +const page_bytes = 8; + +/// A family whose device memory is one word a page: a copy writes the live state's word, a restore reads it back. +const Family = struct { + pool: *Pool, + mem: [64]u32 = @splat(0), + live: [2]u32 = .{ 0, 0 }, // each group's live state + restored: [2]?u32 = .{ null, null }, + fences: u32 = 0, + copies: u32 = 0, + rows_copied: u32 = 0, + fail_copy: bool = false, + fail_rows: bool = false, + queued: []const u32 = &.{}, // the pages the last copy wrote (queued work, until a fence) + held_at_fence: bool = true, // every fence saw the queued pages still held: nothing went back before it + hold_on_restore: usize = 0, // pages the restore takes from the pool (its reclaim runs) + + fn of(p: *anyopaque) *Family { + return @ptrCast(@alignCast(p)); + } + fn bytes(_: *anyopaque, _: u32, _: u32) u64 { + return 20; // three pages of 8 bytes + } + fn copy(p: *anyopaque, _: ?*anyopaque, group: u32, _: u32, dst: radix.Parts.Part) anyerror!void { + const f = of(p); + for (dst.ids) |id| f.mem[id] = f.live[group]; + f.queued = dst.ids; + if (f.fail_copy) return error.CopyFailed; // after its work is queued + f.copies += 1; + } + fn restore(p: *anyopaque, _: ?*anyopaque, group: u32, src: radix.Parts.Part) anyerror!void { + const f = of(p); + if (f.hold_on_restore > 0) { + var ids: [8]u32 = undefined; + if (!f.pool.hold(ids[0..f.hold_on_restore])) return error.NoPages; + for (ids[0..f.hold_on_restore]) |id| f.pool.release(id); + } + f.restored[group] = f.mem[src.ids[0]]; + } + fn copyRows(p: *anyopaque, _: ?*anyopaque, _: u32, src: u32, dst: u32, rows: u32) anyerror!void { + const f = of(p); + if (f.fail_rows) return error.CopyFailed; + f.mem[dst] = f.mem[src]; + f.rows_copied += rows; + } + fn fence(p: *anyopaque, _: ?*anyopaque) void { + const f = of(p); + f.fences += 1; + for (f.queued) |id| f.held_at_fence = f.held_at_fence and f.pool.holders(id) > 0; + f.queued = &.{}; + } + fn hooks(f: *Family) radix.PartHooks { + return .{ .ptr = f, .vtable = &.{ .bytes = bytes, .copy = copy, .restore = restore, .copy_rows = copyRows, .fence = fence } }; + } + + fn snapBytes(_: *anyopaque, _: u32) u64 { + return 0; + } + fn save(_: *anyopaque, _: ?*anyopaque, _: u32) anyerror!radix.Saved { + return error.NoLegacy; + } + fn snapRestore(_: *anyopaque, _: ?*anyopaque, _: radix.Saved) anyerror!void {} + fn drop(_: *anyopaque, _: radix.Saved) void {} + fn snapshots(f: *Family) radix.Snapshots { + return .{ .ptr = f, .vtable = &.{ .bytes = snapBytes, .save = save, .restore = snapRestore, .drop = drop } }; + } +}; + +const copied_state = [_]radix.Parts.PartSpec{.{ .kind = .state, .make = .copied }}; +const window3 = [_]radix.Parts.PartSpec{.{ .kind = .window, .make = .retained, .rows = 3 }}; +const both = [_]radix.Parts.PartSpec{ .{ .kind = .window, .make = .retained, .rows = 3 }, .{ .kind = .state, .make = .copied } }; + +const rules: radix.Rules = .{ .page = page, .lookahead = 0, .min_gap = 1, .min_prompt = 0 }; + +const World = struct { + pool: Pool, + f: Family = undefined, + store: radix.Store = undefined, + reclaims: u32 = 0, // the pool asked the store + + fn init(w: *World, specs: []const radix.Parts.PartSpec, pages: usize) !void { + w.pool = try Pool.init(testing.allocator, pages, page_bytes); + w.f = .{ .pool = &w.pool }; + w.store = try radix.Store.initGroups(testing.allocator, w.f.snapshots(), radix.pagesOf(Pool, &w.pool, page_bytes), rules, .{ .parts = specs, .legacy = false }, w.f.hooks()); + w.store.limit(.{ .pages = pages, .snaps = 8 }); + try w.pool.joinReclaim(.{ .ctx = &w.store, .call = reclaim }); + } + fn deinit(w: *World) void { + w.pool.leaveReclaim(&w.store); + w.store.deinit(); + expect(w.pool.idle() and w.pool.available() == w.pool.total()) catch @panic("pages left after close"); + w.pool.deinit(); + } + fn reclaim(ctx: *anyopaque, want: usize) bool { + const s: *radix.Store = @ptrCast(@alignCast(ctx)); + @as(*World, @fieldParentPtr("store", s)).reclaims += 1; + return s.reclaim(want); + } + /// A stream's lease of `n` pages. + fn lease(w: *World, n: usize) !Lease { + var l: Lease = .{}; + try w.pool.promise(&l, n); + try w.pool.cover(&l, n, null); + return l; + } +}; + +const prompt = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14 }; + +test "a copied part: held from the pool at keep, restored and fenced at a resume, back at eviction" { + var w: World = .{ .pool = undefined }; + try w.init(&copied_state, 16); + defer w.deinit(); + var l = try w.lease(2); + w.f.live[0] = 77; + var path: [2]u32 = undefined; + const kept = w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &.{.{ .base = 0, .ids = &.{} }}, &path); + try expect(kept.held); + try expectEqual(@as(usize, 5), w.store.tree.heldPages()); // two path pages, three part pages + w.pool.finish(&l); + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const plan = try w.store.begin(testing.allocator, .of(&prompt), 14, &.{}, null, &adopt); + defer testing.allocator.free(plan.marks); + try expectEqual(@as(u32, 8), plan.from); + try expectEqual(@as(?u32, 77), w.f.restored[0]); + try expect(w.f.fences >= 2); // the keep's copies, then the restore's + for (adopt.items) |id| w.pool.release(id); + const before = w.pool.available(); + w.store.limit(.{ .pages = 16, .snaps = 0 }); // the entry goes, its parts with it + try expectEqual(before + 5, w.pool.available()); // the part pages and the path the entry alone kept +} + +test "a retained window with other rows a page: whole pages shared, the partial one copied, mapped at their logical index" { + var w: World = .{ .pool = undefined }; + try w.init(&window3, 32); + defer w.deinit(); + var l = try w.lease(2); // the prefix group + var win = try w.lease(3); // the window: rows 0..8 on 3-row pages, the last one partial (2 rows) + for (win.ids.items, 0..) |id, i| w.f.mem[id] = @intCast(100 + i); + var path: [2]u32 = undefined; + const r = [_]radix.Retained{.{ .base = 1, .ids = win.ids.items[1..3] }}; // its window omits page 0 + try expect(w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &r, &path).held); + try expectEqual(@as(u32, 2), w.f.rows_copied); + try expectEqual(@as(u32, 2), w.pool.holders(win.ids.items[1])); // the stream and the store + try expectEqual(@as(u32, 1), w.pool.holders(win.ids.items[2])); // the stream's partial page: never the store's + const last_page = win.ids.items[2]; + w.f.mem[last_page] = 999; // the stream writes on: the checkpoint does not change + w.pool.finish(&l); + w.pool.finish(&win); + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const plan = try w.store.begin(testing.allocator, .of(&prompt), 14, &.{}, null, &adopt); + defer plan.deinit(testing.allocator); + try expectEqual(@as(usize, 1), plan.parts.len); + try expectEqual(@as(u32, 1), plan.parts[0].base); + try expectEqual(@as(usize, 1), plan.parts[0].ids.len); // page 1: whole + try expectEqual(@as(usize, 1), plan.copies.len); + try expectEqual(@as(u32, 2), plan.copies[0].index); // logical page 2, rows 0..2 + try expectEqual(@as(u32, 2), plan.copies[0].rows); + try expectEqual(@as(u32, 102), w.f.mem[plan.copies[0].src]); // the rows as they were at the keep + for (adopt.items) |id| w.pool.release(id); + for (plan.parts[0].ids) |id| w.pool.release(id); + const fences = w.f.fences; + w.store.finish(null, plan.copies); + try expectEqual(fences + 1, w.f.fences); +} + +test "a part copy that fails publishes nothing: the path stays, the part pages go back after a fence" { + var w: World = .{ .pool = undefined }; + try w.init(&both, 32); + defer w.deinit(); + var l = try w.lease(2); + defer w.pool.finish(&l); + var win = try w.lease(3); + defer w.pool.finish(&win); + const r = [_]radix.Retained{ .{ .base = 0, .ids = win.ids.items }, .{ .base = 0, .ids = &.{} } }; + var path: [2]u32 = undefined; + for ([_]bool{ true, false }) |rows_fail| { + w.f.fail_rows = rows_fail; + w.f.fail_copy = !rows_fail; + const fences = w.f.fences; + const kept = w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &r, &path); + try expect(!kept.held); + try expectEqual(@as(usize, 2), kept.shared); // the path is the store's + try expectEqual(@as(usize, 2), w.store.tree.heldPages()); // and nothing else + try expectEqual(fences + 1, w.f.fences); + try expect(w.f.held_at_fence); // the fence came before the pages went back + try expectEqual(@as(usize, 0), w.store.tree.snaps); + for (win.ids.items) |id| try expectEqual(@as(u32, 1), w.pool.holders(id)); + } + try expectEqual(@as(u64, 2), w.store.tree.counts.failed); +} + +test "two entries whose windows overlap hold each shared window page once" { + var w: World = .{ .pool = undefined }; + try w.init(&window3, 32); + defer w.deinit(); + var l = try w.lease(3); + defer w.pool.finish(&l); + var win = try w.lease(4); // rows 0..12 on 3-row pages + defer w.pool.finish(&win); + var path: [3]u32 = undefined; + try expect(w.store.keepParts(.of(&prompt), 4, null, l.ids.items[0..1], &.{.{ .base = 0, .ids = win.ids.items[0..2] }}, path[0..1]).held); + try expect(w.store.keepParts(.of(&prompt), 12, null, l.ids.items[0..3], &.{.{ .base = 0, .ids = win.ids.items[0..4] }}, path[0..3]).held); + try expectEqual(@as(u32, 2), w.pool.holders(win.ids.items[0])); // named by both entries, held once by the store +} + +test "reclaim frees a copied part while its path and retained part are a stream's; the stream keeps its pages" { + var w: World = .{ .pool = undefined }; + try w.init(&both, 32); + defer w.deinit(); + var l = try w.lease(2); + defer w.pool.finish(&l); + var win = try w.lease(3); + defer w.pool.finish(&win); + var path: [2]u32 = undefined; + const r = [_]radix.Retained{ .{ .base = 0, .ids = win.ids.items[0..3] }, .{ .base = 0, .ids = &.{} } }; + try expect(w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &r, &path).held); + const before = w.pool.available(); + try expect(w.store.reclaim(before + 3)); // the copied part's 3 pages (and the partial copy) are store-only + try expect(w.pool.available() >= before + 3); + try expectEqual(@as(usize, 0), w.store.tree.snaps); + for (l.ids.items) |id| try expect(w.pool.holders(id) >= 1); // the stream's pages stay + for (win.ids.items) |id| try expectEqual(@as(u32, 1), w.pool.holders(id)); +} + +test "a restore that takes pool pages runs reclaim, which never evicts the entry being restored" { + var w: World = .{ .pool = undefined }; + try w.init(&copied_state, 8); + defer w.deinit(); + var l = try w.lease(2); + w.f.live[0] = 5; + var path: [2]u32 = undefined; + try expect(w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &.{.{ .base = 0, .ids = &.{} }}, &path).held); + w.pool.finish(&l); + w.f.hold_on_restore = w.pool.available() + 1; // one more than free: the pool asks the store + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const plan = try w.store.begin(testing.allocator, .of(&prompt), 14, &.{}, null, &adopt); + defer testing.allocator.free(plan.marks); + try expectEqual(@as(u32, 0), plan.from); // the restore could not get its pages: a prefill from the start + try expectEqual(@as(?u32, null), w.f.restored[0]); + try expectEqual(@as(usize, 0), adopt.items.len); +} + +test "parts need hooks, a pool that holds pages, and rows for a retained group" { + var pool = try Pool.init(testing.allocator, 4, page_bytes); + defer pool.deinit(); + var f: Family = .{ .pool = &pool }; + const pages = radix.pagesOf(Pool, &pool, page_bytes); + try testing.expectError(error.PartsNeedHooks, radix.Store.initGroups(testing.allocator, f.snapshots(), pages, rules, .{ .parts = &copied_state, .legacy = false }, null)); + var no_hold = pages; + var vt = pages.vtable.*; + vt.hold = null; + no_hold.vtable = &vt; + try testing.expectError(error.PartsNeedHold, radix.Store.initGroups(testing.allocator, f.snapshots(), no_hold, rules, .{ .parts = &copied_state, .legacy = false }, f.hooks())); + const zero = [_]radix.Parts.PartSpec{.{ .kind = .window, .make = .retained }}; + try testing.expectError(error.RetainedRows, radix.Store.initGroups(testing.allocator, f.snapshots(), pages, rules, .{ .parts = &zero, .legacy = false }, f.hooks())); +} + +test "an allocation failing at each step of a keep with parts: before the path, nothing changes; after it, the path stays alone" { + var k: usize = 0; + while (true) : (k += 1) { + var pool = try Pool.init(testing.allocator, 32, page_bytes); + defer pool.deinit(); + var f: Family = .{ .pool = &pool }; + var fa = std.testing.FailingAllocator.init(testing.allocator, .{}); + var store = try radix.Store.initGroups(fa.allocator(), f.snapshots(), radix.pagesOf(Pool, &pool, page_bytes), rules, .{ .parts = &both, .legacy = false }, f.hooks()); + store.limit(.{ .pages = 32, .snaps = 8 }); + var l: Lease = .{}; + try pool.promise(&l, 2); + try pool.cover(&l, 2, null); + var win: Lease = .{}; + try pool.promise(&win, 3); + try pool.cover(&win, 3, null); + const r = [_]radix.Retained{ .{ .base = 0, .ids = win.ids.items }, .{ .base = 0, .ids = &.{} } }; + var path: [2]u32 = undefined; + fa.fail_index = fa.alloc_index + k; + const kept = store.keepParts(.of(&prompt), 8, null, l.ids.items, &r, &path); + const failed = fa.has_induced_failure; + if (failed) { + try expect(!kept.held); + try expectEqual(@as(usize, 0), store.tree.snaps); + const path_pages: usize = if (kept.shared == 2) 2 else 0; // after the insert the path is the store's + try expectEqual(path_pages, store.tree.heldPages()); + for (win.ids.items) |id| try expectEqual(@as(u32, 1), pool.holders(id)); // no part outlived the failure + } else try expect(kept.held); + pool.finish(&l); + pool.finish(&win); + store.deinit(); + try expect(pool.idle() and pool.available() == pool.total()); + if (!failed) break; + } + try expect(k > 3); // the insert, the parts, the copies and the prompt each failed once +} + +test "a retained group without its pages through the mark is refused, not counted" { + var w: World = .{ .pool = undefined }; + try w.init(&window3, 16); + defer w.deinit(); + var l = try w.lease(2); + defer w.pool.finish(&l); + var path: [2]u32 = undefined; + try expect(!w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &.{.{ .base = 3, .ids = &.{} }}, &path).held); + try expectEqual(@as(usize, 0), w.store.tree.heldPages()); +} + +test "a keep whose parts need more than the free pages reclaims, but never the path it is building" { + var w: World = .{ .pool = undefined }; + try w.init(&copied_state, 6); + defer w.deinit(); + var l = try w.lease(2); + var path: [2]u32 = undefined; + w.f.fail_copy = true; // a failed checkpoint leaves its path alone in the tree + try expect(!w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &.{.{ .base = 0, .ids = &.{} }}, &path).held); + w.pool.finish(&l); // the path is now the store's alone: idle, reclaimable + var other = try w.lease(2); // 6 pages: 2 the path's, 2 this stream's, 2 free; the part needs 3 + defer w.pool.finish(&other); + w.f.fail_copy = false; + const asked = w.reclaims; + const kept = w.store.keepParts(.of(&prompt), 8, null, other.ids.items, &.{.{ .base = 0, .ids = &.{} }}, &path); + try expect(w.reclaims > asked); // the part's hold ran the reclaim chain + try expect(!kept.held); // the only reclaimable pages were the path being built: refused, not pruned + try expectEqual(@as(usize, 2), w.store.tree.heldPages()); // the path is still the tree's + try expectEqual(@as(usize, 2), kept.shared); +} + +test "a plan's retained ids are its own: an eviction between begin and adoption leaves them valid" { + var w: World = .{ .pool = undefined }; + try w.init(&both, 32); + defer w.deinit(); + var l = try w.lease(2); + var win = try w.lease(3); + var path: [2]u32 = undefined; + const r = [_]radix.Retained{ .{ .base = 0, .ids = win.ids.items }, .{ .base = 0, .ids = &.{} } }; + try expect(w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &r, &path).held); + const whole = [_]u32{ win.ids.items[0], win.ids.items[1] }; + w.pool.finish(&l); + w.pool.finish(&win); + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const plan = try w.store.begin(testing.allocator, .of(&prompt), 14, &.{}, null, &adopt); + defer plan.deinit(testing.allocator); + w.store.limit(.{ .pages = 32, .snaps = 0 }); // the entry goes before the stream maps anything + try expectEqualSlices(u32, &whole, plan.parts[0].ids); + for (plan.parts[0].ids) |id| try expectEqual(@as(u32, 1), w.pool.holders(id)); // the stream's holds alone + for (plan.parts[0].ids) |id| w.pool.release(id); + w.store.finish(null, plan.copies); + for (adopt.items) |id| w.pool.release(id); +} + +test "a begin whose hand-out allocation fails gives back the pages it adopted" { + var w: World = .{ .pool = undefined }; + try w.init(&window3, 32); + defer w.deinit(); + var l = try w.lease(2); + var win = try w.lease(3); + var path: [2]u32 = undefined; + try expect(w.store.keepParts(.of(&prompt), 8, null, l.ids.items, &.{.{ .base = 0, .ids = win.ids.items }}, &path).held); + w.pool.finish(&l); + w.pool.finish(&win); + var k: usize = 0; + while (true) : (k += 1) { + var fa = std.testing.FailingAllocator.init(testing.allocator, .{ .fail_index = k }); + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const census = w.pool.census(); + const plan = w.store.begin(fa.allocator(), .of(&prompt), 14, &.{}, null, &adopt) catch { + try expectEqual(@as(usize, 0), adopt.items.len); + try expectEqual(census, w.pool.census()); // every hold this begin took went back + continue; + }; + for (adopt.items) |id| w.pool.release(id); + for (plan.parts[0].ids) |id| w.pool.release(id); + w.store.finish(null, plan.copies); + plan.deinit(fa.allocator()); + break; + } + try expect(k >= 3); // the marks, the parts, their ids and the copies each failed once +} diff --git a/zig/src/core/prompt_cache.zig b/zig/src/core/prompt_cache.zig index ec73923dc58..54d945e3dd0 100644 --- a/zig/src/core/prompt_cache.zig +++ b/zig/src/core/prompt_cache.zig @@ -72,6 +72,8 @@ pub const Pages = struct { holders: *const fn (ptr: *anyopaque, id: u32) u32, /// Pages free and not promised to a stream. available: *const fn (ptr: *anyopaque) usize, + /// `out.len` free pages the store holds itself, reclaiming first (a checkpoint's copies); null: none. + hold: ?*const fn (ptr: *anyopaque, out: []u32) bool = null, }; }; /// What a family's states depend on beyond their position. diff --git a/zig/src/core/prompt_radix.zig b/zig/src/core/prompt_radix.zig index 3376ad0adaa..0e0249607ff 100644 --- a/zig/src/core/prompt_radix.zig +++ b/zig/src/core/prompt_radix.zig @@ -4,6 +4,7 @@ const std = @import("std"); const pc = @import("prompt_cache.zig"); const tree_mod = @import("prompt_radix_tree.zig"); const parts = @import("prompt_radix_parts.zig"); +const make_mod = @import("prompt_radix_make.zig"); const Allocator = std.mem.Allocator; pub const Saved = tree_mod.Saved; @@ -11,6 +12,7 @@ pub const Snapshots = tree_mod.Snapshots; pub const Pages = tree_mod.Pages; pub const Counts = tree_mod.Counts; pub const Node = tree_mod.Node; +pub const Entry = tree_mod.Entry; pub const Tree = tree_mod.Tree; pub const Groups = parts.Groups; @@ -20,7 +22,7 @@ pub const Rules = pc.Rules; /// What the tree may hold: pages, and snapshots. pub const Limits = struct { pages: usize = 0, snaps: usize = 0 }; -/// An owner's pages as the store holds them: `T` has retain, release, holders and available (kv_pool.Pool). +/// An owner's pages as the store holds them: `T` has retain, release, holders, available and maybe hold (kv_pool.Pool). pub fn pagesOf(comptime T: type, owner: *T, bytes: u64) Pages { const V = struct { fn of(p: *anyopaque) *T { @@ -38,11 +40,35 @@ pub fn pagesOf(comptime T: type, owner: *T, bytes: u64) Pages { fn available(p: *anyopaque) usize { return of(p).available(); } + fn hold(p: *anyopaque, out: []u32) bool { + return of(p).hold(out); + } }; - return .{ .ptr = owner, .bytes = bytes, .vtable = &.{ .retain = V.retain, .release = V.release, .holders = V.holders, .available = V.available } }; + const has_hold = @hasDecl(T, "hold"); + return .{ .ptr = owner, .bytes = bytes, .vtable = &.{ .retain = V.retain, .release = V.release, .holders = V.holders, .available = V.available, .hold = if (has_hold) V.hold else null } }; } -pub const Plan = pc.Plan; +/// Where a pass resumes and the marks it keeps at (in the caller's allocator), and its retained parts to map (C3). +pub const Plan = struct { + from: u32 = 0, + marks: []const u32 = &.{}, + /// Each retained group's whole pages, held once for the stream: it adopts each at `base + i` (or releases it). + parts: []parts.Retained = &.{}, + /// Partial pages to copy into pages of the stream's promise; `finish` ends them. + copies: []parts.Copy = &.{}, + + /// The plan's slices back to `a` (the ids' holds are the stream's to adopt or release). + pub fn deinit(p: Plan, a: Allocator) void { + a.free(p.marks); + for (p.parts) |r| a.free(r.ids); + a.free(p.parts); + a.free(p.copies); + } +}; +pub const Parts = parts; +pub const PartHooks = parts.PartHooks; +pub const Retained = parts.Retained; +pub const Copy = parts.Copy; /// What a pass asks the store about: the words the tree compares, and what reuse must also match (enforced from C6). pub const Query = struct { @@ -75,18 +101,31 @@ pub const Store = struct { limits: Limits = .{}, /// The tokens of recently planned shared cuts, hashed (at most SHARED_KEYS). shared_keys: std.ArrayList(u64) = .empty, + /// The checkpoint groups (C3) and the family's hooks for them; none for a store of legacy snapshots. + decl: make_mod.Decl = .{}, + /// Whether an entry holds a legacy snapshot. + legacy: bool = true, const SHARED_KEYS = 64; pub fn init(gpa: Allocator, family: Snapshots, pages: Pages, rules: Rules) !Store { - return initGroups(gpa, family, pages, rules, .{}); + return initGroups(gpa, family, pages, rules, .{}, null); } - /// `init` for a family that declares its cache groups (C3, #2444): `P` prefix groups page-major on the path. - pub fn initGroups(gpa: Allocator, family: Snapshots, pages: Pages, rules: Rules, groups: Groups) !Store { + /// `init` with the family's cache groups (C3, #2444) and its part hooks; `groups`' slices outlive the store. + pub fn initGroups(gpa: Allocator, family: Snapshots, pages: Pages, rules: Rules, groups: Groups, hooks: ?PartHooks) !Store { if (rules.page == 0 or (rules.planned and (rules.grid == 0 or rules.grid % rules.page != 0))) return error.BadPromptRadixRules; for (groups.prefix_rows) |rows| if (rows != rules.page) return error.UnequalPrefixRows; // one key a page for all - return .{ .tree = try Tree.init(gpa, family, pages, rules.page, rules.lookahead, groups.prefixCount()), .rules = rules }; + if (groups.parts.len > 0 and hooks == null) return error.PartsNeedHooks; + if (groups.parts.len > 0 and pages.vtable.hold == null) return error.PartsNeedHold; // copies the store holds + for (groups.parts) |spec| if (spec.make == .retained and spec.rows == 0) return error.RetainedRows; + if (groups.none()) return error.PageOnlyStore; // a store with no checkpoint at all: not yet served + return .{ + .tree = try Tree.init(gpa, family, pages, rules.page, rules.lookahead, groups.prefixCount()), + .rules = rules, + .decl = .{ .specs = groups.parts, .hooks = hooks }, + .legacy = groups.legacy, + }; } /// Page ids a logical page: one per prefix group. @@ -143,19 +182,48 @@ pub const Store = struct { } s.reserve(prompt, found, marks_); var from: u32 = 0; + var out: Plan = .{ .marks = marks_ }; if (found) |n| { - const family = s.tree.family; - const ok = if (family.vtable.restore(family.ptr, owner, n.entry.?.saved)) |_| true else |err| blk: { + const ok = if (s.restore(owner, n)) |_| true else |err| blk: { note("restoring {d} tokens failed ({s}); prefilling from the start", .{ n.end, @errorName(err) }); break :blk false; }; if (ok) { + const had = adopt.items.len; try s.tree.take(n, adopt); from = n.end; + if (n.entry.?.parts.len > 0) { + const handed = s.decl.handOut(&s.tree, a, n.end, n.entry.?.parts) catch |err| { + for (adopt.items[had..]) |id| s.tree.pages.vtable.release(s.tree.pages.ptr, id); // this begin's + adopt.shrinkRetainingCapacity(had); + return err; + }; + out.parts = handed.parts; + out.copies = handed.copies; + } } s.resumed(n, prompt, ok); } - return .{ .from = from, .marks = marks_ }; + out.from = from; + return out; + } + + /// The entry at `n` into `owner`, snapshot then copied parts; pinned, so reclaim cannot evict its own source. + fn restore(s: *Store, owner: ?*anyopaque, n: *Node) !void { + const e = &n.entry.?; + n.pins += 1; + defer n.pins -= 1; + const family = s.tree.family; + if (e.saved) |saved| try family.vtable.restore(family.ptr, owner, saved); + try s.decl.restore(owner, e.parts); + } + + /// A resume's copies are done (the stream mapped its own pages from them): each source's hold goes back. + pub fn finish(s: *Store, owner: ?*anyopaque, copies: []const Copy) void { + if (copies.len == 0) return; + const h = s.decl.hooks.?; + h.vtable.fence(h.ptr, owner); // the stream's copies out of the sources complete first + for (copies) |c| s.tree.pages.vtable.release(s.tree.pages.ptr, c.src); } /// A restored state: it is now the prompt's (a failed restore drops it). @@ -251,10 +319,16 @@ pub const Store = struct { /// At `at` the tree takes the ids of `mine` (page-major, `per` a page) it lacks, and the snapshot; `path` its ids. pub fn keep(s: *Store, q: Query, at: u32, owner: ?*anyopaque, mine: []const u32, path: []u32) Kept { + return s.keepParts(q, at, owner, mine, &.{}, path); + } + + /// `keep` with the retained groups' pages through `at` (one entry a declared part; copied ones empty). + pub fn keepParts(s: *Store, q: Query, at: u32, owner: ?*anyopaque, mine: []const u32, retained: []const Retained, path: []u32) Kept { const prompt = q.key; const t = &s.tree; const np = at / s.rules.page; // logical pages below `at` if (!s.on() or at == 0 or at % s.step() != 0 or at + s.rules.lookahead > prompt.len or mine.len != np * s.per() or path.len != mine.len) return .{ .held = false, .shared = 0 }; + if (!s.decl.fits(at, retained)) return .{ .held = false, .shared = 0 }; t.clock += 1; const tokens = prompt[0 .. at + s.rules.lookahead]; // a page past which the lookahead differs is another page const flagged = std.mem.indexOfScalar(u64, s.shared_keys.items, sharedKey(tokens)) != null; @@ -268,7 +342,8 @@ pub const Store = struct { t.touch(last.?); return .{ .held = true, .shared = mine.len }; } - if (t.heldPages() + (np - got) * s.per() <= s.limits.pages and t.snaps + 1 <= s.limits.snaps) break; + const ids = (np - got) * s.per() + s.decl.ids(t, at, retained); // every id the keep adds to the holds + if (t.heldPages() + ids <= s.limits.pages and t.snaps + 1 <= s.limits.snaps) break; const v = s.victim(prompt, null) orelse { t.counts.refused += 1; note("kept nothing at {d} tokens: no room", .{at}); @@ -277,13 +352,22 @@ pub const Store = struct { t.evict(v); } const node = t.insert(tokens, mine, path) catch |err| return s.fail(at, err, 0); + node.pins += 1; // the parts' holds may reclaim: not this node, nor its path + defer node.pins -= 1; + // the path stays whatever follows; the entry is published only whole + const made = s.decl.make(t, owner, at, retained) catch |err| return s.fail(at, err, mine.len); const family = t.family; - const saved = family.vtable.save(family.ptr, owner, at) catch |err| return s.fail(at, err, mine.len); + const saved: ?Saved = if (s.legacy) family.vtable.save(family.ptr, owner, at) catch |err| { + t.dropParts(made); + return s.fail(at, err, mine.len); + } else null; const last = t.gpa.dupe(u32, prompt) catch { - family.vtable.drop(family.ptr, saved); + if (saved) |sv| family.vtable.drop(family.ptr, sv); + t.dropParts(made); return s.fail(at, error.OutOfMemory, mine.len); }; - node.entry = .{ .saved = saved, .bytes = family.vtable.bytes(family.ptr, at), .born = @intCast(prompt.len), .used = t.clock, .last = last, .shared = flagged }; + const bytes = if (s.legacy) family.vtable.bytes(family.ptr, at) else 0; + node.entry = .{ .saved = saved, .parts = made, .bytes = bytes, .born = @intCast(prompt.len), .used = t.clock, .last = last, .shared = flagged }; t.snaps += 1; t.counts.kept += 1; return .{ .held = true, .shared = mine.len }; @@ -318,7 +402,7 @@ pub const Store = struct { } fn pick(n: *Node, prompt: []const u32, skip: ?*Node, moved: *?*Node, oldest: *?*Node) void { - if (n.entry) |e| if (n != skip) { + if (n.entry) |e| if (n != skip and n.pins == 0) { // a state its conversation moved past goes first; a shared cut or a branch point only as the least valuable const gone = !e.shared and !e.branched and ((prompt.len > e.born and prompt.len > n.end and extends(n, prompt)) or superseded(n, e.born)); if (gone and (moved.* == null or e.used < moved.*.?.entry.?.used)) moved.* = n; @@ -338,10 +422,33 @@ pub const Store = struct { /// Evicts until `want` pool pages are free and not promised; false when the tree cannot give that many. pub fn reclaim(s: *Store, want: usize) bool { const pages = s.tree.pages; + if (s.decl.specs.len > 0) while (pages.vtable.available(pages.ptr) < want) { // parts first: their path may be in use + s.tree.evict(s.partVictim() orelse break); + }; while (pages.vtable.available(pages.ptr) < want) if (!s.tree.evictLeaf()) return false; return true; } + /// The least valuable entry with a part page only the store holds, never a pinned one. + fn partVictim(s: *const Store) ?*Node { + var best: ?*Node = null; + s.pickParts(s.tree.root, &best); + return best; + } + + fn pickParts(s: *const Store, n: *Node, best: *?*Node) void { + if (n.entry) |e| if (n.pins == 0 and s.freesPart(e) and (best.* == null or worse(n, best.*.?))) { + best.* = n; + }; + for (n.kids.items) |k| s.pickParts(k, best); + } + + fn freesPart(s: *const Store, e: Entry) bool { + const pages = s.tree.pages; + for (e.parts) |part| for (part.ids) |id| if (pages.vtable.holders(pages.ptr, id) == 1) return true; + return false; + } + /// One log line after a prompt pass: where it resumed, how many states it kept, and what the tree holds. pub fn report(s: *const Store, prompt: usize, from: u32, kept: u64) void { if (@import("builtin").is_test) return; diff --git a/zig/src/core/prompt_radix_groups_test.zig b/zig/src/core/prompt_radix_groups_test.zig index 827672bbb12..d28f4d5084b 100644 --- a/zig/src/core/prompt_radix_groups_test.zig +++ b/zig/src/core/prompt_radix_groups_test.zig @@ -57,7 +57,7 @@ const Fake = struct { const two: radix.Groups = .{ .prefix_rows = &.{ 4, 4 } }; fn open(f: *Fake, limit_ids: usize) !radix.Store { - var s = try radix.Store.initGroups(testing.allocator, f.family(), f.pages(), .{ .page = 4, .lookahead = 0, .min_gap = 1, .min_prompt = 0 }, two); + var s = try radix.Store.initGroups(testing.allocator, f.family(), f.pages(), .{ .page = 4, .lookahead = 0, .min_gap = 1, .min_prompt = 0 }, two, null); s.limit(.{ .pages = limit_ids, .snaps = 4 }); return s; } @@ -121,5 +121,5 @@ test "two prefix groups: the page limit counts ids, so a keep needing more ids t test "prefix groups with other rows a page than the store's are refused when it opens" { var f: Fake = .{}; - try testing.expectError(error.UnequalPrefixRows, radix.Store.initGroups(testing.allocator, f.family(), f.pages(), .{ .page = 4, .min_prompt = 0 }, .{ .prefix_rows = &.{ 4, 8 } })); + try testing.expectError(error.UnequalPrefixRows, radix.Store.initGroups(testing.allocator, f.family(), f.pages(), .{ .page = 4, .min_prompt = 0 }, .{ .prefix_rows = &.{ 4, 8 } }, null)); } diff --git a/zig/src/core/prompt_radix_make.zig b/zig/src/core/prompt_radix_make.zig new file mode 100644 index 00000000000..4a6803dba4f --- /dev/null +++ b/zig/src/core/prompt_radix_make.zig @@ -0,0 +1,156 @@ +//! A checkpoint's parts (C3, #2444): made at a keep, restored at a resume, all or nothing, the copies fenced. +const std = @import("std"); +const parts = @import("prompt_radix_parts.zig"); +const tree_mod = @import("prompt_radix_tree.zig"); +const Allocator = std.mem.Allocator; + +const Part = parts.Part; +const Retained = parts.Retained; +const Copy = parts.Copy; +const Tree = tree_mod.Tree; + +/// A store's part declaration: its groups' specs, and the family's hooks for them. +pub const Decl = struct { + specs: []const parts.PartSpec = &.{}, + hooks: ?parts.PartHooks = null, + + /// Logical pages of a retained group below or at `at`: whole ones, and one more when its rows fall short. + fn spans(rows: u32, at: u32) u32 { + return (at + rows - 1) / rows; + } + + /// Whether `retained` gives each retained group its pages through `at` (one entry a part, in order). + pub fn fits(d: Decl, at: u32, retained: []const Retained) bool { + if (retained.len != d.specs.len) return false; + for (d.specs, retained) |spec, r| switch (spec.make) { + .copied => if (r.ids.len != 0) return false, + .retained => if (r.ids.len == 0 or r.base + r.ids.len != spans(spec.rows, at)) return false, + }; + return true; + } + + /// Ids a keep at `at` adds to the tree's holds: copied parts' pages, retained ids not held yet, partial copies. + pub fn ids(d: Decl, t: *const Tree, at: u32, retained: []const Retained) usize { + var n: usize = 0; + for (d.specs, retained, 0..) |spec, r, g| switch (spec.make) { + .copied => n += copiedPages(d, t, @intCast(g), at), + .retained => { + const partial = at % spec.rows != 0; + const whole = r.ids[0 .. r.ids.len - @intFromBool(partial)]; + for (whole) |id| n += @intFromBool(!t.holds.contains(id)); + n += @intFromBool(partial); + }, + }; + return n; + } + + fn copiedPages(d: Decl, t: *const Tree, g: u32, at: u32) usize { + const h = d.hooks.?; + const bytes = h.vtable.bytes(h.ptr, g, at); + return @intCast((bytes + t.pages.bytes - 1) / t.pages.bytes); + } + + /// The parts of a checkpoint at `at`, each id named in the tree's holds; on failure, what it took goes back. + pub fn make(d: Decl, t: *Tree, owner: ?*anyopaque, at: u32, retained: []const Retained) ![]Part { + if (d.specs.len == 0) return &.{}; + const h = d.hooks.?; + const out = try t.gpa.alloc(Part, d.specs.len); + var made: usize = 0; + errdefer { + h.vtable.fence(h.ptr, owner); // queued copies into the pages complete before they go back + t.dropEach(out[0..made]); + t.gpa.free(out); + } + try t.holds.reserve(t.gpa, d.ids(t, at, retained)); + for (d.specs, retained, 0..) |spec, r, gi| { + const g: u32 = @intCast(gi); + switch (spec.make) { + .copied => { + const n = copiedPages(d, t, g, at); + const held = try t.gpa.alloc(u32, n); + if (!hold(t, held)) { + t.gpa.free(held); + return error.NoCachePages; + } + for (held) |id| t.holds.adoptHeldAssumeCapacity(id); + out[made] = .{ .base = 0, .ids = held, .bytes = h.vtable.bytes(h.ptr, g, at) }; + made += 1; + try h.vtable.copy(h.ptr, owner, g, at, out[made - 1]); + }, + .retained => { + const rows_left = at % spec.rows; + const whole = r.ids.len - @intFromBool(rows_left != 0); + var copy: [1]u32 = undefined; // the partial page: a copy the store holds + if (rows_left != 0 and !hold(t, ©)) return error.NoCachePages; + const own = t.gpa.alloc(u32, r.ids.len) catch |err| { + if (rows_left != 0) t.pages.vtable.release(t.pages.ptr, copy[0]); + return err; + }; + for (r.ids[0..whole], own[0..whole]) |id, *o| { + t.holds.shareAssumeCapacity(t.pages, id); + o.* = id; + } + if (rows_left != 0) { + own[whole] = copy[0]; + t.holds.adoptHeldAssumeCapacity(copy[0]); + } + out[made] = .{ .base = r.base, .ids = own, .bytes = 0 }; + made += 1; + if (rows_left != 0) try h.vtable.copy_rows(h.ptr, owner, g, r.ids[whole], copy[0], rows_left); + }, + } + } + h.vtable.fence(h.ptr, owner); // every copy into the parts is done: they are the checkpoint's from now on + return out; + } + + fn hold(t: *Tree, out: []u32) bool { + const f = t.pages.vtable.hold orelse return false; + return f(t.pages.ptr, out); + } + + /// The copied parts of `e` into the owner's live state, fenced; the store keeps their pages. + pub fn restore(d: Decl, owner: ?*anyopaque, ps: []const Part) !void { + if (d.specs.len == 0) return; + const h = d.hooks.?; + defer h.vtable.fence(h.ptr, owner); // the copies out of the store's pages are done, failed or not + for (d.specs, ps, 0..) |spec, part, g| if (spec.make == .copied) try h.vtable.restore(h.ptr, owner, @intCast(g), part); + } + + /// The retained parts a resume at `at` maps: whole pages held once for the stream, partial ones as copies to make. + pub fn handOut(d: Decl, t: *Tree, a: Allocator, at: u32, ps: []const Part) !struct { parts: []Retained, copies: []Copy } { + var nr: usize = 0; + var nc: usize = 0; + for (d.specs) |spec| if (spec.make == .retained) { + nr += 1; + nc += @intFromBool(at % spec.rows != 0); + }; + const out = try a.alloc(Retained, nr); + var owned: usize = 0; // retained entries whose ids `a` owns + errdefer { + for (out[0..owned]) |r| a.free(r.ids); + a.free(out); + } + const copies = try a.alloc(Copy, nc); + errdefer a.free(copies); + nr = 0; + for (d.specs, ps) |spec, part| if (spec.make == .retained) { + const whole = part.ids.len - @intFromBool(at % spec.rows != 0); + out[nr] = .{ .base = part.base, .ids = try a.dupe(u32, part.ids[0..whole]) }; // the plan's own: eviction frees part.ids + nr += 1; + owned = nr; + }; + nr = 0; + nc = 0; + for (d.specs, ps, 0..) |spec, part, g| if (spec.make == .retained) { + const rows_left = at % spec.rows; + const whole = part.ids.len - @intFromBool(rows_left != 0); + for (part.ids) |id| t.pages.vtable.retain(t.pages.ptr, id); // the stream's hold on each, until it adopts or finishes + nr += 1; + if (rows_left == 0) continue; + copies[nc] = .{ .group = @intCast(g), .index = at / spec.rows, .src = part.ids[whole], .rows = rows_left }; + nc += 1; + }; + return .{ .parts = out, .copies = copies }; + } +}; diff --git a/zig/src/core/prompt_radix_parts.zig b/zig/src/core/prompt_radix_parts.zig index 8985549a2f2..39f6991a7a9 100644 --- a/zig/src/core/prompt_radix_parts.zig +++ b/zig/src/core/prompt_radix_parts.zig @@ -41,6 +41,25 @@ pub const Retained = struct { base: u32, ids: []const u32 }; /// A copy a stream makes from a held page into one of its own: rows [0, rows) of `src` to the group's page `index`. pub const Copy = struct { group: u32, index: u32, src: u32, rows: u32 }; +/// What the store asks a family that declares parts: the bytes, and the device copies (each on the owner's stream). +pub const PartHooks = struct { + ptr: *anyopaque, + vtable: *const VTable, + + pub const VTable = struct { + /// A copied part's bytes at `at` tokens. + bytes: *const fn (ptr: *anyopaque, group: u32, at: u32) u64, + /// The owner's live bytes of `group` after `at` tokens into the held pages of `dst`. + copy: *const fn (ptr: *anyopaque, owner: ?*anyopaque, group: u32, at: u32, dst: Part) anyerror!void, + /// The held pages of `src` into the owner's live `group` (the next prompt chunk starts at its position). + restore: *const fn (ptr: *anyopaque, owner: ?*anyopaque, group: u32, src: Part) anyerror!void, + /// Rows [0, rows) of every layer of `group` from page `src` into page `dst`. + copy_rows: *const fn (ptr: *anyopaque, owner: ?*anyopaque, group: u32, src: u32, dst: u32, rows: u32) anyerror!void, + /// Every copy queued for `owner` has completed, on every rank that runs it. + fence: *const fn (ptr: *anyopaque, owner: ?*anyopaque) void, + }; +}; + /// The store's references: one pool reference a page, however many nodes and entries name it. pub const Holds = struct { counts: std.AutoHashMapUnmanaged(u32, u32) = .empty, diff --git a/zig/src/core/prompt_radix_tree.zig b/zig/src/core/prompt_radix_tree.zig index eee106ebf31..cd9b50d1e3f 100644 --- a/zig/src/core/prompt_radix_tree.zig +++ b/zig/src/core/prompt_radix_tree.zig @@ -13,7 +13,10 @@ pub const Counts = pc.Counts; /// The state kept at a node's end. pub const Entry = struct { - saved: Saved, + /// The family's legacy snapshot; null when it declares none (its state is in `parts`). + saved: ?Saved, + /// The checkpoint groups' pages (C3), each id named in the tree's `holds`. + parts: []parts.Part = &.{}, bytes: u64, /// The length of the prompt that kept it: a later turn's is longer. born: u32, @@ -39,6 +42,8 @@ pub const Node = struct { kids: std.ArrayList(*Node) = .empty, entry: ?Entry = null, used: u64 = 0, + /// A keep or a restore in progress here: no eviction or prune takes this node, its entry or its path. + pins: u32 = 0, }; pub const Tree = struct { @@ -83,10 +88,7 @@ pub const Tree = struct { fn drop(t: *Tree, n: *Node) void { for (n.kids.items) |k| t.drop(k); n.kids.deinit(t.gpa); - if (n.entry) |e| { - t.family.vtable.drop(t.family.ptr, e.saved); - t.gpa.free(e.last); - } + if (n.entry) |e| t.dropEntry(e); for (n.pages) |id| t.holds.drop(t.pages, id); t.gpa.free(n.tokens); t.gpa.free(n.pages); @@ -175,9 +177,11 @@ pub const Tree = struct { const c = t.child(node, tokens[pi * t.page ..]) orelse { const leaf = try t.gpa.create(Node); errdefer t.gpa.destroy(leaf); - leaf.* = .{ .parent = node, .end = @intCast(np * t.page), .tokens = try t.gpa.dupe(u32, tokens[pi * t.page ..]), .pages = try t.gpa.dupe(u32, mine[pi * t.per ..]) }; - errdefer t.gpa.free(leaf.tokens); - errdefer t.gpa.free(leaf.pages); + const own_tokens = try t.gpa.dupe(u32, tokens[pi * t.page ..]); + errdefer t.gpa.free(own_tokens); + const own_pages = try t.gpa.dupe(u32, mine[pi * t.per ..]); + errdefer t.gpa.free(own_pages); + leaf.* = .{ .parent = node, .end = @intCast(np * t.page), .tokens = own_tokens, .pages = own_pages }; try t.holds.reserve(t.gpa, leaf.pages.len); // before the leaf is linked: what follows cannot fail try node.kids.append(t.gpa, leaf); for (leaf.pages) |id| t.holds.shareAssumeCapacity(t.pages, id); @@ -221,11 +225,32 @@ pub const Tree = struct { return up; } + /// An entry's snapshot, parts and prompt go. + fn dropEntry(t: *Tree, e: Entry) void { + if (e.saved) |saved| t.family.vtable.drop(t.family.ptr, saved); + t.dropParts(e.parts); + t.gpa.free(e.last); + } + + /// Parts made by a keep go back: each id's name in `holds`, then the slices. + pub fn dropParts(t: *Tree, ps: []parts.Part) void { + t.dropEach(ps); + t.gpa.free(ps); + } + + /// Each part's ids go back and its id slice is freed; the caller frees `ps` itself. + pub fn dropEach(t: *Tree, ps: []const parts.Part) void { + for (ps) |part| { + for (part.ids) |id| t.holds.drop(t.pages, id); + t.gpa.free(part.ids); + } + } + /// Frees the state at `n`, and the nodes below nothing keeps. pub fn evict(t: *Tree, n: *Node) void { const e = n.entry orelse return; - t.family.vtable.drop(t.family.ptr, e.saved); - t.gpa.free(e.last); + std.debug.assert(n.pins == 0); // a keep or a restore in progress + t.dropEntry(e); n.entry = null; t.snaps -= 1; t.counts.evicted += 1; @@ -236,7 +261,7 @@ pub const Tree = struct { pub fn prune(t: *Tree, start: *Node) void { var n = start; while (n.parent) |parent| { - if (n.kids.items.len > 0 or n.entry != null) return; + if (n.kids.items.len > 0 or n.entry != null or n.pins > 0) return; const slot = std.mem.indexOfScalar(*Node, parent.kids.items, n).?; _ = parent.kids.swapRemove(slot); t.drop(n); @@ -259,7 +284,7 @@ pub const Tree = struct { } fn weakestLeaf(t: *const Tree, n: *Node, best: *?*Node) void { - if (n.parent != null and n.kids.items.len == 0 and t.idle(n) and (best.* == null or lessValuable(n, best.*.?))) best.* = n; + if (n.parent != null and n.kids.items.len == 0 and n.pins == 0 and t.idle(n) and (best.* == null or lessValuable(n, best.*.?))) best.* = n; for (n.kids.items) |k| t.weakestLeaf(k, best); } From 970b40a1207d01d0df29c08c82ed27c91dbf55b8 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 10 Nov 2026 19:55:28 +0000 Subject: [PATCH 42/78] feat(radix): a page-only store resumes at any matched page edge C3 of #2444 defines the store for a family with no checkpoint at all (pure full attention, as vLLM's prefix cache): it keeps the path's pages with no snapshot slot, resumes at the deepest page edge a prompt repeats, inside a node too, and makes room by freeing idle leaves, never the path a pass resumes. No family on the store is such a family today, so every store user is unchanged; the first one (MiniMax M3, P7) needs no store change. --- zig/src/core/prompt_radix.zig | 51 +++++++++++--- zig/src/core/prompt_radix_groups_test.zig | 86 +++++++++++++++++++++++ zig/src/core/prompt_radix_tree.zig | 47 +++++++++++++ 3 files changed, 175 insertions(+), 9 deletions(-) diff --git a/zig/src/core/prompt_radix.zig b/zig/src/core/prompt_radix.zig index 0e0249607ff..c399486b078 100644 --- a/zig/src/core/prompt_radix.zig +++ b/zig/src/core/prompt_radix.zig @@ -105,6 +105,8 @@ pub const Store = struct { decl: make_mod.Decl = .{}, /// Whether an entry holds a legacy snapshot. legacy: bool = true, + /// No checkpoint at all (pure full attention): any matched page edge resumes, and no entry is kept. + page_only: bool = false, const SHARED_KEYS = 64; @@ -119,12 +121,12 @@ pub const Store = struct { if (groups.parts.len > 0 and hooks == null) return error.PartsNeedHooks; if (groups.parts.len > 0 and pages.vtable.hold == null) return error.PartsNeedHold; // copies the store holds for (groups.parts) |spec| if (spec.make == .retained and spec.rows == 0) return error.RetainedRows; - if (groups.none()) return error.PageOnlyStore; // a store with no checkpoint at all: not yet served return .{ .tree = try Tree.init(gpa, family, pages, rules.page, rules.lookahead, groups.prefixCount()), .rules = rules, .decl = .{ .specs = groups.parts, .hooks = hooks }, .legacy = groups.legacy, + .page_only = groups.none(), }; } @@ -157,7 +159,14 @@ pub const Store = struct { /// Whether anything can be kept. pub fn on(s: *const Store) bool { - return s.limits.pages > 0 and s.limits.snaps > 0; + return s.limits.pages > 0 and (s.page_only or s.limits.snaps > 0); + } + + /// The least valuable thing to free: an entry, or for a store without entries an idle leaf; false when none. + fn evictOne(s: *Store, prompt: []const u32, skip: ?*Node) bool { + if (s.page_only) return s.tree.evictLeaf(); + s.tree.evict(s.victim(prompt, skip) orelse return false); + return true; } /// The deepest node whose state `prompt` resumes exactly: its tokens and lookahead a prefix, a token left. @@ -170,6 +179,7 @@ pub const Store = struct { pub fn begin(s: *Store, a: Allocator, q: Query, history_len: u32, shared: []const u32, owner: ?*anyopaque, adopt: *std.ArrayList(u32)) !Plan { const prompt = q.key; if (!s.on()) return .{ .marks = try a.alloc(u32, 0) }; + if (s.page_only) return s.beginPages(a, prompt, history_len, shared, adopt); const found = s.find(prompt); if (found == null) s.tree.counts.misses += 1; const resume_at: u32 = if (found) |n| n.end else 0; @@ -208,6 +218,25 @@ pub const Store = struct { return out; } + /// `begin` for a store with no checkpoint: the deepest matched page edge, its path adopted, nothing restored. + fn beginPages(s: *Store, a: Allocator, prompt: []const u32, history_len: u32, shared: []const u32, adopt: *std.ArrayList(u32)) !Plan { + const hit: ?tree_mod.Hit = if (prompt.len < s.past() + 1) null else s.tree.deepestPage(prompt, (prompt.len - s.past()) / s.rules.page * s.rules.page); + const from: u32 = if (hit) |h| h.at(&s.tree) else 0; + const planned = try s.marks(a, prompt, from, history_len, shared, &.{}); + const marks_ = try s.fitting(a, planned); + errdefer a.free(marks_); + if (hit) |h| h.node.pins += 1; // the room made for the marks never prunes the path this pass resumes + s.reserve(prompt, null, marks_); + if (hit) |h| h.node.pins -= 1; + const h = hit orelse { + s.tree.counts.misses += 1; + return .{ .marks = marks_ }; + }; + try s.tree.takeTo(h, adopt); + s.tree.counts.hits += 1; + return .{ .from = from, .marks = marks_ }; + } + /// The entry at `n` into `owner`, snapshot then copied parts; pinned, so reclaim cannot evict its own source. fn restore(s: *Store, owner: ?*anyopaque, n: *Node) !void { const e = &n.entry.?; @@ -249,8 +278,8 @@ pub const Store = struct { const last = marks_[marks_.len - 1]; while (true) { const new = (last / s.rules.page - s.tree.have(prompt[0 .. last + s.rules.lookahead])) * s.per(); - if (s.tree.heldPages() + new <= s.limits.pages and s.tree.snaps + marks_.len <= s.limits.snaps) return; - s.tree.evict(s.victim(prompt, from) orelse return); + if (s.tree.heldPages() + new <= s.limits.pages and (s.page_only or s.tree.snaps + marks_.len <= s.limits.snaps)) return; + if (!s.evictOne(prompt, from)) return; } } @@ -279,7 +308,7 @@ pub const Store = struct { var out: std.ArrayList(u32) = .empty; errdefer out.deinit(a); for (marks_) |m| { - if (m / s.rules.page * s.per() <= s.limits.pages and s.limits.snaps > 0) try out.append(a, m) else { + if (m / s.rules.page * s.per() <= s.limits.pages and (s.page_only or s.limits.snaps > 0)) try out.append(a, m) else { s.tree.counts.refused += 1; note("kept nothing at {d} tokens: {d} pages pass the limit of {d}", .{ m, m / s.rules.page * s.per(), s.limits.pages }); } @@ -335,6 +364,7 @@ pub const Store = struct { while (true) { var last: ?*Node = null; const got = t.heldPrefix(tokens, np, path, &last); + if (s.page_only and got == np) return .{ .held = true, .shared = mine.len }; // the path is the tree's already if (got == np and last != null and last.?.end == at and last.?.entry != null) { const e = &last.?.entry.?; e.used = t.clock; // the same state again: no copy @@ -343,15 +373,18 @@ pub const Store = struct { return .{ .held = true, .shared = mine.len }; } const ids = (np - got) * s.per() + s.decl.ids(t, at, retained); // every id the keep adds to the holds - if (t.heldPages() + ids <= s.limits.pages and t.snaps + 1 <= s.limits.snaps) break; - const v = s.victim(prompt, null) orelse { + if (t.heldPages() + ids <= s.limits.pages and (s.page_only or t.snaps + 1 <= s.limits.snaps)) break; + if (!s.evictOne(prompt, null)) { t.counts.refused += 1; note("kept nothing at {d} tokens: no room", .{at}); return .{ .held = false, .shared = got * s.per() }; - }; - t.evict(v); + } } const node = t.insert(tokens, mine, path) catch |err| return s.fail(at, err, 0); + if (s.page_only) { + t.counts.kept += 1; + return .{ .held = true, .shared = mine.len }; + } node.pins += 1; // the parts' holds may reclaim: not this node, nor its path defer node.pins -= 1; // the path stays whatever follows; the entry is published only whole diff --git a/zig/src/core/prompt_radix_groups_test.zig b/zig/src/core/prompt_radix_groups_test.zig index d28f4d5084b..62a7d29848e 100644 --- a/zig/src/core/prompt_radix_groups_test.zig +++ b/zig/src/core/prompt_radix_groups_test.zig @@ -123,3 +123,89 @@ test "prefix groups with other rows a page than the store's are refused when it var f: Fake = .{}; try testing.expectError(error.UnequalPrefixRows, radix.Store.initGroups(testing.allocator, f.family(), f.pages(), .{ .page = 4, .min_prompt = 0 }, .{ .prefix_rows = &.{ 4, 8 } }, null)); } + +const page_only: radix.Groups = .{ .legacy = false }; + +fn openPages(f: *Fake, limit_ids: usize) !radix.Store { + var s = try radix.Store.initGroups(testing.allocator, f.family(), f.pages(), .{ .page = 4, .lookahead = 0, .min_gap = 1, .min_prompt = 0 }, page_only, null); + s.limit(.{ .pages = limit_ids, .snaps = 0 }); // no checkpoint: no snapshot slots needed + return s; +} + +test "a store with no checkpoint keeps pages with no snapshot slot and resumes at any matched page edge" { + var f: Fake = .{}; + var s = try openPages(&f, 16); + defer s.deinit(); + try testing.expect(s.on()); + const kept_prompt = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13 }; + const mine = [_]u32{ 0, 1, 2 }; + f.take(&mine); + var path: [3]u32 = undefined; + try testing.expect(s.keep(.of(&kept_prompt), 12, null, &mine, &path).held); + try testing.expectEqual(@as(usize, 0), s.tree.snaps); + for (mine) |id| f.refs[id] -= 1; + // a prompt that diverges inside the kept node's third page resumes after its second, adopting two pages + const other = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 99, 99, 99, 99, 99 }; + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const plan = try s.begin(testing.allocator, .of(&other), 14, &.{}, null, &adopt); + defer testing.allocator.free(plan.marks); + try testing.expectEqual(@as(u32, 8), plan.from); + try testing.expectEqualSlices(u32, &.{ 0, 1 }, adopt.items); + try testing.expectEqual(@as(?u32, null), f.restored); // nothing to restore + for (adopt.items) |id| f.refs[id] -= 1; + adopt.clearRetainingCapacity(); + // a prompt that ends inside it resumes at the last whole page that leaves a token to prefill + const short = [_]u32{ 1, 2, 3, 4, 5, 6, 7 }; + const plan2 = try s.begin(testing.allocator, .of(&short), 7, &.{}, null, &adopt); + defer testing.allocator.free(plan2.marks); + try testing.expectEqual(@as(u32, 4), plan2.from); + try testing.expectEqualSlices(u32, &.{0}, adopt.items); + for (adopt.items) |id| f.refs[id] -= 1; +} + +test "a store with no checkpoint frees an idle leaf to keep another prompt's pages" { + var f: Fake = .{}; + var s = try openPages(&f, 3); + defer s.deinit(); + const a = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 0 }; + const mine = [_]u32{ 0, 1, 2 }; + f.take(&mine); + var path: [3]u32 = undefined; + try testing.expect(s.keep(.of(&a), 12, null, &mine, &path).held); + for (mine) |id| f.refs[id] -= 1; + const b = [_]u32{ 7, 7, 7, 7, 7, 7, 7, 7, 0 }; + const theirs = [_]u32{ 5, 6 }; + f.take(&theirs); + var path2: [2]u32 = undefined; + try testing.expect(s.keep(.of(&b), 8, null, &theirs, &path2).held); // room: the idle leaf of `a` went + try testing.expect(s.tree.heldPages() <= 3); + for (theirs) |id| f.refs[id] -= 1; +} + +test "a store with no checkpoint making room for a pass's marks never frees the path that pass resumes" { + var f: Fake = .{}; + var s = try openPages(&f, 5); + defer s.deinit(); + const a = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 0 }; + const mine = [_]u32{ 0, 1, 2 }; + f.take(&mine); + var path: [3]u32 = undefined; + try testing.expect(s.keep(.of(&a), 12, null, &mine, &path).held); + for (mine) |id| f.refs[id] -= 1; // the tree's alone: idle, and the oldest leaf + const b = [_]u32{ 7, 7, 7, 7, 7, 7, 7, 7, 0 }; + const theirs = [_]u32{ 5, 6 }; + f.take(&theirs); + var path2: [2]u32 = undefined; + try testing.expect(s.keep(.of(&b), 8, null, &theirs, &path2).held); + for (theirs) |id| f.refs[id] -= 1; // a newer idle leaf: 5 pages held of 5 + const longer = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21 }; + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const plan = try s.begin(testing.allocator, .of(&longer), 21, &.{}, null, &adopt); // its mark at 20 wants 2 more pages + defer testing.allocator.free(plan.marks); + try testing.expectEqual(@as(u32, 12), plan.from); + try testing.expectEqualSlices(u32, &mine, adopt.items); + for (adopt.items) |id| f.refs[id] -= 1; + try testing.expectEqual(@as(u32, 0), f.refs[5]); // the newer leaf went instead +} diff --git a/zig/src/core/prompt_radix_tree.zig b/zig/src/core/prompt_radix_tree.zig index cd9b50d1e3f..81ad18aa5d9 100644 --- a/zig/src/core/prompt_radix_tree.zig +++ b/zig/src/core/prompt_radix_tree.zig @@ -32,6 +32,17 @@ pub const Entry = struct { branched: bool = false, }; +/// A matched position: `pages` logical pages into `node` (1 to its count). +pub const Hit = struct { + node: *Node, + pages: usize, + + /// The tokens below the position. + pub fn at(h: Hit, t: *const Tree) u32 { + return h.node.end - @as(u32, @intCast((t.npages(h.node) - h.pages) * t.page)); + } +}; + pub const Node = struct { parent: ?*Node, /// Tokens up to the end of this node. @@ -130,6 +141,42 @@ pub const Tree = struct { return best; } + /// The deepest page edge `prompt` repeats up to `limit_at`, inside a node too: where a page-only store resumes. + pub fn deepestPage(t: *const Tree, prompt: []const u32, limit_at: usize) ?Hit { + var node = t.root; + var at: usize = 0; + var best: ?Hit = null; + while (at + t.page <= limit_at) { + const c = t.child(node, prompt[at..]) orelse break; + const p = t.common(c, prompt[at..], (limit_at - at) / t.page); + if (p == 0) break; + best = .{ .node = c, .pages = p }; + if (p < t.npages(c)) break; + at += p * t.page; + node = c; + } + return best; + } + + /// The ids of the path to `h`, each held once more for the caller; the nodes marked used. + pub fn takeTo(t: *Tree, h: Hit, adopt: *std.ArrayList(u32)) !void { + var path: std.ArrayList(*Node) = .empty; + defer path.deinit(t.gpa); + var w: ?*Node = h.node.parent; + while (w) |x| : (w = x.parent) if (x.parent != null) try path.append(t.gpa, x); + try adopt.ensureUnusedCapacity(t.gpa, h.at(t) / t.page * t.per); + while (path.pop()) |x| t.hand(x, x.pages, adopt); + t.hand(h.node, h.node.pages[0 .. h.pages * t.per], adopt); + } + + fn hand(t: *Tree, n: *Node, ids: []const u32, adopt: *std.ArrayList(u32)) void { + for (ids) |id| { + t.pages.vtable.retain(t.pages.ptr, id); + adopt.appendAssumeCapacity(id); + } + t.touch(n); + } + /// Logical pages of `tokens` (`np` whole pages and `ahead` more) the tree holds, their ids into `path` if any. pub fn heldPrefix(t: *const Tree, tokens: []const u32, np: usize, path: ?[]u32, last: *?*Node) usize { var node = t.root; From 1ae4fca83b577790a72ecd41bfc785633dd01fc5 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 10 Nov 2026 21:54:14 +0000 Subject: [PATCH 43/78] fix(radix): count part pages when planning a pass's marks A mark whose path and checkpoint parts passed the page limit was planned anyway; making room for it evicted a kept checkpoint before the keep itself was refused. Planning now counts the part pages it knows before the keep (copied parts and partial retained pages), so such a mark is refused up front. Long trailing comments in the lines this branch added move above their lines. --- zig/src/core/kv_pool_parts_test.zig | 18 ++++++++++++++++++ zig/src/core/prompt_radix.zig | 19 +++++++++++++------ zig/src/core/prompt_radix_groups_test.zig | 6 ++++-- zig/src/core/prompt_radix_make.zig | 16 ++++++++++++++-- 4 files changed, 49 insertions(+), 10 deletions(-) diff --git a/zig/src/core/kv_pool_parts_test.zig b/zig/src/core/kv_pool_parts_test.zig index eebc961d2a7..70a567f8997 100644 --- a/zig/src/core/kv_pool_parts_test.zig +++ b/zig/src/core/kv_pool_parts_test.zig @@ -391,3 +391,21 @@ test "a begin whose hand-out allocation fails gives back the pages it adopted" { } try expect(k >= 3); // the marks, the parts, their ids and the copies each failed once } + +test "a mark whose path and parts pass the page limit is refused when planned, so it never evicts a checkpoint first" { + var w: World = .{ .pool = undefined }; + try w.init(&copied_state, 16); + defer w.deinit(); + w.store.limit(.{ .pages = 5, .snaps = 8 }); + var l = try w.lease(1); + var path: [1]u32 = undefined; + try expect(w.store.keepParts(.of(&prompt), 4, null, l.ids.items, &.{.{ .base = 0, .ids = &.{} }}, &path).held); + w.pool.finish(&l); // one prefix page and three copied: 4 of 5 + const other = [_]u32{ 9, 9, 9, 9, 8, 8, 8, 8, 7, 7, 7, 7, 6, 6 }; + var adopt: std.ArrayList(u32) = .empty; + defer adopt.deinit(testing.allocator); + const plan = try w.store.begin(testing.allocator, .of(&other), 14, &.{}, null, &adopt); + defer plan.deinit(testing.allocator); + try expectEqual(@as(usize, 0), plan.marks.len); // its mark at 12 needs 3 path pages and 3 part pages: 6 of 5 + try expectEqual(@as(usize, 1), w.store.tree.snaps); // the kept checkpoint was not evicted for it +} diff --git a/zig/src/core/prompt_radix.zig b/zig/src/core/prompt_radix.zig index c399486b078..aec5202eafb 100644 --- a/zig/src/core/prompt_radix.zig +++ b/zig/src/core/prompt_radix.zig @@ -117,7 +117,8 @@ pub const Store = struct { /// `init` with the family's cache groups (C3, #2444) and its part hooks; `groups`' slices outlive the store. pub fn initGroups(gpa: Allocator, family: Snapshots, pages: Pages, rules: Rules, groups: Groups, hooks: ?PartHooks) !Store { if (rules.page == 0 or (rules.planned and (rules.grid == 0 or rules.grid % rules.page != 0))) return error.BadPromptRadixRules; - for (groups.prefix_rows) |rows| if (rows != rules.page) return error.UnequalPrefixRows; // one key a page for all + // one key a page for all prefix groups + for (groups.prefix_rows) |rows| if (rows != rules.page) return error.UnequalPrefixRows; if (groups.parts.len > 0 and hooks == null) return error.PartsNeedHooks; if (groups.parts.len > 0 and pages.vtable.hold == null) return error.PartsNeedHold; // copies the store holds for (groups.parts) |spec| if (spec.make == .retained and spec.rows == 0) return error.RetainedRows; @@ -276,8 +277,11 @@ pub const Store = struct { fn reserve(s: *Store, prompt: []const u32, from: ?*Node, marks_: []const u32) void { if (marks_.len == 0) return; const last = marks_[marks_.len - 1]; + var parts_ids: usize = 0; // each mark's parts, as far as they are known before the keep + for (marks_) |m| parts_ids += s.decl.planIds(&s.tree, m); while (true) { - const new = (last / s.rules.page - s.tree.have(prompt[0 .. last + s.rules.lookahead])) * s.per(); + const path_ids = (last / s.rules.page - s.tree.have(prompt[0 .. last + s.rules.lookahead])) * s.per(); + const new = path_ids + parts_ids; if (s.tree.heldPages() + new <= s.limits.pages and (s.page_only or s.tree.snaps + marks_.len <= s.limits.snaps)) return; if (!s.evictOne(prompt, from)) return; } @@ -308,9 +312,10 @@ pub const Store = struct { var out: std.ArrayList(u32) = .empty; errdefer out.deinit(a); for (marks_) |m| { - if (m / s.rules.page * s.per() <= s.limits.pages and (s.page_only or s.limits.snaps > 0)) try out.append(a, m) else { + const ids = m / s.rules.page * s.per() + s.decl.planIds(&s.tree, m); // the path and the parts' known pages + if (ids <= s.limits.pages and (s.page_only or s.limits.snaps > 0)) try out.append(a, m) else { s.tree.counts.refused += 1; - note("kept nothing at {d} tokens: {d} pages pass the limit of {d}", .{ m, m / s.rules.page * s.per(), s.limits.pages }); + note("kept nothing at {d} tokens: {d} pages pass the limit of {d}", .{ m, ids, s.limits.pages }); } } return out.toOwnedSlice(a); @@ -364,7 +369,8 @@ pub const Store = struct { while (true) { var last: ?*Node = null; const got = t.heldPrefix(tokens, np, path, &last); - if (s.page_only and got == np) return .{ .held = true, .shared = mine.len }; // the path is the tree's already + // a page-only store: the path is the tree's already + if (s.page_only and got == np) return .{ .held = true, .shared = mine.len }; if (got == np and last != null and last.?.end == at and last.?.entry != null) { const e = &last.?.entry.?; e.used = t.clock; // the same state again: no copy @@ -455,7 +461,8 @@ pub const Store = struct { /// Evicts until `want` pool pages are free and not promised; false when the tree cannot give that many. pub fn reclaim(s: *Store, want: usize) bool { const pages = s.tree.pages; - if (s.decl.specs.len > 0) while (pages.vtable.available(pages.ptr) < want) { // parts first: their path may be in use + // parts first: their path may be in use + if (s.decl.specs.len > 0) while (pages.vtable.available(pages.ptr) < want) { s.tree.evict(s.partVictim() orelse break); }; while (pages.vtable.available(pages.ptr) < want) if (!s.tree.evictLeaf()) return false; diff --git a/zig/src/core/prompt_radix_groups_test.zig b/zig/src/core/prompt_radix_groups_test.zig index 62a7d29848e..e6c8a8dd77c 100644 --- a/zig/src/core/prompt_radix_groups_test.zig +++ b/zig/src/core/prompt_radix_groups_test.zig @@ -102,7 +102,8 @@ test "two prefix groups: a longer prompt shares the first page's ids and adds it try testing.expect(s.keep(.of(&a), 8, null, &mine_a, &path).held); const kept = s.keep(.of(&b), 8, null, &mine_b, &path); try testing.expect(kept.held); - try testing.expectEqualSlices(u32, &.{ 0, 1, 6, 7 }, &path); // the shared first page is a's ids; b's second is its own + // the shared first page is a's ids; b's second is its own + try testing.expectEqualSlices(u32, &.{ 0, 1, 6, 7 }, &path); try testing.expectEqual(@as(usize, 6), s.tree.heldPages()); } @@ -202,7 +203,8 @@ test "a store with no checkpoint making room for a pass's marks never frees the const longer = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21 }; var adopt: std.ArrayList(u32) = .empty; defer adopt.deinit(testing.allocator); - const plan = try s.begin(testing.allocator, .of(&longer), 21, &.{}, null, &adopt); // its mark at 20 wants 2 more pages + // its mark at 20 wants 2 more pages + const plan = try s.begin(testing.allocator, .of(&longer), 21, &.{}, null, &adopt); defer testing.allocator.free(plan.marks); try testing.expectEqual(@as(u32, 12), plan.from); try testing.expectEqualSlices(u32, &mine, adopt.items); diff --git a/zig/src/core/prompt_radix_make.zig b/zig/src/core/prompt_radix_make.zig index 4a6803dba4f..173384f83cf 100644 --- a/zig/src/core/prompt_radix_make.zig +++ b/zig/src/core/prompt_radix_make.zig @@ -44,6 +44,16 @@ pub const Decl = struct { return n; } + /// Ids a keep at `at` adds that a plan knows before the keep: copied parts' pages, one per retained partial page. + pub fn planIds(d: Decl, t: *const Tree, at: u32) usize { + var n: usize = 0; + for (d.specs, 0..) |spec, g| switch (spec.make) { + .copied => n += copiedPages(d, t, @intCast(g), at), + .retained => n += @intFromBool(at % spec.rows != 0), + }; + return n; + } + fn copiedPages(d: Decl, t: *const Tree, g: u32, at: u32) usize { const h = d.hooks.?; const bytes = h.vtable.bytes(h.ptr, g, at); @@ -136,7 +146,8 @@ pub const Decl = struct { nr = 0; for (d.specs, ps) |spec, part| if (spec.make == .retained) { const whole = part.ids.len - @intFromBool(at % spec.rows != 0); - out[nr] = .{ .base = part.base, .ids = try a.dupe(u32, part.ids[0..whole]) }; // the plan's own: eviction frees part.ids + // the plan's own copy: eviction frees part.ids + out[nr] = .{ .base = part.base, .ids = try a.dupe(u32, part.ids[0..whole]) }; nr += 1; owned = nr; }; @@ -145,7 +156,8 @@ pub const Decl = struct { for (d.specs, ps, 0..) |spec, part, g| if (spec.make == .retained) { const rows_left = at % spec.rows; const whole = part.ids.len - @intFromBool(rows_left != 0); - for (part.ids) |id| t.pages.vtable.retain(t.pages.ptr, id); // the stream's hold on each, until it adopts or finishes + // the stream's hold on each, until it adopts or finishes + for (part.ids) |id| t.pages.vtable.retain(t.pages.ptr, id); nr += 1; if (rows_left == 0) continue; copies[nc] = .{ .group = @intCast(g), .index = at / spec.rows, .src = part.ids[whole], .rows = rows_left }; From 05f53427ca6e8baacd0b36184c19a2a6c47156a8 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Wed, 11 Nov 2026 10:30:54 +0000 Subject: [PATCH 44/78] test(radix): a legacy save that reclaims never prunes its leaf The CUDA ring families' snapshot saves take pool pages, so a save can make the pool reclaim from the store in the middle of a keep. On a keep onto a leaf an earlier failed save left with no entry, that reclaim could prune the very leaf the keep then writes its entry to. The keep's node pin prevents it; this test holds a save that needs one page more than free over exactly that leaf, and it crashes with the pin removed. --- zig/src/core/kv_pool_store_test.zig | 74 +++++++++++++++++++++++++++++ 1 file changed, 74 insertions(+) diff --git a/zig/src/core/kv_pool_store_test.zig b/zig/src/core/kv_pool_store_test.zig index fa3f5b18123..7f3699267dc 100644 --- a/zig/src/core/kv_pool_store_test.zig +++ b/zig/src/core/kv_pool_store_test.zig @@ -212,3 +212,77 @@ test "page accounting: random admits, keeps, resumes, finishes, cancels and fail } try expect(refused > 0 and failed > 0 and adopted > 0); // the runs met full pools, failed uploads and resumes } + +/// A legacy family whose save holds pool pages for its copy (as the CUDA ring families' blobs do), one more than free. +const HoldingFamily = struct { + pool: *Pool, + fail: bool = false, + live: usize = 0, + + const State = struct { ids: [8]u32, n: usize }; + + fn of(ptr: *anyopaque) *HoldingFamily { + return @ptrCast(@alignCast(ptr)); + } + fn snapshots(f: *HoldingFamily) radix.Snapshots { + return .{ .ptr = f, .vtable = &.{ .bytes = bytes, .save = save, .restore = restore, .drop = drop } }; + } + fn bytes(_: *anyopaque, _: u32) u64 { + return 1; + } + fn save(ptr: *anyopaque, _: ?*anyopaque, _: u32) anyerror!radix.Saved { + const f = of(ptr); + if (f.fail) return error.CopyFailed; + const st = try std.testing.allocator.create(State); + st.n = f.pool.available() + 1; // the pool must reclaim from the store to give them + if (!f.pool.hold(st.ids[0..st.n])) { + std.testing.allocator.destroy(st); + return error.NoCachePages; + } + f.live += 1; + return st; + } + fn restore(_: *anyopaque, _: ?*anyopaque, _: radix.Saved) anyerror!void {} + fn drop(ptr: *anyopaque, saved: radix.Saved) void { + const f = of(ptr); + const st: *State = @ptrCast(@alignCast(saved)); + for (st.ids[0..st.n]) |id| f.pool.release(id); + std.testing.allocator.destroy(st); + f.live -= 1; + } +}; + +fn reclaimStore(ctx: *anyopaque, want: usize) bool { + const s: *radix.Store = @ptrCast(@alignCast(ctx)); + return s.reclaim(want); +} + +test "a legacy save that reclaims pool pages while keeping onto a leaf with no entry never prunes that leaf" { + var pool = try Pool.init(std.testing.allocator, 6, 1); + defer pool.deinit(); + var f: HoldingFamily = .{ .pool = &pool }; + var store = try radix.Store.init(std.testing.allocator, f.snapshots(), radix.pagesOf(Pool, &pool, 1), .{ .page = page, .min_gap = 1, .min_prompt = 0 }); + store.limit(.{ .pages = 6, .snaps = 4 }); + try pool.joinReclaim(.{ .ctx = &store, .call = reclaimStore }); + const prompt = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8, 9 }; + var path: [2]u32 = undefined; + var first: Lease = .{}; + try pool.promise(&first, 2); + try pool.cover(&first, 2, null); + f.fail = true; // a failed save leaves the path in the tree, a leaf with no entry + try expect(!store.keep(.of(&prompt), 8, null, first.ids.items, &path).held); + pool.finish(&first); // the leaf's pages are the store's alone: the only ones reclaim could free + try expectEqual(@as(usize, 2), store.tree.heldPages()); + var second: Lease = .{}; + try pool.promise(&second, 2); + try pool.cover(&second, 2, null); + f.fail = false; + const kept = store.keep(.of(&prompt), 8, null, second.ids.items, &path); // onto that leaf; its save reclaims + try expect(!kept.held); // the pinned leaf could not be taken, so the save found no page + try expectEqual(@as(usize, 2), store.tree.heldPages()); // the leaf survived the keep that was writing to it + try expectEqual(@as(usize, 0), f.live); + pool.finish(&second); + pool.leaveReclaim(&store); + store.deinit(); + try expect(pool.idle() and pool.available() == pool.total()); +} From 782d58c6a60c505da059bcc65c772f1002a39808 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 13:24:45 +0000 Subject: [PATCH 45/78] changelog: #2875 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 24ad3b85410..14c6fc5b247 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Prompt cache on the core pool (the `TF_KV_SHARE` families, Qwen 3.5/3.6 on CUDA, and Qwen 3.5 and Nemotron on HIP): the radix store holds each pool page once, however many entries name it, and pins the node a keep or restore is using. A snapshot save that takes pool pages (Gemma 4, Kolibri-1, Spark2.5, Dory) can no longer make reclaim prune the leaf it is writing to. Replies and cached counts are unchanged. The store also learns checkpoint parts and prefix groups for later steps to use (#2875, part of #2444). Thanks to @laudney. + - GLM-5.3-Flash on CUDA at four ranks or more: when the default split prefill's start-up exchange with every rank fails, as on a switchless ring of four, the failing rank now prints a refusal naming `TF_GLM_PREFILL_REDUCE=gather`, and rank 0 names it too when another rank fails first. Before, every rank stopped on a bare `ncclGroupEnd` error. Start-ups that succeed are unchanged (#2868). Thanks to @Bizuayeu. - HIP on gfx1151: the routed prompt tile runs an expert of at most 8 rows in K-parallel passes inside the same launch, with the matrix path's bytes; with Qwen3.6-35B-A3B on a Radeon 8060S, served cold prompts of 64 to 256 tokens are 11% to 22% faster (1.5% at 1,024), and logits and served tokens equal main's. Other cards' tables leave it off (#2878, fixes #2795). Thanks to @AbrahamCete. From d378271fce3b79a782af89db5b58cc4b0daa1d51 Mon Sep 17 00:00:00 2001 From: ThinkOffApp <233579950+ThinkOffApp@users.noreply.github.com> Date: Mon, 9 Nov 2026 23:32:54 +0200 Subject: [PATCH 46/78] qwen27: slot 0 keeps its rewind point through the warm pass, and a background pass's chunks narrow with depth A 20 GiB prompt cache holds two 128k states. The state kept at a reply's end pushed out the one at the history end, so a client whose replayed reply differs from the generated one resumed a turn further back. Slot 0's in-slot rewind buffer now exists whenever there is room, also beside a cache that keeps two full-context states, where needsRewind alone makes none, and its point is carried through passes that only extend slot 0, so the next turn rewinds in place. Beside such a wide cache, which keeps the reply's end, the warm pass over slot 0 runs when it carries that point; beside a smaller cache it still yields, as on main. Background chunks scale down with depth so a waiting request is held for less. --- zig/src/native/qwen27_prompt_rows.zig | 21 +++++--- zig/src/native/qwen27_resume.zig | 72 ++++++++++++++++++++++++++- zig/src/native/qwen27_served.zig | 2 +- zig/src/native/qwen27_slots.zig | 32 +++++++----- zig/src/qwen27_rows_test.zig | 47 +++++++++++++---- 5 files changed, 141 insertions(+), 33 deletions(-) diff --git a/zig/src/native/qwen27_prompt_rows.zig b/zig/src/native/qwen27_prompt_rows.zig index 9f885bb3a33..893c77e99de 100644 --- a/zig/src/native/qwen27_prompt_rows.zig +++ b/zig/src/native/qwen27_prompt_rows.zig @@ -1,9 +1,14 @@ -//! Each prompt chunk chooses its width from the streams that can decode beside it. -pub fn choose(b: anytype, current: anytype) usize { - for (b.owners[0..b.runner.slots]) |owner| if (owner) |stream| { - if (stream != current and !stream.finished and !stream.paused and !b.fills.contains(stream)) { - return @min(b.chunk, 512); - } - }; - return b.chunk; +//! Each prompt chunk chooses its width from the streams that can decode beside it, and a background pass's depth. +pub fn choose(b: anytype, current: anytype, at: usize) usize { + const width = for (b.owners[0..b.runner.slots]) |owner| { + const stream = owner orelse continue; + if (stream != current and !stream.finished and !stream.paused and !b.fills.contains(stream)) break @min(b.chunk, 512); + } else b.chunk; + return if (current.yield_check != null) @min(width, deep(b.chunk, at)) else width; +} + +/// A background chunk's width at offset `at`: narrower with depth, each about as long (256 at least). +pub fn deep(chunk: usize, at: usize) usize { + const scaled = @min(chunk, chunk * 16384 / @max(at, 16384)); + return @max(256, scaled / 128 * 128); } diff --git a/zig/src/native/qwen27_resume.zig b/zig/src/native/qwen27_resume.zig index 29e1e543b53..7a45c8e5f40 100644 --- a/zig/src/native/qwen27_resume.zig +++ b/zig/src/native/qwen27_resume.zig @@ -27,11 +27,13 @@ pub const Standing = struct { tokens: []u32 = &.{}, spans: []modes.Span = &.{}, rewind_at: u32 = 0, // where the rewind buffer holds slot 0's DeltaNet states and taps (0: none) + wide: bool = false, // the prompt cache keeps two full-context states: it holds a reply's end beside slot 0 + /// Frees the record; `wide`, the cache's size, stays. pub fn deinit(st: *Standing, gpa: std.mem.Allocator) void { gpa.free(st.tokens); gpa.free(st.spans); - st.* = .{}; + st.* = .{ .wide = st.wide }; } /// What `l` left in slot 0 of `context` (its rows' spans `spans`), when a later prompt may resume it; else nothing. @@ -41,7 +43,20 @@ pub const Standing = struct { if (!l.kept() or l.at > context.len) return; const canonical = (if (l.at <= l.prompt_len) modes.prefix(gpa, spans, l.at) else modes.reply(gpa, spans, l.prompt_len, l.at)) catch return; const tokens = gpa.dupe(u32, context[0..l.at]) catch return gpa.free(canonical); - st.* = .{ .tokens = tokens, .spans = canonical, .rewind_at = if (rewind < l.at) rewind else 0 }; + st.* = .{ .tokens = tokens, .spans = canonical, .rewind_at = if (rewind < l.at) rewind else 0, .wide = st.wide }; + } + + /// The rewind point a pass resuming slot 0 at `n` keeps while it only extends the slot (0: none). + pub fn carry(st: *const Standing, n: u32) u32 { + return if (st.rewind_at > 0 and st.rewind_at <= n) st.rewind_at else 0; + } + + /// The rewind point a pass that `r` resumes in slot 0 carries through it (0: none, or it does not resume there). + pub fn carried(st: *const Standing, r: Resume) u32 { + return switch (r) { + .none => 0, + .here, .back => |n| st.carry(n), + }; } /// How `ids` (its rows' spans `spans`) resumes, past the cache's own restore point `kept_at`, by the cache's rule. @@ -184,6 +199,59 @@ test "qwen27 resume: slot 0 keeps what a fresh pass would stand at, and the poin try t.expectEqual(@as(u32, 0), st.rewind_at); } +test "qwen27 resume: a warm pass that only extends slot 0 keeps the reply's rewind point for a diverging next turn" { + var st: Standing = .{ .rewind_at = 4 }; // the turn's pass saved slot 0's states at its history's end + defer st.deinit(t.allocator); + const context = [_]u32{ 1, 2, 3, 4, 20, 21, 22, 23 }; // history 1-4, the generation prompt 20, a reply 21-23 + st.leave(t.allocator, .{ .at = 7, .prompt_len = 5, .context_len = 8, .fill_at = 5, .ended = true }, &context, &.{}); + const warm = [_]u32{ 1, 2, 3, 4, 20, 21, 22, 24, 25, 26 }; // a warm pass past the slot's 7: its history is all + const n = switch (st.plan(&warm, &.{.{ 5, 7 }}, 0)) { + .here => |x| x, + else => return error.TestUnexpectedResult, + }; + try t.expectEqual(@as(u32, 7), n); + const kept = st.carry(n); // as takeStanding: the record is used up, the point stays + st.deinit(t.allocator); + st.rewind_at = kept; + st.leave(t.allocator, .{ .at = 9, .prompt_len = 10, .context_len = 10, .fill_at = 9, .ended = false }, &warm, &.{.{ 5, 7 }}); + try t.expectEqual(@as(u32, 4), st.rewind_at); // a yield at a landed chunk keeps it too + st.leave(t.allocator, .{ .at = 10, .prompt_len = 10, .context_len = 10, .fill_at = 10, .ended = true }, &warm, &.{.{ 5, 7 }}); + try t.expectEqualSlices(u32, &warm, st.tokens); + try t.expectEqual(@as(u32, 4), st.rewind_at); + const next = [_]u32{ 1, 2, 3, 4, 30, 23, 31 }; // the reply rendered without its reasoning + try t.expectEqual(Resume{ .back = 4 }, st.plan(&next, &.{.{ 5, 6 }}, 0)); + try t.expectEqual(Resume{ .here = 10 }, st.plan(&(warm ++ [_]u32{32}), &.{.{ 5, 7 }}, 0)); +} + +test "qwen27 resume: a rewind point past where the pass resumes, or none, is not carried" { + const st: Standing = .{ .rewind_at = 4 }; + try t.expectEqual(@as(u32, 4), st.carry(4)); // a `.back` pass resumes at the point itself + try t.expectEqual(@as(u32, 4), st.carry(9)); + try t.expectEqual(@as(u32, 0), st.carry(3)); + try t.expectEqual(@as(u32, 0), st.carry(0)); + try t.expectEqual(@as(u32, 0), (Standing{}).carry(9)); +} + +test "qwen27 resume: a warm pass over slot 0's conversation runs where it carries the rewind point, else steps aside" { + var st: Standing = .{ .rewind_at = 4 }; + defer st.deinit(t.allocator); + const context = [_]u32{ 1, 2, 3, 4, 20, 21, 22, 23 }; // history 1-4, the generation prompt 20, a reply 21-23 + st.leave(t.allocator, .{ .at = 7, .prompt_len = 5, .context_len = 8, .fill_at = 5, .ended = true }, &context, &.{}); + const warm = [_]u32{ 1, 2, 3, 4, 30, 23, 31 }; // the reply as the template renders it: it goes back to the point + try t.expectEqual(Resume{ .back = 4 }, st.plan(&warm, &.{.{ 5, 6 }}, 0)); + try t.expectEqual(@as(u32, 4), st.carried(st.plan(&warm, &.{.{ 5, 6 }}, 0))); // it runs + try t.expectEqual(@as(u32, 0), st.carried(st.plan(&warm, &.{.{ 5, 6 }}, 4))); // the cache's state is further: aside + const kept = st.carried(st.plan(&warm, &.{.{ 5, 6 }}, 0)); // as takeStanding, then the pass to its end + st.deinit(t.allocator); + st.rewind_at = kept; + st.leave(t.allocator, .{ .at = 7, .prompt_len = 7, .context_len = 7, .fill_at = 7, .ended = true }, &warm, &.{.{ 5, 6 }}); + try t.expectEqual(Resume{ .here = 7 }, st.plan(&(warm ++ [_]u32{40}), &.{.{ 5, 6 }}, 0)); // a client replaying it + try t.expectEqual(Resume{ .back = 4 }, st.plan(&.{ 1, 2, 3, 4, 20, 21, 22, 23, 40 }, &.{}, 0)); // one that does not + st.rewind_at = 0; // a slot with no rewind point would keep none through the pass: it steps aside + try t.expectEqual(@as(u32, 0), st.carried(st.plan(&(warm ++ [_]u32{40}), &.{.{ 5, 6 }}, 0))); + try t.expectEqual(@as(u32, 0), st.carried(.none)); +} + test { _ = @import("qwen27_resume_cache_test.zig"); } diff --git a/zig/src/native/qwen27_served.zig b/zig/src/native/qwen27_served.zig index 2e24131464e..b307e195c09 100644 --- a/zig/src/native/qwen27_served.zig +++ b/zig/src/native/qwen27_served.zig @@ -421,7 +421,7 @@ pub fn attach(parent: Allocator, io: std.Io, model: *q.model.Model, drafter: ?*q if (model.ane) |share| share.footprint_growth = platform.memory_policy.growth(before, cache_fit.footprint()); } ane.end(); - // Slot 0's history rewind when two full-context copies do not fit, wired before cache sizing. + // Slot 0's history rewind, wired before cache sizing: even a cache with two full-context copies may evict it. const rewind_room = laterRoom(h); if (slot_state.rewindInit(h.back, rewind_room, if (o.cache_gib) |g| gibBytes(g) else cache_fit.defaultBudget(rewind_room))) |buf| if (h.resident) |*r| r.addNow(&.{buf}); try enableCache(h, o.cache_gib, o.cache_over_cap, a, why); diff --git a/zig/src/native/qwen27_slots.zig b/zig/src/native/qwen27_slots.zig index 2df612a9944..c893a9c0788 100644 --- a/zig/src/native/qwen27_slots.zig +++ b/zig/src/native/qwen27_slots.zig @@ -27,13 +27,14 @@ pub fn opening(b: *Metal, s: *lanes.Stream, slot: u32) !Fill { const ids = s.prompt(); if (ids.len == 0 or ids.len + s.max_new + b.window > b.runner.capacity) return error.PromptTooLong; if (s.yields()) return error.Yielded; // a background pass steps aside before it touches the slot - // as does one over the conversation slot 0 holds when the cache cannot keep its end: it would leave nothing there - if (slot == 0 and s.yield_check != null and s.max_new == 0 and !keepsEnd(s) and resumable(b, s, 0) != .none) return error.Yielded; + const kept_at: u32 = if (s.reuse.saved) |saved| @as(*const q.snapshot.State, @ptrCast(@alignCast(saved))).at else 0; + // as does one over slot 0's conversation when the cache cannot keep its end, unless a wide one and the point allow + if (slot == 0 and s.yield_check != null and s.max_new == 0 and !keepsEnd(s) and resumable(b, s, 0) != .none and + (!b.standing.wide or b.standing.carried(resumable(b, s, kept_at)) == 0)) return error.Yielded; if (b.vision != null and s.images.len > 0 and b.image_fill != null and b.image_fill != s) return error.LaneMemory; try recover(b, s); b.runner.rope[slot] = try images.layout(b.vision, s); const d = b.drafter(slot); - const kept_at: u32 = if (s.reuse.saved) |saved| @as(*const q.snapshot.State, @ptrCast(@alignCast(saved))).at else 0; const resumed = if (slot == 0) try takeStanding(b, s, kept_at) else 0; // slot 0's own state: nothing copied if (resumed == 0) { if (d) |x| try x.reset() else try b.runner.reset(slot); @@ -87,7 +88,7 @@ pub fn chunk(b: *Metal, s: *lanes.Stream, f: *Fill, _: u32) !bool { return error.Yielded; } const at = f.at; - var end = resume_rules.chunkEnd(at, @import("qwen27_prompt_rows.zig").choose(b, s), ids.len, s.chunks, s.decode_spans); + var end = resume_rules.chunkEnd(at, @import("qwen27_prompt_rows.zig").choose(b, s, at), ids.len, s.chunks, s.decode_spans); if (f.checkpoint > at and f.checkpoint < end) end = f.checkpoint; end = resume_rules.rewindEnd(at, end, f.rewind); var decoded = false; @@ -233,21 +234,27 @@ fn resumable(b: *const Metal, s: *const lanes.Stream, kept_at: u32) resume_rules return b.standing.plan(s.prompt(), s.decode_spans, kept_at); } -/// Slot 0's position for `s` when it resumes there (0: not); the record is used up, as the pass changes the slot. +/// Slot 0's position for `s` when it resumes there (0: not); the record is used up, its rewind point carried on. fn takeStanding(b: *Metal, s: *const lanes.Stream, kept_at: u32) !u32 { - defer b.standing.deinit(b.gpa); - switch (resumable(b, s, kept_at)) { + const plan = resumable(b, s, kept_at); + const kept = b.standing.carried(plan); + b.standing.deinit(b.gpa); + // rows [0, rewind_at) and the buffer hold only while slot 0 just extends past it: a restore, reset or move drops it + switch (plan) { .none => return 0, .here => |n| { - const d = b.draft orelse return n; - try d.model.reset(); // a new session over the kept taps, as after a restore - d.model.committed_end = n; + if (b.draft) |d| { + try d.model.reset(); // a new session over the kept taps, as after a restore + d.model.committed_end = n; + } + b.standing.rewind_at = kept; return n; }, .back => |h| { const buf = b.rewind orelse return 0; if (b.draft) |d| try d.model.reset(); q.snapshot.rewindSlot(b.gpa, b.runner, 0, b.ringOf(0), h, buf) catch return 0; // the pass then starts at 0 + b.standing.rewind_at = kept; return h; }, } @@ -274,11 +281,12 @@ fn saveRewind(b: *Metal, at: u32) void { b.standing.rewind_at = at; } -/// The history rewind where two full-context states do not fit in `cache`, provided `room` holds its bytes. +/// The history rewind, where `room` holds it; beside a `cache` that keeps two full-context states it is also `wide`. pub fn rewindInit(b: *Metal, room: u64, cache: u64) ?mtl.Buffer { if (!slotResume(b) or b.runner.ready == 0) return null; const need = q.snapshot.rewindBytes(b.runner, b.ringOf(0)); - if (!resume_rules.needsRewind(cache, q.snapshot.Spares.size(q.snapshot.slotBytes(b.runner, b.ringOf(0), b.runner.capacity)))) return null; + // such a cache keeps the reply's end but may still evict the history's end: slot 0 keeps that one in place + b.standing.wide = !resume_rules.needsRewind(cache, q.snapshot.Spares.size(q.snapshot.slotBytes(b.runner, b.ringOf(0), b.runner.capacity))); if (need > room) { std.log.info("qwen27: no rewind point for slot 0 ({d} MiB, {d} MiB left under the load limit): a next turn that drops the last reply's reasoning reads its prompt again", .{ need >> 20, room >> 20 }); return null; diff --git a/zig/src/qwen27_rows_test.zig b/zig/src/qwen27_rows_test.zig index 43c94d1f9b1..4d672e3568d 100644 --- a/zig/src/qwen27_rows_test.zig +++ b/zig/src/qwen27_rows_test.zig @@ -43,32 +43,59 @@ test "qwen27 rows: a decoder arriving between chunks caps the next one and leavi var b: Backend = .{ .owners = .{ &prompt, null, null } }; defer b.fills.deinit(a); try b.fills.put(a, &prompt, true); - try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt)); + try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt, 0)); b.owners[1] = &other; try b.fills.put(a, &other, true); - try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt)); + try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt, 0)); try std.testing.expect(b.fills.remove(&other)); - try std.testing.expectEqual(@as(usize, 512), policy.choose(&b, &prompt)); + try std.testing.expectEqual(@as(usize, 512), policy.choose(&b, &prompt, 0)); other.paused = true; - try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt)); + try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt, 0)); other.paused = false; - try std.testing.expectEqual(@as(usize, 512), policy.choose(&b, &prompt)); + try std.testing.expectEqual(@as(usize, 512), policy.choose(&b, &prompt, 0)); other.finished = true; - try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt)); + try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt, 0)); other.finished = false; for ([_]usize{ 16, 128, 256, 512, 1024, 2048 }) |capacity| { b.chunk = capacity; - try std.testing.expectEqual(@min(capacity, @as(usize, 512)), policy.choose(&b, &prompt)); + try std.testing.expectEqual(@min(capacity, @as(usize, 512)), policy.choose(&b, &prompt, 0)); } b.chunk = 2048; b.owners[1] = null; b.owners[2] = &other; b.runner.slots = 2; - try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt)); + try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt, 0)); b.runner.slots = 3; - try std.testing.expectEqual(@as(usize, 512), policy.choose(&b, &prompt)); + try std.testing.expectEqual(@as(usize, 512), policy.choose(&b, &prompt, 0)); b.owners[2] = null; - try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt)); + try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt, 0)); +} + +fn never(_: *anyopaque) bool { + return false; +} + +test "qwen27 rows: a background pass narrows its chunks with depth and a foreground one does not" { + var prompt = try Stream.init(a, .{ .id = "prompt", .prompt = &.{ 1, 2 }, .max_new = 4, .drafts = false }); + defer prompt.deinit(a); + var b: Backend = .{ .owners = .{ &prompt, null, null } }; + defer b.fills.deinit(a); + try b.fills.put(a, &prompt, true); + try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt, 120000)); // a foreground pass + prompt.yield_check = .{ .ptr = &b, .check = never }; + for ([_]usize{ 0, 16384 }) |at| try std.testing.expectEqual(@as(usize, 2048), policy.choose(&b, &prompt, at)); + try std.testing.expectEqual(@as(usize, 1664), policy.choose(&b, &prompt, 20000)); // 1677, a whole 128 rows + try std.testing.expectEqual(@as(usize, 1024), policy.choose(&b, &prompt, 32768)); + try std.testing.expectEqual(@as(usize, 256), policy.choose(&b, &prompt, 131072)); + try std.testing.expectEqual(@as(usize, 256), policy.choose(&b, &prompt, 262144)); // never below 256 + b.chunk = 128; // nor wider than the chunk + try std.testing.expectEqual(@as(usize, 128), policy.choose(&b, &prompt, 131072)); + b.chunk = 2048; + var other = try Stream.init(a, .{ .id = "other", .prompt = &.{ 1, 2 }, .max_new = 4, .drafts = false }); + defer other.deinit(a); + b.owners[1] = &other; // a decoder beside it caps the chunk first + try std.testing.expectEqual(@as(usize, 512), policy.choose(&b, &prompt, 20000)); + try std.testing.expectEqual(@as(usize, 256), policy.choose(&b, &prompt, 131072)); } const floor = @import("native/qwen27_prompt_floor.zig"); From 9ff4d7504e311f9a1796cc2bee930ad2333e3acf Mon Sep 17 00:00:00 2001 From: ThinkOffApp <233579950+ThinkOffApp@users.noreply.github.com> Date: Tue, 10 Nov 2026 13:04:54 +0200 Subject: [PATCH 47/78] qwen27: the warm pass's yield, the rewind carry and the rewind buffer's choice are tested rules Standing.warmYields holds the warm pass's yield in opening, Standing.take holds takeStanding's use of the record and its carried rewind point, and makesRewind holds rewindInit's choice: the buffer wherever room holds it, tight where needsRewind says the cache cannot keep two full-context states, else wide. Each has a test that fails when its rule is reverted. The changelog takes the line for #2723 under a new Unreleased heading. --- CHANGELOG.md | 2 + zig/src/native/qwen27_resume.zig | 71 +++++++++++++++++++++++++++++--- zig/src/native/qwen27_slots.zig | 26 +++++++----- 3 files changed, 82 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 14c6fc5b247..c8165118cdd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -151,6 +151,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - HIP on gfx12 under `--policy matrix=on`: Qwen3 dense checkpoints load each layer's gate and up weights interleaved as one weight, and one WMMA product writes silu(gate) x up in its epilogue instead of two products and `tf_silu_mul`, with the same output bytes: Qwen3-Reranker-8B ranks 16 documents in 0.657 to 0.661 s instead of 0.687 to 0.697 s and Qwen3-Embedding-8B embeds 100 short inputs in 0.901 to 0.911 s instead of 0.940 to 0.948 s on a Radeon AI PRO R9700. The kernel contract checks both fused blocks byte for byte against the unfused launches (part of #2245). Thanks to @laudney. - Nemotron on Macs: the masked prompt shader calls the shared `tf_paged_row` page helper again, so its prompt sources compile; #2322 had used the name #2283 retired the same day (#2548). +- Qwen3.8-27B on Metal: a follow-up turn whose client sends the last reply without its reasoning resumes slot 0 in place from its rewind point, now also beside a prompt cache large enough for two whole contexts (which can still evict the history's end), instead of a turn further back; there the warm pass keeps that point, and a background pass at depth runs in narrower chunks so a waiting request is held for less (#2723). Thanks to @ThinkOffApp. + ## 1.2.7 - GLM-5.3-Flash on CUDA: with `TF_GLM_LAYERS=N`, serve's start-up memory check counts the weights by tensor from the shard headers, only those the load reads (the first N layers, the MTP layer when drafts are on, the embedding, norm and head), where it counted every checkpoint file. A 20-layer thin load of the 163.6 GiB checkpoint now counts 66.6 GiB without drafts. Unset, the count is the files' as before (#2788, fixes #2745). Thanks to @wojo. diff --git a/zig/src/native/qwen27_resume.zig b/zig/src/native/qwen27_resume.zig index 7a45c8e5f40..f8ac6cefce3 100644 --- a/zig/src/native/qwen27_resume.zig +++ b/zig/src/native/qwen27_resume.zig @@ -59,6 +59,18 @@ pub const Standing = struct { }; } + /// Uses up the record for a pass that `r` resumes in slot 0, keeping the rewind point the pass carries (0: none). + pub fn take(st: *Standing, gpa: std.mem.Allocator, r: Resume) void { + const kept = st.carried(r); + st.deinit(gpa); + st.rewind_at = kept; + } + + /// Whether a warm pass over slot 0 (`own`) yields: not if the cache keeps its end, nor `wide` with `from` carrying. + pub fn warmYields(st: *const Standing, keeps_end: bool, own: Resume, from: Resume) bool { + return !keeps_end and own != .none and (!st.wide or st.carried(from) == 0); + } + /// How `ids` (its rows' spans `spans`) resumes, past the cache's own restore point `kept_at`, by the cache's rule. pub fn plan(st: *const Standing, ids: []const u32, spans: []const modes.Span, kept_at: u32) Resume { const n = st.tokens.len; @@ -101,6 +113,15 @@ pub fn rewindEnd(at: usize, end: usize, rewind: u32) usize { return if (rewind > at and rewind < end) rewind else end; } +/// Slot 0's `need`-byte history rewind: none, a `tight` one, or a `wide` one beside a cache that keeps two of them. +pub const Rewind = enum { none, tight, wide }; + +/// Slot 0 makes its rewind wherever `room` holds it: a `cache` keeping two `full`-byte states may still evict it. +pub fn makesRewind(need: u64, room: u64, cache: u64, full: u64) Rewind { + if (need > room) return .none; + return if (needsRewind(cache, full)) .tight else .wide; +} + /// A prompt's next physical chunk end, before cache marks or the slot's rewind point cut it further. pub fn chunkEnd(at: usize, chunk: usize, len: usize, starts: []const u32, spans: []const modes.Span) usize { var end = @min(len, at + chunk); @@ -210,9 +231,7 @@ test "qwen27 resume: a warm pass that only extends slot 0 keeps the reply's rewi else => return error.TestUnexpectedResult, }; try t.expectEqual(@as(u32, 7), n); - const kept = st.carry(n); // as takeStanding: the record is used up, the point stays - st.deinit(t.allocator); - st.rewind_at = kept; + st.take(t.allocator, .{ .here = n }); // as takeStanding: the record is used up, the point stays st.leave(t.allocator, .{ .at = 9, .prompt_len = 10, .context_len = 10, .fill_at = 9, .ended = false }, &warm, &.{.{ 5, 7 }}); try t.expectEqual(@as(u32, 4), st.rewind_at); // a yield at a landed chunk keeps it too st.leave(t.allocator, .{ .at = 10, .prompt_len = 10, .context_len = 10, .fill_at = 10, .ended = true }, &warm, &.{.{ 5, 7 }}); @@ -241,9 +260,7 @@ test "qwen27 resume: a warm pass over slot 0's conversation runs where it carrie try t.expectEqual(Resume{ .back = 4 }, st.plan(&warm, &.{.{ 5, 6 }}, 0)); try t.expectEqual(@as(u32, 4), st.carried(st.plan(&warm, &.{.{ 5, 6 }}, 0))); // it runs try t.expectEqual(@as(u32, 0), st.carried(st.plan(&warm, &.{.{ 5, 6 }}, 4))); // the cache's state is further: aside - const kept = st.carried(st.plan(&warm, &.{.{ 5, 6 }}, 0)); // as takeStanding, then the pass to its end - st.deinit(t.allocator); - st.rewind_at = kept; + st.take(t.allocator, st.plan(&warm, &.{.{ 5, 6 }}, 0)); // as takeStanding, then the pass to its end st.leave(t.allocator, .{ .at = 7, .prompt_len = 7, .context_len = 7, .fill_at = 7, .ended = true }, &warm, &.{.{ 5, 6 }}); try t.expectEqual(Resume{ .here = 7 }, st.plan(&(warm ++ [_]u32{40}), &.{.{ 5, 6 }}, 0)); // a client replaying it try t.expectEqual(Resume{ .back = 4 }, st.plan(&.{ 1, 2, 3, 4, 20, 21, 22, 23, 40 }, &.{}, 0)); // one that does not @@ -252,6 +269,48 @@ test "qwen27 resume: a warm pass over slot 0's conversation runs where it carrie try t.expectEqual(@as(u32, 0), st.carried(.none)); } +test "qwen27 resume: rule a, a warm pass carrying slot 0's rewind point runs beside a wide cache, else it yields" { + var st: Standing = .{ .rewind_at = 4 }; + defer st.deinit(t.allocator); + const context = [_]u32{ 1, 2, 3, 4, 20, 21, 22, 23 }; // history 1-4, the generation prompt 20, a reply 21-23 + st.leave(t.allocator, .{ .at = 7, .prompt_len = 5, .context_len = 8, .fill_at = 5, .ended = true }, &context, &.{}); + const warm = [_]u32{ 1, 2, 3, 4, 30, 23, 31 }; // back to the point at 4 + const own = st.plan(&warm, &.{.{ 5, 6 }}, 0); + try t.expect(st.warmYields(false, own, own)); // a tight cache would lose the reply's end it leaves: aside + st.wide = true; // the cache keeps two full-context states + try t.expect(!st.warmYields(false, own, own)); // it carries the point: it runs + try t.expect(st.warmYields(false, own, st.plan(&warm, &.{.{ 5, 6 }}, 4))); // the cache's state is further: aside + try t.expect(!st.warmYields(true, own, .none)); // the cache keeps the conversation's end + try t.expect(!st.warmYields(false, .none, .none)); // it would not resume in slot 0 + st.rewind_at = 0; + try t.expect(st.warmYields(false, own, own)); // no point to carry: it would leave nothing in slot 0 + st.deinit(t.allocator); + try t.expect(st.wide); // the record goes, the cache's size stays +} + +test "qwen27 resume: rule b, taking slot 0's record for a pass that extends it keeps the rewind point" { + var st: Standing = .{ .rewind_at = 4 }; + defer st.deinit(t.allocator); + const context = [_]u32{ 1, 2, 3, 4, 20, 21, 22, 23 }; + st.leave(t.allocator, .{ .at = 7, .prompt_len = 5, .context_len = 8, .fill_at = 5, .ended = true }, &context, &.{}); + st.take(t.allocator, st.plan(&[_]u32{ 1, 2, 3, 4, 20, 21, 22, 24 }, &.{.{ 5, 7 }}, 0)); // `.here` at 7 + try t.expectEqual(@as(usize, 0), st.tokens.len); // the record is used up + try t.expectEqual(@as(u32, 4), st.rewind_at); // the point stays + st.take(t.allocator, .none); // a pass that does not resume in slot 0 drops it + try t.expectEqual(@as(u32, 0), st.rewind_at); +} + +test "qwen27 resume: rule c, slot 0 makes its rewind buffer where room holds it, even beside a cache with two states" { + const full: u64 = 2496 << 20; // one full-context state at 33024 rows + try t.expectEqual(Rewind.wide, makesRewind(1 << 30, 2 << 30, 5 << 30, full)); // needsRewind says no: made anyway + try t.expectEqual(Rewind.wide, makesRewind(1 << 30, 2 << 30, 8 << 30, full)); + try t.expectEqual(Rewind.tight, makesRewind(1 << 30, 2 << 30, 4 << 30, full)); // needsRewind's own case + try t.expectEqual(Rewind.tight, makesRewind(1 << 30, 1 << 30, full, full)); + try t.expectEqual(Rewind.tight, makesRewind(1 << 30, 2 << 30, 0, full)); // no cache at all + try t.expectEqual(Rewind.none, makesRewind(2 << 30, 1 << 30, 0, full)); // no room + try t.expectEqual(Rewind.none, makesRewind(2 << 30, 1 << 30, 8 << 30, full)); +} + test { _ = @import("qwen27_resume_cache_test.zig"); } diff --git a/zig/src/native/qwen27_slots.zig b/zig/src/native/qwen27_slots.zig index c893a9c0788..e13b6755e1d 100644 --- a/zig/src/native/qwen27_slots.zig +++ b/zig/src/native/qwen27_slots.zig @@ -29,8 +29,8 @@ pub fn opening(b: *Metal, s: *lanes.Stream, slot: u32) !Fill { if (s.yields()) return error.Yielded; // a background pass steps aside before it touches the slot const kept_at: u32 = if (s.reuse.saved) |saved| @as(*const q.snapshot.State, @ptrCast(@alignCast(saved))).at else 0; // as does one over slot 0's conversation when the cache cannot keep its end, unless a wide one and the point allow - if (slot == 0 and s.yield_check != null and s.max_new == 0 and !keepsEnd(s) and resumable(b, s, 0) != .none and - (!b.standing.wide or b.standing.carried(resumable(b, s, kept_at)) == 0)) return error.Yielded; + if (slot == 0 and s.yield_check != null and s.max_new == 0 and + b.standing.warmYields(keepsEnd(s), resumable(b, s, 0), resumable(b, s, kept_at))) return error.Yielded; if (b.vision != null and s.images.len > 0 and b.image_fill != null and b.image_fill != s) return error.LaneMemory; try recover(b, s); b.runner.rope[slot] = try images.layout(b.vision, s); @@ -237,8 +237,8 @@ fn resumable(b: *const Metal, s: *const lanes.Stream, kept_at: u32) resume_rules /// Slot 0's position for `s` when it resumes there (0: not); the record is used up, its rewind point carried on. fn takeStanding(b: *Metal, s: *const lanes.Stream, kept_at: u32) !u32 { const plan = resumable(b, s, kept_at); - const kept = b.standing.carried(plan); - b.standing.deinit(b.gpa); + b.standing.take(b.gpa, plan); + errdefer b.standing.rewind_at = 0; // rows [0, rewind_at) and the buffer hold only while slot 0 just extends past it: a restore, reset or move drops it switch (plan) { .none => return 0, @@ -247,19 +247,23 @@ fn takeStanding(b: *Metal, s: *const lanes.Stream, kept_at: u32) !u32 { try d.model.reset(); // a new session over the kept taps, as after a restore d.model.committed_end = n; } - b.standing.rewind_at = kept; return n; }, .back => |h| { - const buf = b.rewind orelse return 0; + const buf = b.rewind orelse return dropRewind(b); if (b.draft) |d| try d.model.reset(); - q.snapshot.rewindSlot(b.gpa, b.runner, 0, b.ringOf(0), h, buf) catch return 0; // the pass then starts at 0 - b.standing.rewind_at = kept; + q.snapshot.rewindSlot(b.gpa, b.runner, 0, b.ringOf(0), h, buf) catch return dropRewind(b); // starts at 0 return h; }, } } +/// A pass that could not rewind slot 0 starts it afresh, keeping no rewind point (0: resumed nowhere). +fn dropRewind(b: *Metal) u32 { + b.standing.rewind_at = 0; + return 0; +} + /// A stream leaves slot 0: what it stands at stays for the next prompt that extends it, if a fresh pass stands there. fn leave(b: *Metal, s: *const lanes.Stream) void { const at = b.runner.offsets[0]; @@ -285,9 +289,9 @@ fn saveRewind(b: *Metal, at: u32) void { pub fn rewindInit(b: *Metal, room: u64, cache: u64) ?mtl.Buffer { if (!slotResume(b) or b.runner.ready == 0) return null; const need = q.snapshot.rewindBytes(b.runner, b.ringOf(0)); - // such a cache keeps the reply's end but may still evict the history's end: slot 0 keeps that one in place - b.standing.wide = !resume_rules.needsRewind(cache, q.snapshot.Spares.size(q.snapshot.slotBytes(b.runner, b.ringOf(0), b.runner.capacity))); - if (need > room) { + const kind = resume_rules.makesRewind(need, room, cache, q.snapshot.Spares.size(q.snapshot.slotBytes(b.runner, b.ringOf(0), b.runner.capacity))); + b.standing.wide = kind == .wide; // such a cache keeps a reply's end, but may evict the history's end + if (kind == .none) { std.log.info("qwen27: no rewind point for slot 0 ({d} MiB, {d} MiB left under the load limit): a next turn that drops the last reply's reasoning reads its prompt again", .{ need >> 20, room >> 20 }); return null; } From 9a79f7b2d3832047b88cdf95cd645e48ddabd761 Mon Sep 17 00:00:00 2001 From: ThinkOffApp <233579950+ThinkOffApp@users.noreply.github.com> Date: Tue, 10 Nov 2026 22:35:51 +0200 Subject: [PATCH 48/78] qwen27: a turn that resumed inside a reply kept as generated makes no new rewind point Slot 0 remembers the rows of the reply it generated. A follow-up whose prompt replays that reply as generated and resumes, by the slot or a cache state, inside it or at its end skips the extra rewind stop at its history end, so its fill is the same as main's. A client that drops the reply's reasoning, a cold turn and a resume that goes back still make a fresh rewind point. --- zig/src/native/qwen27_resume.zig | 104 ++++++++++++++++++++++++++++++- zig/src/native/qwen27_slots.zig | 11 ++-- 2 files changed, 108 insertions(+), 7 deletions(-) diff --git a/zig/src/native/qwen27_resume.zig b/zig/src/native/qwen27_resume.zig index f8ac6cefce3..bba7f7d6bb7 100644 --- a/zig/src/native/qwen27_resume.zig +++ b/zig/src/native/qwen27_resume.zig @@ -27,6 +27,7 @@ pub const Standing = struct { tokens: []u32 = &.{}, spans: []modes.Span = &.{}, rewind_at: u32 = 0, // where the rewind buffer holds slot 0's DeltaNet states and taps (0: none) + reply: modes.Span = .{ 0, 0 }, // rows of a reply these tokens hold exactly as it was generated (empty: none) wide: bool = false, // the prompt cache keeps two full-context states: it holds a reply's end beside slot 0 /// Frees the record; `wide`, the cache's size, stays. @@ -39,11 +40,13 @@ pub const Standing = struct { /// What `l` left in slot 0 of `context` (its rows' spans `spans`), when a later prompt may resume it; else nothing. pub fn leave(st: *Standing, gpa: std.mem.Allocator, l: Leaving, context: []const u32, spans: []const modes.Span) void { const rewind = st.rewind_at; // a point this stream's own pass saved + const reply = st.reply; // a reply the pass resumed past and only extended st.deinit(gpa); if (!l.kept() or l.at > context.len) return; const canonical = (if (l.at <= l.prompt_len) modes.prefix(gpa, spans, l.at) else modes.reply(gpa, spans, l.prompt_len, l.at)) catch return; const tokens = gpa.dupe(u32, context[0..l.at]) catch return gpa.free(canonical); - st.* = .{ .tokens = tokens, .spans = canonical, .rewind_at = if (rewind < l.at) rewind else 0, .wide = st.wide }; + const own: modes.Span = if (l.at > l.prompt_len) .{ l.prompt_len, l.at } else if (reply[1] <= l.at) reply else .{ 0, 0 }; + st.* = .{ .tokens = tokens, .spans = canonical, .rewind_at = if (rewind < l.at) rewind else 0, .reply = own, .wide = st.wide }; } /// The rewind point a pass resuming slot 0 at `n` keeps while it only extends the slot (0: none). @@ -59,11 +62,25 @@ pub const Standing = struct { }; } - /// Uses up the record for a pass that `r` resumes in slot 0, keeping the rewind point the pass carries (0: none). + /// Uses up the record for a pass that `r` resumes in slot 0, keeping the rewind point and reply it carries. pub fn take(st: *Standing, gpa: std.mem.Allocator, r: Resume) void { const kept = st.carried(r); + const reply: modes.Span = switch (r) { + .here => |n| if (st.reply[1] <= n) st.reply else .{ 0, 0 }, + .back, .none => .{ 0, 0 }, + }; st.deinit(gpa); st.rewind_at = kept; + st.reply = reply; + } + + /// Before `take`: where `ids` (spans `spans`) sends back the reply these tokens hold, as it was generated. + pub fn replayed(st: *const Standing, ids: []const u32, spans: []const modes.Span) Replay { + const from, const end = st.reply; + if (end <= from or end > st.tokens.len) return .{}; + const same: u32 = @intCast(std.mem.findDiff(u32, st.tokens[0..end], ids) orelse end); + if (same <= from or !modes.equal(st.spans, spans, same)) return .{}; + return .{ .from = from, .to = if (same == end) std.math.maxInt(u32) else same }; } /// Whether a warm pass over slot 0 (`own`) yields: not if the cache keeps its end, nor `wide` with `from` carrying. @@ -122,6 +139,17 @@ pub fn makesRewind(need: u64, room: u64, cache: u64, full: u64) Rewind { return if (needsRewind(cache, full)) .tight else .wide; } +/// Resume rows (from, to] past a reply's prompt whose reply rows before them a prompt sends back as generated. +pub const Replay = struct { + from: u32 = 0, + to: u32 = 0, + + /// A pass from `at` (slot or cache) resumed in a reply sent back as generated: it makes no point (unused there). + pub fn holds(p: Replay, at: usize) bool { + return at > p.from and at <= p.to; + } +}; + /// A prompt's next physical chunk end, before cache marks or the slot's rewind point cut it further. pub fn chunkEnd(at: usize, chunk: usize, len: usize, starts: []const u32, spans: []const modes.Span) usize { var end = @min(len, at + chunk); @@ -311,6 +339,78 @@ test "qwen27 resume: rule c, slot 0 makes its rewind buffer where room holds it, try t.expectEqual(Rewind.none, makesRewind(2 << 30, 1 << 30, 8 << 30, full)); } +test "qwen27 resume: a follow-up sending the reply back as generated, restored at the cache's reply end, makes no point" { + var st: Standing = .{ .rewind_at = 4 }; // the turn's pass saved slot 0's states at its history's end + defer st.deinit(t.allocator); + const context = [_]u32{ 1, 2, 3, 4, 20, 21, 22, 23 }; // history 1-4, the generation prompt 20, a reply 21-23 + st.leave(t.allocator, .{ .at = 7, .prompt_len = 5, .context_len = 8, .fill_at = 5, .ended = true }, &context, &.{}); + try t.expectEqual(modes.Span{ 5, 7 }, st.reply); // the stream left one short of its reply's end + const next = context ++ [_]u32{ 30, 31 }; // the reply as generated, then a new user turn + const spans = [_]modes.Span{.{ 5, 8 }}; + try t.expectEqual(@as(Resume, .none), st.plan(&next, &spans, 7)); // the cache keeps the reply's end: its own path + const replay = st.replayed(&next, &spans); // as opening, before takeStanding uses the record up + st.take(t.allocator, .none); + try t.expectEqual(@as(u32, 0), st.rewind_at); + try t.expect(replay.holds(7)); // restored at 7, inside the generated rows 5-7: no stop at the history's end + try t.expect(!replay.holds(4)); // a cache state at the history's mark is prompt rows: the point is made + try t.expect(!replay.holds(0)); + const turn = next ++ [_]u32{ 40, 41, 42 }; // that pass, with no point, and its reply 40-42 + st.leave(t.allocator, .{ .at = 12, .prompt_len = 10, .context_len = 13, .fill_at = 10, .ended = true }, &turn, &spans); + try t.expectEqual(modes.Span{ 10, 12 }, st.reply); + const again = [_]modes.Span{ .{ 5, 8 }, .{ 10, 13 } }; + try t.expect(st.replayed(&(turn ++ [_]u32{50}), &again).holds(12)); // the next follow-up skips it too + const diverging = next ++ [_]u32{ 60, 42, 61 }; // the last reply rendered without its reasoning + const cut = [_]modes.Span{ .{ 5, 8 }, .{ 10, 12 } }; + try t.expectEqual(@as(Resume, .none), st.plan(&diverging, &cut, 9)); // no point: the cache's state, as main + try t.expect(!st.replayed(&diverging, &cut).holds(9)); // and that pass makes one + try t.expectEqual(Replay{}, st.replayed(&diverging, &.{})); // reply rows sent as prompt rows are not as generated +} + +test "qwen27 resume: a follow-up resuming in slot 0 past a reply kept as generated makes no point, carrying the old one" { + var st: Standing = .{ .rewind_at = 4 }; + defer st.deinit(t.allocator); + const context = [_]u32{ 1, 2, 3, 4, 20, 21, 22, 23 }; + st.leave(t.allocator, .{ .at = 7, .prompt_len = 5, .context_len = 8, .fill_at = 5, .ended = true }, &context, &.{}); + const warm = context ++ [_]u32{ 30, 31 }; // a warm pass that ran: it extends the slot past the reply + st.take(t.allocator, st.plan(&warm, &.{.{ 5, 8 }}, 0)); // `.here` at 7 + st.leave(t.allocator, .{ .at = 9, .prompt_len = 10, .context_len = 10, .fill_at = 9, .ended = false }, &warm, &.{.{ 5, 8 }}); + try t.expectEqual(modes.Span{ 5, 7 }, st.reply); // a yield at a landed chunk keeps it + st.leave(t.allocator, .{ .at = 10, .prompt_len = 10, .context_len = 10, .fill_at = 10, .ended = true }, &warm, &.{.{ 5, 8 }}); + try t.expectEqual(modes.Span{ 5, 7 }, st.reply); // and so does the pass to its prompt's end + const follow = warm ++ [_]u32{32}; + try t.expectEqual(Resume{ .here = 10 }, st.plan(&follow, &.{.{ 5, 8 }}, 0)); + const replay = st.replayed(&follow, &.{.{ 5, 8 }}); + try t.expect(replay.holds(10)); // in the slot, or a cache state at the warm pass's end + try t.expectEqual(Resume{ .back = 4 }, st.plan(&.{ 1, 2, 3, 4, 50, 23, 51 }, &.{.{ 5, 6 }}, 0)); // reply dropped + st.take(t.allocator, .{ .here = 10 }); + try t.expectEqual(@as(u32, 4), st.rewind_at); // the older point stays for such a turn + st.leave(t.allocator, .{ .at = 3, .prompt_len = 5, .context_len = 5, .fill_at = 3, .ended = false }, &context, &.{}); + try t.expectEqual(modes.Span{ 0, 0 }, st.reply); // a reply past where the slot stands is not kept +} + +test "qwen27 resume: a client dropping the reasoning, a cold turn or a pass that resumes nowhere makes a fresh point" { + var st: Standing = .{ .rewind_at = 4 }; + defer st.deinit(t.allocator); + try t.expect(!st.replayed(&.{ 1, 2, 3 }, &.{}).holds(0)); // cold: nothing stands + const context = [_]u32{ 1, 2, 3, 4, 20, 21, 22, 23 }; + const reply: Leaving = .{ .at = 7, .prompt_len = 5, .context_len = 8, .fill_at = 5, .ended = true }; + st.leave(t.allocator, reply, &context, &.{}); + const warm = [_]u32{ 1, 2, 3, 4, 30, 23, 31 }; // the reply rendered without its reasoning + try t.expect(!st.replayed(&(warm ++ [_]u32{32}), &.{.{ 5, 6 }}).holds(4)); // its follow-up goes back to 4 + const back = st.plan(&warm, &.{.{ 5, 6 }}, 0); // the server's warm pass over it + try t.expectEqual(Resume{ .back = 4 }, back); + st.take(t.allocator, back); + try t.expectEqual(modes.Span{ 0, 0 }, st.reply); // the rows past the point are rewritten + st.leave(t.allocator, .{ .at = 7, .prompt_len = 7, .context_len = 7, .fill_at = 7, .ended = true }, &warm, &.{.{ 5, 6 }}); + const follow = warm ++ [_]u32{ 32, 33 }; + try t.expectEqual(Resume{ .here = 7 }, st.plan(&follow, &.{.{ 5, 6 }}, 0)); // past the carried point at 4 + try t.expect(!st.replayed(&follow, &.{.{ 5, 6 }}).holds(7)); // prompt rows: the new history's end is a point + st.leave(t.allocator, reply, &context, &.{}); + st.take(t.allocator, .none); + try t.expectEqual(modes.Span{ 0, 0 }, st.reply); + try t.expectEqual(@as(u32, 0), st.rewind_at); +} + test { _ = @import("qwen27_resume_cache_test.zig"); } diff --git a/zig/src/native/qwen27_slots.zig b/zig/src/native/qwen27_slots.zig index e13b6755e1d..d82558d64da 100644 --- a/zig/src/native/qwen27_slots.zig +++ b/zig/src/native/qwen27_slots.zig @@ -35,6 +35,7 @@ pub fn opening(b: *Metal, s: *lanes.Stream, slot: u32) !Fill { try recover(b, s); b.runner.rope[slot] = try images.layout(b.vision, s); const d = b.drafter(slot); + const replay: resume_rules.Replay = if (slot == 0) b.standing.replayed(ids, s.decode_spans) else .{}; // before take const resumed = if (slot == 0) try takeStanding(b, s, kept_at) else 0; // slot 0's own state: nothing copied if (resumed == 0) { if (d) |x| try x.reset() else try b.runner.reset(slot); @@ -58,7 +59,7 @@ pub fn opening(b: *Metal, s: *lanes.Stream, slot: u32) !Fill { s.fill_at.store(@intCast(at), .monotonic); // what the restore kept: 0 when it failed const encoded = if (b.vision) |v| try images.encode(v, b.runner.model.queue, s, slot) else null; if (encoded != null) b.image_fill = s; // the vision rows are this pass's until it closes - return .{ .slot = slot, .at = at, .from = at, .images = encoded, .checkpoint = resume_rules.cachePoint(slot == 0 and slotResume(b) and plain(s) and s.max_new > 0 and s.reuse.hook != null and b.checkpoint_hook != null, at, b.chunk, ids.len, s.history_len, s.reuse.marks, s.chunks, s.decode_spans), .rewind = rewindPoint(b, s, slot, at) }; + return .{ .slot = slot, .at = at, .from = at, .images = encoded, .checkpoint = resume_rules.cachePoint(slot == 0 and slotResume(b) and plain(s) and s.max_new > 0 and s.reuse.hook != null and b.checkpoint_hook != null, at, b.chunk, ids.len, s.history_len, s.reuse.marks, s.chunks, s.decode_spans), .rewind = rewindPoint(b, s, slot, at, replay) }; } /// The pass is over (done, failed or cancelled): its images' rows are free for the next image prompt. @@ -238,7 +239,7 @@ fn resumable(b: *const Metal, s: *const lanes.Stream, kept_at: u32) resume_rules fn takeStanding(b: *Metal, s: *const lanes.Stream, kept_at: u32) !u32 { const plan = resumable(b, s, kept_at); b.standing.take(b.gpa, plan); - errdefer b.standing.rewind_at = 0; + errdefer b.standing = .{ .wide = b.standing.wide }; // tokens and spans freed by take: the point and reply go too // rows [0, rewind_at) and the buffer hold only while slot 0 just extends past it: a restore, reset or move drops it switch (plan) { .none => return 0, @@ -272,9 +273,9 @@ fn leave(b: *Metal, s: *const lanes.Stream) void { b.standing.leave(b.gpa, .{ .at = at, .prompt_len = @intCast(s.prompt_len), .context_len = @intCast(s.context.items.len), .fill_at = s.fill_at.load(.monotonic), .ended = ended }, s.context.items, s.decode_spans); } -/// Where slot 0's pass keeps its rewind point: its rendered history's end, which a next turn shares (0: none). -fn rewindPoint(b: *const Metal, s: *const lanes.Stream, slot: u32, at: usize) u32 { - if (slot != 0 or b.rewind == null or !slotResume(b) or !plain(s)) return 0; +/// Where slot 0's pass from `at` keeps its rewind point: its history's end (0: none, or it resumed in a `replay`). +fn rewindPoint(b: *const Metal, s: *const lanes.Stream, slot: u32, at: usize, replay: resume_rules.Replay) u32 { + if (slot != 0 or b.rewind == null or !slotResume(b) or !plain(s) or replay.holds(at)) return 0; return if (s.history_len > at and s.history_len < s.prompt().len) s.history_len else 0; } From 4957025953c71ac9bbf0248f46a0e957dcfa23da Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 13:24:47 +0000 Subject: [PATCH 49/78] changelog: #2723's line from lines/2723.txt in place of its own --- CHANGELOG.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c8165118cdd..cbd9a4b64fe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Qwen3.8-27B on Metal: slot 0's rewind point is now made even beside a prompt cache that can hold a whole-context state, and the warm pass keeps it, so a client that sends the last reply back without its reasoning resumes in place at its history's end instead of a turn further back (on an M5 Max at 32k with a 5 GiB cache, turns 1-6 of a 7-turn chat took 16.5 s and 18.9 s of summed time to first token against 217 s, with identical replies); a client that sends the reply back as generated makes the same fill calls as before; a background pass at depth runs in narrower chunks, so a request arriving mid-pass at 128k waited 0.5 s against 4.1 s; slot 0's 146 MiB rewind buffer counts against `--prompt-cache-gib` room (#2723). Thanks to @ThinkOffApp. + - Prompt cache on the core pool (the `TF_KV_SHARE` families, Qwen 3.5/3.6 on CUDA, and Qwen 3.5 and Nemotron on HIP): the radix store holds each pool page once, however many entries name it, and pins the node a keep or restore is using. A snapshot save that takes pool pages (Gemma 4, Kolibri-1, Spark2.5, Dory) can no longer make reclaim prune the leaf it is writing to. Replies and cached counts are unchanged. The store also learns checkpoint parts and prefix groups for later steps to use (#2875, part of #2444). Thanks to @laudney. - GLM-5.3-Flash on CUDA at four ranks or more: when the default split prefill's start-up exchange with every rank fails, as on a switchless ring of four, the failing rank now prints a refusal naming `TF_GLM_PREFILL_REDUCE=gather`, and rank 0 names it too when another rank fails first. Before, every rank stopped on a bare `ncclGroupEnd` error. Start-ups that succeed are unchanged (#2868). Thanks to @Bizuayeu. @@ -151,7 +153,6 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - HIP on gfx12 under `--policy matrix=on`: Qwen3 dense checkpoints load each layer's gate and up weights interleaved as one weight, and one WMMA product writes silu(gate) x up in its epilogue instead of two products and `tf_silu_mul`, with the same output bytes: Qwen3-Reranker-8B ranks 16 documents in 0.657 to 0.661 s instead of 0.687 to 0.697 s and Qwen3-Embedding-8B embeds 100 short inputs in 0.901 to 0.911 s instead of 0.940 to 0.948 s on a Radeon AI PRO R9700. The kernel contract checks both fused blocks byte for byte against the unfused launches (part of #2245). Thanks to @laudney. - Nemotron on Macs: the masked prompt shader calls the shared `tf_paged_row` page helper again, so its prompt sources compile; #2322 had used the name #2283 retired the same day (#2548). -- Qwen3.8-27B on Metal: a follow-up turn whose client sends the last reply without its reasoning resumes slot 0 in place from its rewind point, now also beside a prompt cache large enough for two whole contexts (which can still evict the history's end), instead of a turn further back; there the warm pass keeps that point, and a background pass at depth runs in narrower chunks so a waiting request is held for less (#2723). Thanks to @ThinkOffApp. ## 1.2.7 From 7280f2dd0cd2db8ac89172a253e24587145f0f74 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Wed, 11 Nov 2026 22:05:57 +0900 Subject: [PATCH 50/78] cuda source check: declare the SDK names main's kernels use check_cuda_sources.py failed 66 of its 918 checks on main: eleven units use SDK names that the fixed declarations in tools/zig/cuda_host lack. Declare each as CUDA 13.0 does: __syncthreads_or and __double2float_rz from sm_20_intrinsics.h, __float2int_rn from crt/device_functions.h, __float2bfloat16_rz from cuda_bf16.h, __hmul on half from cuda_fp16.h, ldexpf from crt/math_functions.h, and the __maxnreg__ macro from crt/host_defines.h, whose maxnreg attribute Clang ignores with a warning. Only the check reads these headers: no kernel, engine or test binary changes. --- tools/zig/cuda_host/cuda_bf16.h | 1 + tools/zig/cuda_host/cuda_fp16.h | 1 + tools/zig/cuda_host/cuda_runtime.h | 3 +++ tools/zig/cuda_host/math.h | 1 + 4 files changed, 6 insertions(+) diff --git a/tools/zig/cuda_host/cuda_bf16.h b/tools/zig/cuda_host/cuda_bf16.h index 9f96566052c..3d215164188 100644 --- a/tools/zig/cuda_host/cuda_bf16.h +++ b/tools/zig/cuda_host/cuda_bf16.h @@ -4,6 +4,7 @@ struct __nv_bfloat16 { unsigned short x; }; struct __nv_bfloat162 { __nv_bfloat16 x,y; }; __device__ float __bfloat162float(__nv_bfloat16); __device__ __nv_bfloat16 __float2bfloat16(float); __device__ __nv_bfloat16 __float2bfloat16_rn(float); __device__ __nv_bfloat16 __ushort_as_bfloat16(unsigned short); +__device__ __nv_bfloat16 __float2bfloat16_rz(float); __device__ unsigned short __bfloat16_as_ushort(__nv_bfloat16); __device__ float2 __bfloat1622float2(__nv_bfloat162); __device__ __nv_bfloat162 __floats2bfloat162_rn(float,float); diff --git a/tools/zig/cuda_host/cuda_fp16.h b/tools/zig/cuda_host/cuda_fp16.h index d423d31f12d..44654bcdacc 100644 --- a/tools/zig/cuda_host/cuda_fp16.h +++ b/tools/zig/cuda_host/cuda_fp16.h @@ -14,3 +14,4 @@ __device__ half2 __halves2half2(half,half); __device__ half __low2half(half2); _ __device__ float __low2float(half2); __device__ float __high2float(half2); __device__ half2 __hadd2(half2,half2); __device__ half2 __hsub2(half2,half2); __device__ half2 __hmul2(half2,half2); __device__ half2 __hfma2(half2,half2,half2); __device__ half __hadd(half,half); +__device__ half __hmul(half,half); diff --git a/tools/zig/cuda_host/cuda_runtime.h b/tools/zig/cuda_host/cuda_runtime.h index a6fd8663516..2a63e2b086e 100644 --- a/tools/zig/cuda_host/cuda_runtime.h +++ b/tools/zig/cuda_host/cuda_runtime.h @@ -8,6 +8,7 @@ #define __noinline__ __attribute__((noinline)) #define __align__(n) __attribute__((aligned(n))) #define __launch_bounds__(...) __attribute__((launch_bounds(__VA_ARGS__))) +#define __maxnreg__(a) __attribute__((maxnreg(a))) #define __grid_constant__ #define __CUDACC__ 1 #include @@ -31,6 +32,7 @@ template __device__ T __shfl_xor_sync(unsigned,T,int,int=32); template __device__ T __shfl_up_sync(unsigned,T,unsigned,int=32); template __device__ T __shfl_down_sync(unsigned,T,unsigned,int=32); __device__ void __syncthreads(); __device__ int __syncthreads_count(int); __device__ void __syncwarp(unsigned=~0u); +__device__ int __syncthreads_or(int); __device__ void __threadfence(); __device__ void __threadfence_system(); __device__ void __trap(); __device__ void __nanosleep(unsigned); __device__ unsigned long long clock64(); __device__ unsigned __ballot_sync(unsigned,int); __device__ int __any_sync(unsigned,int); @@ -49,6 +51,7 @@ __device__ int __float_as_int(float); __device__ unsigned __float_as_uint(float); __device__ float __int_as_float(int); __device__ float __uint_as_float(unsigned); __device__ double __longlong_as_double(long long); __device__ long long __double2ll_rz(double); __device__ float __double2float_rn(double); +__device__ float __double2float_rz(double); __device__ int __float2int_rn(float); __device__ float __fadd_rn(float,float); __device__ float __fsub_rn(float,float); __device__ float __fmul_rn(float,float); __device__ float __fdiv_rn(float,float); __device__ float __fsqrt_rn(float); __device__ float __fmaf_rn(float,float,float); __device__ float __expf(float); diff --git a/tools/zig/cuda_host/math.h b/tools/zig/cuda_host/math.h index 4aade403d68..65cf5c4b1c8 100644 --- a/tools/zig/cuda_host/math.h +++ b/tools/zig/cuda_host/math.h @@ -19,5 +19,6 @@ template __host__ __device__ bool isfinite(T); #define M_2_SQRTPI 1.12837916709551257390 #define M_SQRT1_2 0.70710678118654752440 __host__ __device__ float erff(float); __host__ __device__ float rintf(float); +__host__ __device__ float ldexpf(float,int); __host__ __device__ double pow(double,double); __host__ __device__ double sin(double); __host__ __device__ double cos(double); template __host__ __device__ bool signbit(T); From 305cecd425291bb85b1e8a4df57b0d8c75a7f9fb Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 13:24:49 +0000 Subject: [PATCH 51/78] changelog: #2902 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index cbd9a4b64fe..92e52fa03c7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Tests: the CUDA source check (`tools/zig/check_cuda_sources.py`) declares the SDK names main's kernels use, so it passes all 918 checks again (#2902). Thanks to @Bizuayeu. + - Qwen3.8-27B on Metal: slot 0's rewind point is now made even beside a prompt cache that can hold a whole-context state, and the warm pass keeps it, so a client that sends the last reply back without its reasoning resumes in place at its history's end instead of a turn further back (on an M5 Max at 32k with a 5 GiB cache, turns 1-6 of a 7-turn chat took 16.5 s and 18.9 s of summed time to first token against 217 s, with identical replies); a client that sends the reply back as generated makes the same fill calls as before; a background pass at depth runs in narrower chunks, so a request arriving mid-pass at 128k waited 0.5 s against 4.1 s; slot 0's 146 MiB rewind buffer counts against `--prompt-cache-gib` room (#2723). Thanks to @ThinkOffApp. - Prompt cache on the core pool (the `TF_KV_SHARE` families, Qwen 3.5/3.6 on CUDA, and Qwen 3.5 and Nemotron on HIP): the radix store holds each pool page once, however many entries name it, and pins the node a keep or restore is using. A snapshot save that takes pool pages (Gemma 4, Kolibri-1, Spark2.5, Dory) can no longer make reclaim prune the leaf it is writing to. Replies and cached counts are unchanged. The store also learns checkpoint parts and prefix groups for later steps to use (#2875, part of #2444). Thanks to @laudney. From 5e1ca3c235948f062c80c01d4c9013abd6997464 Mon Sep 17 00:00:00 2001 From: Shahaf Antwarg <10106174+kotevcode@users.noreply.github.com> Date: Wed, 11 Nov 2026 20:35:35 +0700 Subject: [PATCH 52/78] native: the CUDA kept-state identity is written once through any hasher, its keys unchanged --- CHANGELOG.md | 1 + zig/src/cuda/aot_identity.zig | 9 +- zig/src/native/cuda_disk.zig | 59 +++++++-- zig/src/native/cuda_disk_golden_test.zig | 143 +++++++++++++++++++++ zig/src/native/cuda_disk_identity_test.zig | 4 + zig/src/native/disk_identity.zig | 10 +- zig/src/platform/file_identity.zig | 12 +- 7 files changed, 215 insertions(+), 23 deletions(-) create mode 100644 zig/src/native/cuda_disk_golden_test.zig diff --git a/CHANGELOG.md b/CHANGELOG.md index 92e52fa03c7..b6eddd97525 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -126,6 +126,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Pipeline stages: the shared stage driver pipelines rounds, so rank 0 sends round r + 1 while the later stages still run round r, and a stage that echoes a round short or empty is refused by name; tokens are checked equal to one engine on the same binary (#2579, part of #2005). Thanks to @jschmied. +- Internal: the CUDA kept-state identity (binary, checkpoint stamps, kernel images, captured cubins, GPU, driver, family facts, probe and resolved profile values) is written once through any hasher, so the execution witness can take the same components through SHA-256; every kept key stays the same bytes, pinned against main's sequence and a fixed key (refs #2741). - GLM-5.3-Flash on CUDA: tensor-parallel shares and layer sums go through one mesh layer (declared shares, one reduction boundary over a `Collective`), with the same gathers, bytes and order and tokens equal to main at TP=2; an aborted `tp_group` call now ends as `Aborted` instead of stopping the process (#2789, part of #2560). Thanks to @laudney. - Tensor parallel on CUDA: with `TF_ROCE=1` on every rank, three ranks or more take the one-shot RoCE all-gather for decode-sized gathers, byte-equal to NCCL's, falling back to NCCL with one warning when a peer can't be reached, and GLM-5.3-Flash's DFlash2 drafter MLP splits across every rank; without it they keep NCCL, two ranks run as before, and MiniMax-M3 keeps NCCL at three ranks or more. With `TF_ROCE=1`, GLM-5.3-Flash decode is about 6% faster at one and two streams on three GB10s (four streams level) and 12.9%, 10.4% and 6.6% faster at one, two and four streams on four GB10s, with the target's tokens unchanged; drafter acceptance can move slightly, since the split MLP sums in another order (#2205). Thanks to @shanedertrain. diff --git a/zig/src/cuda/aot_identity.zig b/zig/src/cuda/aot_identity.zig index 6c6e8e001af..a1b21900072 100644 --- a/zig/src/cuda/aot_identity.zig +++ b/zig/src/cuda/aot_identity.zig @@ -3,13 +3,19 @@ const std = @import("std"); const aot = @import("aot.zig"); pub fn hash(a: std.mem.Allocator, io: std.Io, dir: []const u8) !u64 { + var h = std.hash.Wyhash.init(0x616f74); + try write(&h, a, io, dir); + return h.final(); +} + +/// The manifest's text, then each cubin it names by its SHA-256, into any hasher with `update(bytes)`. +pub fn write(h: anytype, a: std.mem.Allocator, io: std.Io, dir: []const u8) !void { const path = try std.fs.path.join(a, &.{ dir, "aot.json" }); defer a.free(path); const text = try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 24)); defer a.free(text); const parsed = try aot.parseSpecs(a, text); defer parsed.deinit(); - var h = std.hash.Wyhash.init(0x616f74); h.update(text); for (parsed.value.kernels) |k| { const cubin = try std.fmt.allocPrint(a, "{s}/cubins/{s}.cubin", .{ dir, k.hash }); @@ -17,7 +23,6 @@ pub fn hash(a: std.mem.Allocator, io: std.Io, dir: []const u8) !u64 { const digest = try @import("sys").file_identity.hash(io, cubin, 1 << 26); h.update(&digest); } - return h.final(); } test "persistence identity: captured CUDA cubin replacement moves an unchanged manifest key" { diff --git a/zig/src/native/cuda_disk.zig b/zig/src/native/cuda_disk.zig index 4ff1ee1e14b..0c6c03d8219 100644 --- a/zig/src/native/cuda_disk.zig +++ b/zig/src/native/cuda_disk.zig @@ -133,24 +133,57 @@ pub fn keptKey(comptime K: type, a: Allocator, io: std.Io, g: *const cuda.Contex /// The identity's hash: the checkpoint, kernels, device, driver, a family's facts, then its probe's hash. pub fn identityOf(a: Allocator, io: std.Io, g: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, facts: []const u64, probed: u64, profile: disk_identity.Profile) !u64 { var h = std.hash.Wyhash.init(0x6375); - try @import("sys").file_identity.current(&h, a, io); - try disk_identity.checkpoint(&h, a, io, dir); + const binary = try @import("sys").file_identity.running(a, io); + var name_buf: [256]u8 = undefined; + try components(&h, a, io, .{ .binary = binary, .checkpoint = dir, .kernels = kernels, .device = try deviceOf(g, &name_buf), .facts = facts, .probed = probed, .profile = profile }); + return h.final(); +} + +/// The GPU's name (or "GPU" when the driver gives none), compute capability and driver version. +pub const Device = struct { name: []const u8, capability: u64, driver: u64 }; + +pub fn deviceOf(g: *const cuda.Context, name_buf: []u8) !Device { + const name = g.name(name_buf) catch "GPU"; + return .{ .name = name, .capability = try g.capability(), .driver = @intCast(try g.d.version()) }; +} + +/// What a kept key binds; a null part is bound otherwise by its caller (the witness binds bytes, not stamps). +pub const Components = struct { + binary: ?[32]u8, + checkpoint: ?[]const u8, + kernels: ?[]const u8, + device: Device, + facts: []const u64, + probed: ?u64, + profile: disk_identity.Profile, +}; + +/// The components in one order into any hasher with `update(bytes)`: a kept key's Wyhash, or a SHA-256. +pub fn components(h: anytype, a: Allocator, io: std.Io, c: Components) !void { + if (c.binary) |digest| @import("sys").file_identity.bind(h, digest); + if (c.checkpoint) |dir| try disk_identity.checkpoint(h, a, io, dir); inline for (comptime std.meta.declarations(cuda.kernels)) |d| { const v = @field(cuda.kernels, d); if (comptime @TypeOf(v) == []const u8) h.update(v); } - if (kernels) |k| { - const captured = try cuda.aot.identity.hash(a, io, k); - h.update(std.mem.asBytes(&captured)); - } - var name_buf: [256]u8 = undefined; - h.update(g.name(&name_buf) catch "GPU"); - const device = [_]u64{ try g.capability(), @intCast(try g.d.version()) }; + if (c.kernels) |k| try captured(h, a, io, k); + h.update(c.device.name); + const device = [_]u64{ c.device.capability, c.device.driver }; h.update(std.mem.sliceAsBytes(&device)); - h.update(std.mem.sliceAsBytes(facts)); - h.update(std.mem.asBytes(&probed)); - disk_identity.profile(&h, profile); - return h.final(); + h.update(std.mem.sliceAsBytes(c.facts)); + if (c.probed) |p| h.update(std.mem.asBytes(&p)); + disk_identity.profile(h, c.profile); +} + +/// A captured set as a nested digest: a kept key folds its 64-bit Wyhash, a SHA-256 its SHA-256. +fn captured(h: anytype, a: Allocator, io: std.Io, dir: []const u8) !void { + if (@TypeOf(h.*) == std.hash.Wyhash) { + const v = try cuda.aot.identity.hash(a, io, dir); + return h.update(std.mem.asBytes(&v)); + } + var s = std.crypto.hash.sha2.Sha256.init(.{}); + try cuda.aot.identity.write(&s, a, io, dir); + h.update(&s.finalResult()); } /// A fixed prompt's state at a mark inside its last chunk, hashed (a kernel, driver or arithmetic change moves it). diff --git a/zig/src/native/cuda_disk_golden_test.zig b/zig/src/native/cuda_disk_golden_test.zig new file mode 100644 index 00000000000..9e5ce7b5a12 --- /dev/null +++ b/zig/src/native/cuda_disk_golden_test.zig @@ -0,0 +1,143 @@ +//! Kept keys stay byte for byte: the identity, however it is written, matches the sequence main hashed. +const std = @import("std"); +const cuda = @import("cuda"); +const api = @import("engine_api"); +const disk = @import("cuda_disk.zig"); +const identity = @import("disk_identity.zig"); +const t = std.testing; + +const Driver = struct { + fn name(out: [*]u8, _: c_int, _: cuda.abi.Device) callconv(.c) cuda.abi.Result { + @memcpy(out[0..5], "fake\x00"); + return 0; + } + fn attribute(out: *c_int, which: cuda.abi.DeviceAttribute, _: cuda.abi.Device) callconv(.c) cuda.abi.Result { + out.* = if (which == .compute_capability_major) 12 else 1; + return 0; + } + fn version(out: *c_int) callconv(.c) cuda.abi.Result { + out.* = 13000; + return 0; + } + fn environment(_: [:0]const u8) ?[:0]const u8 { + return "1"; + } +}; + +const row: api.profile.Entry = .{ .id = .TF_QWEN_KV_PAGES_ALL, .families = &.{"qwen3_5"}, .class = .arithmetic, .syntax = .one, .effect = "test arithmetic", .evidence = "test fixture" }; +const profile: identity.Profile = .{ .family = "qwen3_5", .backend = "cuda", .entries = &.{row}, .environment = Driver.environment }; +/// The key main (93a2e93ec) gives the fixed inputs below. +const main_key: u64 = 9113307822166948912; +const manifest = "{\"kernels\":[{\"fn\":\"f\",\"hash\":\"h\",\"name\":\"k\",\"num_warps\":4,\"params\":[],\"consts\":{}}]}"; + +/// main's identityOf, update for update, with the binary's digest and the device's facts given. +fn reference(a: std.mem.Allocator, binary: [32]u8, dir: []const u8, kernels: ?[]const u8, facts: []const u64, probed: u64) !u64 { + var h = std.hash.Wyhash.init(0x6375); + @import("sys").file_identity.bind(&h, binary); + try identity.checkpoint(&h, a, t.io, dir); + inline for (comptime std.meta.declarations(cuda.kernels)) |d| { + const v = @field(cuda.kernels, d); + if (comptime @TypeOf(v) == []const u8) h.update(v); + } + if (kernels) |k| { + const captured = try cuda.aot.identity.hash(a, t.io, k); + h.update(std.mem.asBytes(&captured)); + } + h.update("fake"); + const device = [_]u64{ 121, 13000 }; + h.update(std.mem.sliceAsBytes(&device)); + h.update(std.mem.sliceAsBytes(facts)); + h.update(std.mem.asBytes(&probed)); + identity.profile(&h, profile); + return h.final(); +} + +/// A checkpoint folder whose identity reads no stamps (config.json alone), and a captured set beside it. +const Fixture = struct { + tmp: t.TmpDir, + dir: []const u8, + kernels: []const u8, + + fn init(a: std.mem.Allocator) !Fixture { + var tmp = t.tmpDir(.{}); + errdefer tmp.cleanup(); + try tmp.dir.createDir(t.io, "model", .default_dir); + try tmp.dir.writeFile(t.io, .{ .sub_path = "model/config.json", .data = "{\"model_type\":\"qwen3_5\"}" }); + try tmp.dir.createDir(t.io, "aot", .default_dir); + try tmp.dir.createDir(t.io, "aot/cubins", .default_dir); + try tmp.dir.writeFile(t.io, .{ .sub_path = "aot/aot.json", .data = manifest }); + try tmp.dir.writeFile(t.io, .{ .sub_path = "aot/cubins/h.cubin", .data = "cubin bytes" }); + const root = try tmp.dir.realPathFileAlloc(t.io, ".", a); + return .{ .tmp = tmp, .dir = try std.fs.path.join(a, &.{ root, "model" }), .kernels = try std.fs.path.join(a, &.{ root, "aot" }) }; + } +}; + +test "disk identity: identityOf hashes main's sequence over this binary, a checkpoint and a captured set" { + var arena: std.heap.ArenaAllocator = .init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var f = try Fixture.init(a); + defer f.tmp.cleanup(); + try f.tmp.dir.writeFile(t.io, .{ .sub_path = "model/model.safetensors", .data = "shard" }); // its stamps are hashed too + var driver: cuda.Driver = undefined; + driver.api.cuDeviceGetName = Driver.name; + driver.api.cuDeviceGetAttribute = Driver.attribute; + driver.api.cuDriverGetVersion = Driver.version; + const gpu: cuda.Context = .{ .d = &driver, .device = 0, .handle = null }; + const exe = try std.process.executablePathAlloc(t.io, a); + const binary = try @import("sys").file_identity.hash(t.io, exe, 1 << 31); + for ([_]?[]const u8{ null, f.kernels }) |kernels| + try t.expectEqual(try reference(a, binary, f.dir, kernels, &.{ 1, 2, 3 }, 42), try disk.identityOf(a, t.io, &gpu, f.dir, kernels, &.{ 1, 2, 3 }, 42, profile)); +} + +test "disk identity: a fixed binary, checkpoint, captured set, device, facts, probe and profile give main's key" { + if (cuda.kernels.available) return error.SkipZigTest; // the constant is a host build's, whose kernel images are empty + var arena: std.heap.ArenaAllocator = .init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var f = try Fixture.init(a); + defer f.tmp.cleanup(); + try t.expectEqual(main_key, try reference(a, @splat(0x5a), f.dir, f.kernels, &.{ 1, 2, 3 }, 42)); + var h = std.hash.Wyhash.init(0x6375); + try disk.components(&h, a, t.io, fixed(f)); + try t.expectEqual(main_key, h.final()); +} + +fn fixed(f: Fixture) disk.Components { + return .{ .binary = @splat(0x5a), .checkpoint = f.dir, .kernels = f.kernels, .device = .{ .name = "fake", .capability = 121, .driver = 13000 }, .facts = &.{ 1, 2, 3 }, .probed = 42, .profile = profile }; +} + +fn sha(a: std.mem.Allocator, c: disk.Components) ![32]u8 { + var h = std.crypto.hash.sha2.Sha256.init(.{}); + try disk.components(&h, a, t.io, c); + return h.finalResult(); +} + +test "disk identity: the same components through SHA-256 move with a captured cubin, a profile value and the device" { + var arena: std.heap.ArenaAllocator = .init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var f = try Fixture.init(a); + defer f.tmp.cleanup(); + var c = fixed(f); + c.binary = null; // the witness binds the binary and the checkpoint's bytes elsewhere + c.checkpoint = null; + c.probed = null; + const base = try sha(a, c); + try t.expectEqualSlices(u8, &base, &try sha(a, c)); + try f.tmp.dir.writeFile(t.io, .{ .sub_path = "aot/cubins/h.cubin", .data = "other cubin" }); + const cubin = try sha(a, c); + try t.expect(!std.mem.eql(u8, &base, &cubin)); + var off = c; + off.profile.environment = Unset.environment; + try t.expect(!std.mem.eql(u8, &cubin, &try sha(a, off))); + var other = c; + other.device.capability = 80; + try t.expect(!std.mem.eql(u8, &cubin, &try sha(a, other))); +} + +const Unset = struct { + fn environment(_: [:0]const u8) ?[:0]const u8 { + return null; + } +}; diff --git a/zig/src/native/cuda_disk_identity_test.zig b/zig/src/native/cuda_disk_identity_test.zig index 107fd664053..7bfcdf9a5b6 100644 --- a/zig/src/native/cuda_disk_identity_test.zig +++ b/zig/src/native/cuda_disk_identity_test.zig @@ -6,6 +6,10 @@ const disk = @import("cuda_disk.zig"); const identity = @import("disk_identity.zig"); const t = std.testing; +test { + _ = @import("cuda_disk_golden_test.zig"); +} + const Recording = struct { var calls: [16]enum { name, major, minor, version } = undefined; var count: usize = 0; diff --git a/zig/src/native/disk_identity.zig b/zig/src/native/disk_identity.zig index 3126a14864b..46ff8cb44a2 100644 --- a/zig/src/native/disk_identity.zig +++ b/zig/src/native/disk_identity.zig @@ -4,7 +4,7 @@ const api = @import("engine_api"); const Allocator = std.mem.Allocator; /// Config/index bytes, stamps of every indexed shard (IndexedShardMissing if absent) and top-level safetensors. -pub fn checkpoint(h: *std.hash.Wyhash, a: Allocator, io: std.Io, dir: []const u8) !void { +pub fn checkpoint(h: anytype, a: Allocator, io: std.Io, dir: []const u8) !void { var arena: std.heap.ArenaAllocator = .init(a); defer arena.deinit(); const scratch = arena.allocator(); @@ -46,7 +46,7 @@ pub fn checkpoint(h: *std.hash.Wyhash, a: Allocator, io: std.Io, dir: []const u8 } /// This Mac's OS build and chip: a macOS update or another chip can compile a kernel to other bits. -pub fn machine(h: *std.hash.Wyhash, chip: []const u8) void { +pub fn machine(h: anytype, chip: []const u8) void { var os: [64]u8 = undefined; h.update(api.prompt_imprint.osBuild(&os)); h.update(chip); @@ -65,7 +65,7 @@ pub const Profile = struct { }; /// Only applicable non-identical entries partition kept state; defaults and `on_with` resolve before values are hashed. -pub fn profile(h: *std.hash.Wyhash, p: Profile) void { +pub fn profile(h: anytype, p: Profile) void { var began = false; for (p.entries) |e| { if (e.class == .identical or !e.applies(p.family, p.backend)) continue; @@ -95,7 +95,7 @@ pub fn profile(h: *std.hash.Wyhash, p: Profile) void { } } -fn valueField(h: *std.hash.Wyhash, syntax: api.profile.Syntax, value: ?[]const u8) void { +fn valueField(h: anytype, syntax: api.profile.Syntax, value: ?[]const u8) void { switch (syntax) { .one => field(h, if (value != null and std.mem.eql(u8, value.?, "1")) "1" else "0"), .auto => field(h, value orelse "auto"), @@ -117,7 +117,7 @@ fn valueField(h: *std.hash.Wyhash, syntax: api.profile.Syntax, value: ?[]const u } } -fn field(h: *std.hash.Wyhash, bytes: []const u8) void { +fn field(h: anytype, bytes: []const u8) void { var len: [8]u8 = undefined; std.mem.writeInt(u64, &len, @intCast(bytes.len), .little); h.update(&len); diff --git a/zig/src/platform/file_identity.zig b/zig/src/platform/file_identity.zig index fc40f934da8..4c126ed767a 100644 --- a/zig/src/platform/file_identity.zig +++ b/zig/src/platform/file_identity.zig @@ -22,14 +22,20 @@ pub fn hash(io: std.Io, path: []const u8, maximum: u64) ![32]u8 { return h.finalResult(); } -pub fn bind(h: *std.hash.Wyhash, digest: [32]u8) void { +/// `h`: any hasher with `update(bytes)`, a kept key's Wyhash or a SHA-256. +pub fn bind(h: anytype, digest: [32]u8) void { h.update("serving-binary-v1"); h.update(&digest); } /// The running binary includes compiled host code, embedded GPU images, compiler options and snapshot layouts. -pub fn current(h: *std.hash.Wyhash, a: std.mem.Allocator, io: std.Io) !void { +pub fn current(h: anytype, a: std.mem.Allocator, io: std.Io) !void { + bind(h, try running(a, io)); +} + +/// The running binary's SHA-256, read by its path. +pub fn running(a: std.mem.Allocator, io: std.Io) ![32]u8 { const path = try std.process.executablePathAlloc(io, a); defer a.free(path); - bind(h, try hash(io, path, 1 << 31)); + return hash(io, path, 1 << 31); } From 08295050e182c9e97f6bf79a186f02a7bc4559d2 Mon Sep 17 00:00:00 2001 From: Shahaf Antwarg <10106174+kotevcode@users.noreply.github.com> Date: Wed, 11 Nov 2026 20:54:52 +0700 Subject: [PATCH 53/78] changelog: #2903's line names its pull request --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index b6eddd97525..c2619030b6f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -127,6 +127,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Pipeline stages: the shared stage driver pipelines rounds, so rank 0 sends round r + 1 while the later stages still run round r, and a stage that echoes a round short or empty is refused by name; tokens are checked equal to one engine on the same binary (#2579, part of #2005). Thanks to @jschmied. - Internal: the CUDA kept-state identity (binary, checkpoint stamps, kernel images, captured cubins, GPU, driver, family facts, probe and resolved profile values) is written once through any hasher, so the execution witness can take the same components through SHA-256; every kept key stays the same bytes, pinned against main's sequence and a fixed key (refs #2741). +- Internal: the CUDA kept-state identity (binary, checkpoint stamps, kernel images, captured cubins, GPU, driver, family facts, probe and resolved profile values) is written once through any hasher, so the execution witness can take the same components through SHA-256; every kept key stays the same bytes, pinned against main's sequence and a fixed key (#2903, refs #2741). Thanks to @kotevcode. - GLM-5.3-Flash on CUDA: tensor-parallel shares and layer sums go through one mesh layer (declared shares, one reduction boundary over a `Collective`), with the same gathers, bytes and order and tokens equal to main at TP=2; an aborted `tp_group` call now ends as `Aborted` instead of stopping the process (#2789, part of #2560). Thanks to @laudney. - Tensor parallel on CUDA: with `TF_ROCE=1` on every rank, three ranks or more take the one-shot RoCE all-gather for decode-sized gathers, byte-equal to NCCL's, falling back to NCCL with one warning when a peer can't be reached, and GLM-5.3-Flash's DFlash2 drafter MLP splits across every rank; without it they keep NCCL, two ranks run as before, and MiniMax-M3 keeps NCCL at three ranks or more. With `TF_ROCE=1`, GLM-5.3-Flash decode is about 6% faster at one and two streams on three GB10s (four streams level) and 12.9%, 10.4% and 6.6% faster at one, two and four streams on four GB10s, with the target's tokens unchanged; drafter acceptance can move slightly, since the split MLP sums in another order (#2205). Thanks to @shanedertrain. From 2806c61c69f72d44459e3b6682ee4ff9b1ddafa6 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 13:24:51 +0000 Subject: [PATCH 54/78] changelog: #2903's line from lines/2903.txt in place of its own --- CHANGELOG.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c2619030b6f..bebfd2ca485 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Internal: the CUDA kept prompt-state identity is written once through any hasher, so the coming execution witness can hash the same parts with SHA-256; every kept key stays byte for byte what it was, pinned by tests against 1.2.7's sequence and a fixed key, so saved prompt states still reopen. Nothing a user sees changes (#2903, part of #2741). Thanks to @kotevcode. + - Tests: the CUDA source check (`tools/zig/check_cuda_sources.py`) declares the SDK names main's kernels use, so it passes all 918 checks again (#2902). Thanks to @Bizuayeu. - Qwen3.8-27B on Metal: slot 0's rewind point is now made even beside a prompt cache that can hold a whole-context state, and the warm pass keeps it, so a client that sends the last reply back without its reasoning resumes in place at its history's end instead of a turn further back (on an M5 Max at 32k with a 5 GiB cache, turns 1-6 of a 7-turn chat took 16.5 s and 18.9 s of summed time to first token against 217 s, with identical replies); a client that sends the reply back as generated makes the same fill calls as before; a background pass at depth runs in narrower chunks, so a request arriving mid-pass at 128k waited 0.5 s against 4.1 s; slot 0's 146 MiB rewind buffer counts against `--prompt-cache-gib` room (#2723). Thanks to @ThinkOffApp. @@ -126,8 +128,6 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Pipeline stages: the shared stage driver pipelines rounds, so rank 0 sends round r + 1 while the later stages still run round r, and a stage that echoes a round short or empty is refused by name; tokens are checked equal to one engine on the same binary (#2579, part of #2005). Thanks to @jschmied. -- Internal: the CUDA kept-state identity (binary, checkpoint stamps, kernel images, captured cubins, GPU, driver, family facts, probe and resolved profile values) is written once through any hasher, so the execution witness can take the same components through SHA-256; every kept key stays the same bytes, pinned against main's sequence and a fixed key (refs #2741). -- Internal: the CUDA kept-state identity (binary, checkpoint stamps, kernel images, captured cubins, GPU, driver, family facts, probe and resolved profile values) is written once through any hasher, so the execution witness can take the same components through SHA-256; every kept key stays the same bytes, pinned against main's sequence and a fixed key (#2903, refs #2741). Thanks to @kotevcode. - GLM-5.3-Flash on CUDA: tensor-parallel shares and layer sums go through one mesh layer (declared shares, one reduction boundary over a `Collective`), with the same gathers, bytes and order and tokens equal to main at TP=2; an aborted `tp_group` call now ends as `Aborted` instead of stopping the process (#2789, part of #2560). Thanks to @laudney. - Tensor parallel on CUDA: with `TF_ROCE=1` on every rank, three ranks or more take the one-shot RoCE all-gather for decode-sized gathers, byte-equal to NCCL's, falling back to NCCL with one warning when a peer can't be reached, and GLM-5.3-Flash's DFlash2 drafter MLP splits across every rank; without it they keep NCCL, two ranks run as before, and MiniMax-M3 keeps NCCL at three ranks or more. With `TF_ROCE=1`, GLM-5.3-Flash decode is about 6% faster at one and two streams on three GB10s (four streams level) and 12.9%, 10.4% and 6.6% faster at one, two and four streams on four GB10s, with the target's tokens unchanged; drafter acceptance can move slightly, since the split MLP sums in another order (#2205). Thanks to @shanedertrain. From 070a97c600b98996bdfa4191503bf13f9417ef0e Mon Sep 17 00:00:00 2001 From: Ash <71890694+ashhart@users.noreply.github.com> Date: Wed, 11 Nov 2026 14:27:55 +0000 Subject: [PATCH 55/78] docs(xpu): the B70 prefill row predates 1.1.0's precise prompt DeltaNet (#2859) --- docs/xpu/README.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/docs/xpu/README.md b/docs/xpu/README.md index bd649f64412..edddef11648 100644 --- a/docs/xpu/README.md +++ b/docs/xpu/README.md @@ -90,6 +90,11 @@ Options through environment variables: | 2 / 4 / 8 simultaneous streams, chat, greedy (aggregate) | 151 / 239 / 300 | 148 / 215 / 208 | 162 / 230 / 224 | 162 / 249 / 227 | | `bench_openai` decode, 1 / 2 / 4 / 8 requests (aggregate) | 173 / 243 / 343 / 350 | 165 / 300 / 337 / 439 | 165 / 336 / 367 / 456 | 179 / 325 / 364 / 463 | +The prefill row predates 1.1.0, which feeds the prompt DeltaNet's fp32 intermediates to the matrix units at close to fp32 +precision (`489baa2673`, `17facba3f8`). On a B70 with GPTQ g128 that costs about 5% of a 32,768-token prompt in the +raw runner (`xpu-qwen27 --bench=32768`: 1,647 to 1,571 tok/s, tokens unchanged; #2859). The table's sizes have not been +re-measured since. + Before the per-stream DFlash2 drafts (copies only in the shared rounds), 2 / 4 / 8 code streams ran at 64 / 120 / 182 tok/s in GPTQ. Validation: `zig build test`, oracle, drafted == serial, multi-stream == serial (chains and trees), and `bench_concurrent --alone --serial` on 1, 2, 4 and 8 streams, greedy and sampled, for the four From 077388fb0c947eea565ffbd81388c2a5c2cb515c Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 13:24:53 +0000 Subject: [PATCH 56/78] changelog: #2906 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index bebfd2ca485..9381b4e5820 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Docs: the Intel Arc (XPU) page notes that its B70 prefill row predates 1.1.0's more precise prompt DeltaNet, which costs about 5% of a 32,768-token prompt with tokens unchanged (#2906, measured in #2859). Thanks to @dinadur. + - Internal: the CUDA kept prompt-state identity is written once through any hasher, so the coming execution witness can hash the same parts with SHA-256; every kept key stays byte for byte what it was, pinned by tests against 1.2.7's sequence and a fixed key, so saved prompt states still reopen. Nothing a user sees changes (#2903, part of #2741). Thanks to @kotevcode. - Tests: the CUDA source check (`tools/zig/check_cuda_sources.py`) declares the SDK names main's kernels use, so it passes all 918 checks again (#2902). Thanks to @Bizuayeu. From e11a0bad750d4a914cdc37e66d3c280d96e90bf5 Mon Sep 17 00:00:00 2001 From: Scott Biggart Date: Tue, 10 Nov 2026 18:31:24 -0800 Subject: [PATCH 57/78] tools: tfvideo_frames.py checks libtfvideo's RGB frames against PyAV's The video byte gate, tf-glm-video with glm_video_fixtures.py, builds only on macOS and needs a checkpoint's processor. This tool needs only a libtfvideo build and clips: it decodes every frame, GLM's 2 frames a second, and the first three and last frames, hashes each frame's RGB rows, and exits nonzero when any frame or return code differs from PyAV's or from a second build's. A clip that will not open or decode fails and the run goes on. --time-only times the decoder alone. Nothing in the engine or the build calls it, so no served path changes. --- tools/zig/tfvideo_frames.py | 177 ++++++++++++++++++++++++++++++++++++ 1 file changed, 177 insertions(+) create mode 100644 tools/zig/tfvideo_frames.py diff --git a/tools/zig/tfvideo_frames.py b/tools/zig/tfvideo_frames.py new file mode 100644 index 00000000000..d53b11506a6 --- /dev/null +++ b/tools/zig/tfvideo_frames.py @@ -0,0 +1,177 @@ +"""libtfvideo's RGB frames against PyAV's (read_video_pyav) or another libtfvideo build, byte for byte, per clip.""" + +import argparse +import ctypes +import hashlib +import json +import sys +import time +from pathlib import Path + +TF_DECODE, TF_SHORT = -5, -6 + + +class Info(ctypes.Structure): + _fields_ = [("frames", ctypes.c_int64), ("rate_num", ctypes.c_int), ("rate_den", ctypes.c_int), + ("width", ctypes.c_int), ("height", ctypes.c_int), ("counted", ctypes.c_int)] + + +FRAME_FN = ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_void_p, ctypes.c_int, ctypes.c_int, + ctypes.POINTER(ctypes.c_uint8), ctypes.c_int, ctypes.c_int, ctypes.c_int) + + +class Library: + """One libtfvideo build, loaded by path so two builds can share the process.""" + + def __init__(self, path: Path): + self.name = str(path) + self.lib = ctypes.CDLL(str(path.resolve()), mode=ctypes.RTLD_LOCAL) + self.lib.tf_video_open.argtypes = [ctypes.c_char_p, ctypes.c_size_t, ctypes.POINTER(ctypes.c_void_p), + ctypes.POINTER(Info)] + self.lib.tf_video_frames.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_int), ctypes.c_int, FRAME_FN, + ctypes.c_void_p] + self.lib.tf_video_close.argtypes = [ctypes.c_void_p] + self.lib.tf_video_version.restype = ctypes.c_char_p + self.version = self.lib.tf_video_version().decode() + + def open(self, data: bytes) -> tuple[int, Info]: + """tf_video_open's return code and the stream's facts.""" + handle, info = ctypes.c_void_p(), Info() + rc = self.lib.tf_video_open(data, len(data), ctypes.byref(handle), ctypes.byref(info)) + if rc == 0: + self.lib.tf_video_close(handle) + return rc, info + + def frames(self, data: bytes, indices: list[int], hashed: bool = True) -> tuple[list, float, int]: + """The frames delivered (hashed, or index only), the seconds tf_video_frames took and its return code.""" + handle, info = ctypes.c_void_p(), Info() + rc = self.lib.tf_video_open(data, len(data), ctypes.byref(handle), ctypes.byref(info)) + if rc != 0: + return [], 0.0, rc + out: list[tuple[int, str]] = [] + + def take(_ctx, _k, index, rgb, width, height, stride): + if not hashed: # timing: the bytes are left alone, so the decoder's threads never wait on Python + out.append((index, "")) + return 0 + raw = ctypes.string_at(rgb, stride * (height - 1) + width * 3) + out.append((index, hashlib.sha256(b"".join(raw[r * stride:r * stride + width * 3] + for r in range(height))).hexdigest())) + return 0 + + callback = FRAME_FN(take) + start = time.perf_counter() + rc = self.lib.tf_video_frames(handle, (ctypes.c_int * len(indices))(*indices), len(indices), callback, None) + seconds = time.perf_counter() - start + self.lib.tf_video_close(handle) + return out, seconds, rc + + +class PyAV: + """transformers' read_video_pyav: decode from the start, keep the wanted indices, to_ndarray(format="rgb24").""" + + def __init__(self): + import av + self.av, self.name, self.version = av, f"PyAV {av.__version__}", str(av.ffmpeg_version_info) + + def open(self, data: bytes) -> tuple[int, None]: + """0 when PyAV opens the container with a video stream, -2 (TF_NO_STREAM) without one, else -1 (TF_INVALID).""" + import io + try: + with self.av.open(io.BytesIO(data)) as c: + return (0 if c.streams.video else -2), None + except self.av.error.FFmpegError: + return -1, None + + def frames(self, data: bytes, indices: list[int], hashed: bool = True) -> tuple[list, float, int]: + """As Library.frames; an FFmpeg error counts as TF_DECODE and an early end as TF_SHORT.""" + import io + want, out, rc = set(indices), [], 0 + start = time.perf_counter() + try: + with self.av.open(io.BytesIO(data)) as c: + c.seek(0) + for i, frame in enumerate(c.decode(video=0)): + if i > indices[-1]: + break + if i in want: + rgb = frame.to_ndarray(format="rgb24") + out.append((i, hashlib.sha256(rgb.tobytes()).hexdigest() if hashed else "")) + except self.av.error.FFmpegError: + rc = TF_DECODE + seconds = time.perf_counter() - start + return out, seconds, rc or (len(out) if len(out) == len(indices) else TF_SHORT) + + +def index_sets(info: Info) -> dict[str, list[int]]: + """Every frame, GLM's 2 frames a second, the first three (an early stop) and the last (the end-of-stream drain).""" + n, rate = int(info.frames), info.rate_num / info.rate_den if info.rate_den else 0.0 + two = sorted({min(n - 1, round(i * rate / 2)) for i in range(int(n / rate * 2) + 1)}) if rate else [0] + return {"all": list(range(n)), "2fps": two, "first3": list(range(min(3, n))), "last": [n - 1]} + + +def best(source, data: bytes, indices: list[int], repeat: int, hashed: bool) -> tuple[list, float, int]: + frames, seconds, rc = source.frames(data, indices, hashed) + for _ in range(repeat - 1): + seconds = min(seconds, source.frames(data, indices, hashed)[1]) + return frames, seconds, rc + + +def compare(lib, ref, clip: Path, repeat: int, rows: dict) -> int: + """One clip's sets against the reference (or timed alone when ref is None); returns the sets that differ.""" + data = clip.read_bytes() + rc, info = lib.open(data) + ref_rc, ref_info = ref.open(data) if ref else (rc, info) + if rc != 0 or ref_rc != 0 or info.frames <= 0: # a damaged clip may not open, or may count no frames + rows["open"] = {"test_rc": rc, "reference_rc": ref_rc, "frames": info.frames} + theirs = (ref_info or info).frames # PyAV counts nothing here: libtfvideo's count stands for both + ok = ref is not None and rc == ref_rc and (info.frames > 0) == (theirs > 0) + print(f"{'ok ' if ok else 'FAIL'} {clip} open: test rc {rc}, reference rc {ref_rc}, {info.frames} frames") + return int(not ok) + failures = 0 + for name, indices in index_sets(info).items(): + mine, mine_s, mine_rc = best(lib, data, indices, repeat, ref is not None) + theirs, theirs_s, theirs_rc = best(ref, data, indices, repeat, True) if ref else (mine, 0.0, mine_rc) + same = sum(a == b for a, b in zip(mine, theirs)) + ok = mine == theirs and mine_rc == theirs_rc and (ref is not None or mine_rc == len(indices)) + failures += not ok + rows[name] = {"indices": indices, "test": mine, "reference": theirs if ref else None, "test_rc": mine_rc, + "reference_rc": theirs_rc, "test_s": mine_s, "reference_s": theirs_s} + codes = f", rc {mine_rc}/{theirs_rc}" if (mine_rc, theirs_rc) != (len(indices), len(indices)) else "" + timing = f"test {mine_s:.3f} s" + (f", reference {theirs_s:.3f} s" if ref else "") + what = "frames equal" if ref else "frames delivered" + print(f"{'ok ' if ok else 'FAIL'} {clip} {name}: {same}/{len(indices)} {what}{codes}, {timing}") + return failures + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--lib", required=True, type=Path, help="the libtfvideo build under test") + ap.add_argument("--against", type=Path, help="another libtfvideo build; PyAV when omitted") + ap.add_argument("--time-only", action="store_true", + help="no reference and no hashing: time tf_video_frames alone (its callback leaves the bytes)") + ap.add_argument("--repeat", type=int, default=1, help="decodes per set; the fastest time is reported") + ap.add_argument("--json", type=Path, help="write every frame's hash, return code and each set's time here") + ap.add_argument("clips", nargs="+", type=Path) + args = ap.parse_args() + lib = Library(args.lib) + ref = None if args.time_only else Library(args.against) if args.against else PyAV() + print(f"test {lib.name} (FFmpeg {lib.version}); reference " + + (f"{ref.name} (FFmpeg {ref.version})" if ref else "none (--time-only)")) + failures, report = 0, {"lib": lib.name, "reference": ref.name if ref else None, "clips": {}} + for clip in args.clips: + rows = report["clips"][str(clip)] = {} + try: + failures += compare(lib, ref, clip, args.repeat, rows) + except Exception as e: # one unreadable clip is a failed clip, not the end of the run + rows["error"] = repr(e) + failures += 1 + print(f"FAIL {clip}: {e!r}") + if args.json: + args.json.write_text(json.dumps(report, indent=1)) + print(f"{failures} set(s) differ" if failures else "every frame equal") + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(main()) From 87dd2f61786e503021b4dd3707265335544b921c Mon Sep 17 00:00:00 2001 From: Scott Biggart Date: Tue, 10 Nov 2026 18:31:24 -0800 Subject: [PATCH 58/78] video: libtfvideo.so exports only tf_video_*, so FFmpeg's x86_64 asm links FFmpeg's x86_64 assembly addresses tables defined in C, such as ff_pw_* and ff_h264_cabac_tables, with 32-bit PC-relative relocations. Linked from the static libav*.a into a shared library, those tables are exported and so preemptible, and the link stopped on R_X86_64_PC32 errors (5,046 with gcc 13.3) although build_ffmpeg.sh passes --enable-pic. On ELF targets the library now links with a version script that exports the four tf_video_* symbols core/tfvideo.zig looks up and makes every other symbol local, so the tables bind inside the library with no text relocations. check-video-exports-linux checks this for aarch64 and x86_64 Linux with a planted table. The loader, build_ffmpeg.sh and its FFmpeg pin, and Mach-O builds are unchanged. --- CHANGELOG.md | 2 + tools/zig/check_tfvideo_exports.py | 70 ++++++++++++++++++++++++++++++ zig/build/video.zig | 17 ++++++++ zig/tests/tfvideo_planted.c | 14 ++++++ 4 files changed, 103 insertions(+) create mode 100644 tools/zig/check_tfvideo_exports.py create mode 100644 zig/tests/tfvideo_planted.c diff --git a/CHANGELOG.md b/CHANGELOG.md index 9381b4e5820..01ca8503596 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -218,6 +218,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Docs: the optimization ledger marks the rows that have open PRs and adds DSpine as row 29 (#2681). Thanks to @jschmied. +- Video on Linux: `libtfvideo.so` links with a version script that exports only its four `tf_video_*` functions, so on x86_64 `zig build video` links `tools/zig/build_ffmpeg.sh`'s FFmpeg with its assembly, where before the link stopped on thousands of `R_X86_64_PC32` relocations against FFmpeg's own tables. Its frames equal PyAV 19.0.1's byte for byte on 23 clips. On aarch64 the library stops exporting FFmpeg's API; real video there is unchecked. macOS is unchanged (#2861). Thanks to @swbiggart. + - GLM-5.3-Flash on Metal with `--vision`: a prompt chunk sends the vision tower only the images that end after its start, so two interleaved image prompts no longer re-encode images their earlier chunks already injected (12 images through the tower instead of 18 in a host test); each image is still encoded on its own and read through its own span, so the rows a chunk reads are unchanged (#2744). - Flash Next on Metal: `TF_FLASHNEXT_SCALE_INDEX=1` (off by default; M1 through M4) has the 4-bit expert kernels read each group's scale and bias through a 16-bit index into its tensor's table of pairs: the same values from 10% fewer weight bytes, so replies keep their tokens. On an M3 Ultra with the MLX 4-bit checkpoint, drafted decode on one stream is 1.3% to 3.9% faster (prose 137.3 to 139.1 tok/s, a copy-heavy edit 262.8 to 273.2) and `"draft": false` 0.8% to 2.2% faster, with every token unchanged. The index sits beside the stored scales and biases, which prompt passes still read: 7.3 GiB more in use and about 4 s more load (#2760, part of #2751). diff --git a/tools/zig/check_tfvideo_exports.py b/tools/zig/check_tfvideo_exports.py new file mode 100644 index 00000000000..c1a9d57db62 --- /dev/null +++ b/tools/zig/check_tfvideo_exports.py @@ -0,0 +1,70 @@ +"""An ELF libtfvideo.so exports exactly the four tf_video_* functions, keeps a planted table local, has no TEXTREL.""" +import struct +import sys +from pathlib import Path + +WANT = {'tf_video_open', 'tf_video_frames', 'tf_video_close', 'tf_video_version'} +SHT_SYMTAB, SHT_DYNAMIC, SHT_DYNSYM = 2, 6, 11 +DT_TEXTREL, DT_FLAGS, DF_TEXTREL = 22, 30, 0x4 + + +def sections(data: bytes) -> list[tuple]: + """(type, offset, size, link, entsize) for each section of a little-endian ELF64 file.""" + if data[:4] != b'\x7fELF' or data[4] != 2 or data[5] != 1: + raise SystemExit('not a little-endian ELF64 file') + shoff, = struct.unpack_from(' dict[str, tuple[int, int]]: + """Defined symbols of the given table: name -> (binding, section index).""" + out = {} + for t, off, size, link, ent in secs: + if t != kind: + continue + strtab = secs[link][1] + for at in range(off + ent, off + size, ent): + name, info, _, shndx = struct.unpack_from('> 4, shndx) + return out + + +def textrel(data: bytes, secs: list[tuple]) -> bool: + """True when .dynamic has DT_TEXTREL or DF_TEXTREL in DT_FLAGS.""" + for t, off, size, _, _ in secs: + if t == SHT_DYNAMIC: + for at in range(off, off + size, 16): + tag, val = struct.unpack_from(' int: + failed = 0 + for arg in sys.argv[1:]: + data = Path(arg).read_bytes() + secs = sections(data) + exported = {n for n, (bind, _) in symbols(data, secs, SHT_DYNSYM).items() if bind in (1, 2)} + planted = symbols(data, secs, SHT_SYMTAB).get('tfvideo_planted_table') + problems = [f'exports {sorted(exported ^ WANT)} beyond or short of the four'] if exported != WANT else [] + if planted is None or planted[0] != 0: + problems.append(f'the planted table is not a local symbol: {planted}') + if textrel(data, secs): + problems.append('text relocations') + machine = {62: 'x86_64', 183: 'aarch64'}.get(struct.unpack_from(' + +const uint16_t tfvideo_planted_table[4] = {1, 2, 3, 4}; + +uint32_t tfvideo_planted_read(void) { + uint32_t v; +#if defined(__x86_64__) + __asm__("movzwl tfvideo_planted_table(%%rip), %0" : "=r"(v)); +#else + v = tfvideo_planted_table[0]; +#endif + return v; +} From 9955a1d0cff6fcf95e09496cee41d77229a3ddeb Mon Sep 17 00:00:00 2001 From: Scott Biggart Date: Wed, 11 Nov 2026 02:03:21 -0800 Subject: [PATCH 59/78] ci: the macOS adapters run check-video-exports-linux beside check-video-layout-linux --- zig/build/host_ci.zig | 1 + 1 file changed, 1 insertion(+) diff --git a/zig/build/host_ci.zig b/zig/build/host_ci.zig index feaa56d56f0..fed6dc007a9 100644 --- a/zig/build/host_ci.zig +++ b/zig/build/host_ci.zig @@ -251,6 +251,7 @@ const macos_adapters = [_][]const u8{ "check-nemotron-prefill-widths-host", "check-qwen35-cache-host", "check-sessions-linux", + "check-video-exports-linux", "check-video-layout-linux", "compile-disk-identity-metal-host", "compile-dsv41-engram-alloc-aarch64", From d815f291c19eb4f54d8b7faba5d3ddf063ce3874 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 13 Nov 2026 13:24:55 +0000 Subject: [PATCH 60/78] changelog: #2861's line from lines/2861.txt in place of its own --- CHANGELOG.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 01ca8503596..31f6983e72c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Video on Linux: `libtfvideo.so` exports only its `tf_video_*` functions, so FFmpeg's x86_64 assembly links into it, and a CI check refuses a build that exports more (#2861). Thanks to @swbiggart. + - Docs: the Intel Arc (XPU) page notes that its B70 prefill row predates 1.1.0's more precise prompt DeltaNet, which costs about 5% of a 32,768-token prompt with tokens unchanged (#2906, measured in #2859). Thanks to @dinadur. - Internal: the CUDA kept prompt-state identity is written once through any hasher, so the coming execution witness can hash the same parts with SHA-256; every kept key stays byte for byte what it was, pinned by tests against 1.2.7's sequence and a fixed key, so saved prompt states still reopen. Nothing a user sees changes (#2903, part of #2741). Thanks to @kotevcode. @@ -218,8 +220,6 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Docs: the optimization ledger marks the rows that have open PRs and adds DSpine as row 29 (#2681). Thanks to @jschmied. -- Video on Linux: `libtfvideo.so` links with a version script that exports only its four `tf_video_*` functions, so on x86_64 `zig build video` links `tools/zig/build_ffmpeg.sh`'s FFmpeg with its assembly, where before the link stopped on thousands of `R_X86_64_PC32` relocations against FFmpeg's own tables. Its frames equal PyAV 19.0.1's byte for byte on 23 clips. On aarch64 the library stops exporting FFmpeg's API; real video there is unchecked. macOS is unchanged (#2861). Thanks to @swbiggart. - - GLM-5.3-Flash on Metal with `--vision`: a prompt chunk sends the vision tower only the images that end after its start, so two interleaved image prompts no longer re-encode images their earlier chunks already injected (12 images through the tower instead of 18 in a host test); each image is still encoded on its own and read through its own span, so the rows a chunk reads are unchanged (#2744). - Flash Next on Metal: `TF_FLASHNEXT_SCALE_INDEX=1` (off by default; M1 through M4) has the 4-bit expert kernels read each group's scale and bias through a 16-bit index into its tensor's table of pairs: the same values from 10% fewer weight bytes, so replies keep their tokens. On an M3 Ultra with the MLX 4-bit checkpoint, drafted decode on one stream is 1.3% to 3.9% faster (prose 137.3 to 139.1 tok/s, a copy-heavy edit 262.8 to 273.2) and `"draft": false` 0.8% to 2.2% faster, with every token unchanged. The index sits beside the stored scales and biases, which prompt passes still read: 7.3 GiB more in use and about 4 s more load (#2760, part of #2751). From 3c86d8da371b1ef94d445cb3257db73e8f98c75f Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Sun, 8 Nov 2026 12:29:10 +0000 Subject: [PATCH 61/78] refactor(core): move the memory counts to engine_status.zig engine_api.zig is at the 600-line limit, and the OCR reader types that follow add to it. Memory and MemorySource describe a backend's status, as HealthSource and Status do, so they move beside them in engine_status.zig and engine_api re-exports them under the same names; no caller changes. --- zig/src/core/engine_api.zig | 10 ++-------- zig/src/core/engine_status.zig | 8 ++++++++ 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index e2ff1527b18..920af8cd006 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -371,18 +371,12 @@ pub const Opened = struct { follow: ?*const fn (*anyopaque) anyerror!void = null, }; -pub const Memory = struct { active: u64 = 0, cache: u64 = 0, peak: u64 = 0, pinned: u64 = 0, limit: ?u64 = null }; - /// A backend's words for its own refusals (a request it cannot serve); null: the error's name. pub const Explain = struct { ctx: ?*anyopaque = null, text: *const fn (ctx: ?*anyopaque, err: anyerror) ?[]const u8 }; -/// A backend's own memory counts for LaneHost; read from HTTP threads, so it must not touch the GPU. -pub const MemorySource = struct { - ctx: ?*anyopaque = null, - read: *const fn (ctx: ?*anyopaque, reset_peak: bool) ?Memory, -}; - pub const HealthSource = @import("engine_status.zig").HealthSource; +pub const Memory = @import("engine_status.zig").Memory; +pub const MemorySource = @import("engine_status.zig").MemorySource; pub const Status = @import("engine_status.zig").Status; pub const SubmitError = error{ Closed, Busy, InvalidSpans, Stale, InvalidRequest, MemoryPressure }; // Stale: put-back has not reached the device diff --git a/zig/src/core/engine_status.zig b/zig/src/core/engine_status.zig index 3a024ac4243..0700fd90a23 100644 --- a/zig/src/core/engine_status.zig +++ b/zig/src/core/engine_status.zig @@ -12,6 +12,14 @@ pub const FsrSource = struct { read: *const fn (*anyopaque) ?Status.FsrCache, }; +pub const Memory = struct { active: u64 = 0, cache: u64 = 0, peak: u64 = 0, pinned: u64 = 0, limit: ?u64 = null }; + +/// A backend's own memory counts for LaneHost; read from HTTP threads, so it must not touch the GPU. +pub const MemorySource = struct { + ctx: ?*anyopaque = null, + read: *const fn (ctx: ?*anyopaque, reset_peak: bool) ?Memory, +}; + pub const Status = struct { healthy: bool = true, /// Engine-thread publication; all retained states, including long-message interior cuts and spare storage. From b6655d40fdf8a950a5f36a49ab97f210289045bd Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Sun, 8 Nov 2026 12:29:11 +0000 Subject: [PATCH 62/78] refactor(server): move the JPEG encoder out of pdf.zig The libjpeg-turbo loader and the BGRA page encoder move to their own module, jpeg_out.zig, with the optional-library loader that PDFium shares, so the OCR route can encode image crops and pdf.zig gains the room its page selection needs later; jpegPixels becomes jpegSize, which returns the width and height a frame header names, for the OCR page planner. Chat PDF input behaves exactly as before, and the new encodeRgb is covered by a round-trip test. --- zig/src/server/jpeg_out.zig | 101 +++++++++++++++++++++++++++++ zig/src/server/pdf.zig | 75 ++++----------------- zig/src/server/pdf_worker.zig | 8 ++- zig/src/server/pdf_worker_test.zig | 4 +- zig/src/server/root.zig | 1 + 5 files changed, 120 insertions(+), 69 deletions(-) create mode 100644 zig/src/server/jpeg_out.zig diff --git a/zig/src/server/jpeg_out.zig b/zig/src/server/jpeg_out.zig new file mode 100644 index 00000000000..0e851bc95b1 --- /dev/null +++ b/zig/src/server/jpeg_out.zig @@ -0,0 +1,101 @@ +//! JPEG output via libjpeg-turbo, loaded on first use: PDF page images (pdf.zig) and OCR image crops (encodeRgb). +const std = @import("std"); +const sys = @import("sys"); +const log = @import("log.zig"); +const Allocator = std.mem.Allocator; + +/// The TurboJPEG 3 entry points the encoders call. +pub const Fns = struct { + tj3Init: *const fn (c_int) callconv(.c) ?*anyopaque, + tj3Destroy: *const fn (?*anyopaque) callconv(.c) void, + tj3Set: *const fn (?*anyopaque, c_int, c_int) callconv(.c) c_int, + tj3Compress8: *const fn (?*anyopaque, [*]const u8, c_int, c_int, c_int, c_int, *?[*]u8, *usize) callconv(.c) c_int, + tj3Free: *const fn (?*anyopaque) callconv(.c) void, +}; + +pub const Error = error{ OutOfMemory, EncodeFailed }; + +const Loaded = union(enum) { unknown, missing, ready: Fns }; + +var state: Loaded = .unknown; +var mutex: std.Io.Mutex = .init; + +/// Looks up every function field of `T` in optional library `id`, optional ones allowed to miss; PDFium uses it too. +pub fn loadDynLib(comptime T: type, comptime id: sys.optional_libraries.Id, a: Allocator) ?T { + const exe = std.process.executableDirPathAlloc(log.io(), a) catch null; + defer if (exe) |p| a.free(p); + var dl = sys.optional_libraries.open(sys.optional_libraries.spec(id), a, sys, exe) orelse return null; + var f: T = undefined; + const info = @typeInfo(T).@"struct"; + inline for (info.field_names, info.field_types) |name, FT| { + if (@typeInfo(FT) == .optional) { + @field(f, name) = dl.lookup(@typeInfo(FT).optional.child, name); + } else { + @field(f, name) = dl.lookup(FT, name) orelse { + dl.close(); + return null; + }; + } + } + return f; +} + +/// libjpeg-turbo's entry points, or null when the library is missing; the first call loads it. +pub fn get(a: Allocator) ?Fns { + const io = log.io(); + mutex.lockUncancelable(io); + defer mutex.unlock(io); + if (state == .unknown) state = if (loadDynLib(Fns, .jpeg, a)) |f| .{ .ready = f } else .missing; + return if (state == .ready) state.ready else null; +} + +/// Encodes PDFium's BGRx bitmap (rows `stride` bytes apart) as a 4:2:0 JPEG at quality 85. +pub fn encodeBgra(a: Allocator, jf: Fns, bgra: []const u8, width: usize, height: usize, stride: usize) Error![]const u8 { + return encode(a, jf, bgra, width, height, stride, 8, 85); // 8: TJPF_BGRA, PDFium's BGRx bitmap +} + +/// Encodes packed RGB rows as a 4:2:0 JPEG at `quality`; NoLibrary when libjpeg-turbo is missing. +pub fn encodeRgb(a: Allocator, rgb: []const u8, width: usize, height: usize, quality: u8) (Error || error{NoLibrary})![]const u8 { + const jf = get(a) orelse return error.NoLibrary; + const pitch = std.math.mul(usize, width, 3) catch return error.EncodeFailed; + if (pitch == 0 or rgb.len / pitch < height) return error.EncodeFailed; + return encode(a, jf, rgb, width, height, pitch, 0, quality); // 0: TJPF_RGB +} + +fn encode(a: Allocator, jf: Fns, pixels: []const u8, width: usize, height: usize, pitch: usize, format: c_int, quality: u8) Error![]const u8 { + const max: usize = std.math.maxInt(c_int); + if (width == 0 or height == 0 or width > max or height > max or pitch > max) return error.EncodeFailed; + const handle = jf.tj3Init(0) orelse return error.EncodeFailed; // 0 = TJINIT_COMPRESS + defer jf.tj3Destroy(handle); + _ = jf.tj3Set(handle, 3, quality); // TJPARAM_QUALITY = 3 + _ = jf.tj3Set(handle, 4, 2); // TJPARAM_SUBSAMP = 4, TJSAMP_420 = 2 + var out_buf: ?[*]u8 = null; + var out_size: usize = 0; + const rc = jf.tj3Compress8(handle, pixels.ptr, @intCast(width), @intCast(pitch), @intCast(height), format, &out_buf, &out_size); + defer if (out_buf) |p| jf.tj3Free(p); + if (rc != 0 or out_buf == null or out_size == 0) return error.EncodeFailed; + return a.dupe(u8, out_buf.?[0..out_size]); +} + +test "jpeg out: encodeRgb round-trips a small image through core.picture" { + const picture = @import("core").picture; + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const w = 24; + const h = 10; + var rgb: [w * h * 3]u8 = undefined; + for (&rgb, 0..) |*px, i| px.* = @truncate(i * 7); + const bytes = encodeRgb(a, &rgb, w, h, 90) catch |err| switch (err) { + error.NoLibrary => return error.SkipZigTest, + else => return err, + }; + try std.testing.expectEqual(picture.Format.jpeg, picture.sniff(bytes)); + const img = picture.decode(std.testing.io, a, bytes) catch |err| switch (err) { + error.NoJpegLibrary, error.JpegLibraryVersion => return error.SkipZigTest, + else => return err, + }; + try std.testing.expectEqual(@as(usize, w), img.width); + try std.testing.expectEqual(@as(usize, h), img.height); + try std.testing.expectError(error.EncodeFailed, encodeRgb(a, rgb[0 .. rgb.len - 1], w, h, 90)); +} diff --git a/zig/src/server/pdf.zig b/zig/src/server/pdf.zig index a134564a3ee..837bbc02661 100644 --- a/zig/src/server/pdf.zig +++ b/zig/src/server/pdf.zig @@ -4,6 +4,7 @@ const sys = @import("sys"); const json = @import("json"); const errors = @import("errors.zig"); const images = @import("images.zig"); +const jpeg_out = @import("jpeg_out.zig"); const log = @import("log.zig"); const worker = @import("pdf_worker.zig"); const Value = json.Value; @@ -40,65 +41,23 @@ pub const PdfiumFns = struct { FPDF_RenderPage_Close: ?*const fn (?*anyopaque) callconv(.c) void = null, }; -pub const JpegFns = struct { - tj3Init: *const fn (c_int) callconv(.c) ?*anyopaque, - tj3Destroy: *const fn (?*anyopaque) callconv(.c) void, - tj3Set: *const fn (?*anyopaque, c_int, c_int) callconv(.c) c_int, - tj3Compress8: *const fn (?*anyopaque, [*]const u8, c_int, c_int, c_int, c_int, *?[*]u8, *usize) callconv(.c) c_int, - tj3Free: *const fn (?*anyopaque) callconv(.c) void, -}; - const DlInfo = extern struct { fname: ?[*:0]const u8, fbase: ?*anyopaque, sname: ?[*:0]const u8, saddr: ?*anyopaque }; extern "c" fn dladdr(addr: *const anyopaque, info: *DlInfo) c_int; const PdfiumLoaded = union(enum) { unknown, missing, ready: PdfiumFns }; -const JpegLoaded = union(enum) { unknown, missing, ready: JpegFns }; var pdfium_state: PdfiumLoaded = .unknown; -var jpeg_state: JpegLoaded = .unknown; var pdfium_mutex: std.Io.Mutex = .init; -var jpeg_mutex: std.Io.Mutex = .init; - -pub fn loadDynLib(comptime T: type, comptime id: sys.optional_libraries.Id, a: Allocator) ?T { - const exe = std.process.executableDirPathAlloc(log.io(), a) catch null; - defer if (exe) |p| a.free(p); - var dl = sys.optional_libraries.open(sys.optional_libraries.spec(id), a, sys, exe) orelse return null; - var f: T = undefined; - const info = @typeInfo(T).@"struct"; - inline for (info.field_names, info.field_types) |name, FT| { - if (@typeInfo(FT) == .optional) { - @field(f, name) = dl.lookup(@typeInfo(FT).optional.child, name); - } else { - @field(f, name) = dl.lookup(FT, name) orelse { - dl.close(); - return null; - }; - } - } - return f; -} fn findPdfium(a: Allocator) PdfiumLoaded { - const f = loadDynLib(PdfiumFns, .pdfium, a) orelse return .missing; + const f = jpeg_out.loadDynLib(PdfiumFns, .pdfium, a) orelse return .missing; if (sys.windows) f.FPDF_InitLibrary(); return .{ .ready = f }; } -fn findJpeg(a: Allocator) JpegLoaded { - return if (loadDynLib(JpegFns, .jpeg, a)) |f| .{ .ready = f } else .missing; -} - -fn getJpeg(a: Allocator) ?JpegFns { - const io = log.io(); - jpeg_mutex.lockUncancelable(io); - defer jpeg_mutex.unlock(io); - if (jpeg_state == .unknown) jpeg_state = findJpeg(a); - return if (jpeg_state == .ready) jpeg_state.ready else null; -} - /// Logs the loaded PDFium path and SHA-256 at startup when --vision-pdf is enabled, or warns if a library is missing. pub fn logStartup(a: Allocator, io: std.Io) void { - if (getJpeg(a) == null) log.line("--vision-pdf: {s}", .{sys.optional_libraries.spec(.jpeg).message}); + if (jpeg_out.get(a) == null) log.line("--vision-pdf: {s}", .{sys.optional_libraries.spec(.jpeg).message}); pdfium_mutex.lockUncancelable(io); defer pdfium_mutex.unlock(io); if (pdfium_state == .unknown) pdfium_state = findPdfium(a); @@ -207,7 +166,7 @@ pub fn decodeFileData(cx: *Cx, raw: []const u8) errors.Refused![]const u8 { /// Renders all pages of `pdf_bytes` within the request-wide `budget` into `image_url` parts appended to `parts`. pub fn expandBytes(cx: *Cx, pdf_bytes: []const u8, image_type: []const u8, max_images: u32, max_pages: u32, budget: *Budget, parts: *std.ArrayList(Value)) errors.Refused!void { - const jf = getJpeg(cx.a) orelse return cx.fail(.unimplemented, "{s}", .{sys.optional_libraries.spec(.jpeg).message}); + const jf = jpeg_out.get(cx.a) orelse return cx.fail(.unimplemented, "{s}", .{sys.optional_libraries.spec(.jpeg).message}); const io = log.io(); pdfium_mutex.lockUncancelable(io); defer pdfium_mutex.unlock(io); @@ -255,7 +214,7 @@ fn renderPageWithDeadline(cx: *Cx, f: PdfiumFns, bmp: ?*anyopaque, page: ?*anyop if (std.Io.Clock.awake.now(log.io()).toNanoseconds() >= deadline_ns) return cx.refuse("PDF rendering timed out"); } -pub fn renderInProcess(cx: *Cx, f: PdfiumFns, jf: JpegFns, pdf_bytes: []const u8, max_images: u32, max_pages: u32, deadline_ns: i96, budget: *Budget, out_jpegs: *std.ArrayList([]const u8)) errors.Refused!void { +pub fn renderInProcess(cx: *Cx, f: PdfiumFns, jf: jpeg_out.Fns, pdf_bytes: []const u8, max_images: u32, max_pages: u32, deadline_ns: i96, budget: *Budget, out_jpegs: *std.ArrayList([]const u8)) errors.Refused!void { if (pdf_bytes.len > std.math.maxInt(c_int)) return cx.refuse("PDF exceeds the encoded byte limit"); if (std.Io.Clock.awake.now(log.io()).toNanoseconds() >= deadline_ns) return cx.refuse("PDF rendering timed out"); const doc = f.FPDF_LoadMemDocument(pdf_bytes.ptr, @intCast(pdf_bytes.len), null) orelse @@ -287,7 +246,7 @@ pub fn renderInProcess(cx: *Cx, f: PdfiumFns, jf: JpegFns, pdf_bytes: []const u8 const total_bytes = @as(usize, @intCast(stride)) * @as(usize, @intCast(dims[1])); @memset(bgra[0..total_bytes], 0xFF); try renderPageWithDeadline(cx, f, bmp, page, dims[0], dims[1], idx, deadline_ns); - const jpeg_bytes = encodeJpegBgra(cx.a, jf, bgra[0..total_bytes], @intCast(dims[0]), @intCast(dims[1]), @intCast(stride)) catch |err| switch (err) { + const jpeg_bytes = jpeg_out.encodeBgra(cx.a, jf, bgra[0..total_bytes], @intCast(dims[0]), @intCast(dims[1]), @intCast(stride)) catch |err| switch (err) { error.OutOfMemory => return error.OutOfMemory, error.EncodeFailed => return cx.refuse("rendered PDF page could not be encoded as JPEG"), }; @@ -315,7 +274,8 @@ fn pageDimensions(pw_raw: f32, ph_raw: f32) [2]c_int { return .{ w, h }; } -pub fn jpegPixels(bytes: []const u8) ?u64 { +/// The width and height a baseline or progressive JPEG's frame header names, or null when it has none. +pub fn jpegSize(bytes: []const u8) ?[2]u32 { if (bytes.len < 4 or bytes[0] != 0xFF or bytes[1] != 0xD8) return null; var i: usize = 2; while (i + 4 <= bytes.len) { @@ -330,28 +290,15 @@ pub fn jpegPixels(bytes: []const u8) ?u64 { if (seg_len < 2 or i + seg_len > bytes.len) return null; if (marker == 0xC0 or marker == 0xC1 or marker == 0xC2) { if (seg_len < 7) return null; - const h: u64 = std.mem.readInt(u16, bytes[i + 3 ..][0..2], .big); - const w: u64 = std.mem.readInt(u16, bytes[i + 5 ..][0..2], .big); - return if (w > 0 and h > 0) w * h else null; + const h: u32 = std.mem.readInt(u16, bytes[i + 3 ..][0..2], .big); + const w: u32 = std.mem.readInt(u16, bytes[i + 5 ..][0..2], .big); + return if (w > 0 and h > 0) .{ w, h } else null; } i += seg_len; } return null; } -fn encodeJpegBgra(a: Allocator, jf: JpegFns, bgra: []const u8, width: usize, height: usize, stride: usize) error{ OutOfMemory, EncodeFailed }![]const u8 { - const handle = jf.tj3Init(0) orelse return error.EncodeFailed; // 0 = TJINIT_COMPRESS - defer jf.tj3Destroy(handle); - _ = jf.tj3Set(handle, 3, 85); // TJPARAM_QUALITY = 3 - _ = jf.tj3Set(handle, 4, 2); // TJPARAM_SUBSAMP = 4, TJSAMP_420 = 2 - var out_buf: ?[*]u8 = null; - var out_size: usize = 0; - const rc = jf.tj3Compress8(handle, bgra.ptr, @intCast(width), @intCast(stride), @intCast(height), 8, &out_buf, &out_size); // 8: TJPF_BGRA, PDFium's BGRx bitmap - defer if (out_buf) |p| jf.tj3Free(p); - if (rc != 0 or out_buf == null or out_size == 0) return error.EncodeFailed; - return a.dupe(u8, out_buf.?[0..out_size]); -} - pub fn makeDataUrl(a: Allocator, mime: []const u8, bytes: []const u8) Allocator.Error![]const u8 { const prefix = try std.fmt.allocPrint(a, "data:{s};base64,", .{mime}); const enc = std.base64.standard.Encoder; diff --git a/zig/src/server/pdf_worker.zig b/zig/src/server/pdf_worker.zig index 248cd387a9a..099b751a447 100644 --- a/zig/src/server/pdf_worker.zig +++ b/zig/src/server/pdf_worker.zig @@ -6,6 +6,7 @@ const errors = @import("errors.zig"); const images = @import("images.zig"); const log = @import("log.zig"); const pdf = @import("pdf.zig"); +const jpeg_out = @import("jpeg_out.zig"); const Cx = errors.Cx; const Budget = pdf.Budget; const Allocator = std.mem.Allocator; @@ -117,7 +118,8 @@ pub fn parseChildReply(cx: *Cx, frame: []const u8, max_images: u32, max_pages: u if (len > images.max_image_bytes) return cx.refuse("rendered PDF page exceeds the 10 MB image byte limit"); if (budget.encoded_bytes + len > images.max_request_bytes) return cx.refuse("the request's images exceed 20 MB encoded"); const jb = frame[pos .. pos + len]; - const px = pdf.jpegPixels(jb) orelse return cx.refuse("rendered PDF page could not be encoded as JPEG"); + const size = pdf.jpegSize(jb) orelse return cx.refuse("rendered PDF page could not be encoded as JPEG"); + const px = @as(u64, size[0]) * size[1]; if (budget.pixels_used + px > pdf.max_request_pixels) return cx.refuse("the request's images exceed the 32-megapixel limit"); budget.pdf_pages_used += 1; budget.images_used += 1; @@ -173,8 +175,8 @@ fn work(cx: *Cx, jpegs: *std.ArrayList([]const u8)) errors.Refused!void { if (!readAll(0, pdf_bytes)) return cx.refuse(failed); applyChildLimits(); const libs = sys.optional_libraries; - const f = pdf.loadDynLib(pdf.PdfiumFns, .pdfium, cx.a) orelse return cx.fail(.unimplemented, "{s}", .{libs.spec(.pdfium).message}); - const jf = pdf.loadDynLib(pdf.JpegFns, .jpeg, cx.a) orelse return cx.fail(.unimplemented, "{s}", .{libs.spec(.jpeg).message}); + const f = jpeg_out.loadDynLib(pdf.PdfiumFns, .pdfium, cx.a) orelse return cx.fail(.unimplemented, "{s}", .{libs.spec(.pdfium).message}); + const jf = jpeg_out.loadDynLib(jpeg_out.Fns, .jpeg, cx.a) orelse return cx.fail(.unimplemented, "{s}", .{libs.spec(.jpeg).message}); f.FPDF_InitLibrary(); var budget: Budget = .{ .images_used = req.images_used, .pdf_pages_used = req.pdf_pages_used, .encoded_bytes = @intCast(req.encoded_bytes), .pixels_used = req.pixels_used }; return pdf.renderInProcess(cx, f, jf, pdf_bytes, req.max_images, req.max_pages, deadline_ns, &budget, jpegs); diff --git a/zig/src/server/pdf_worker_test.zig b/zig/src/server/pdf_worker_test.zig index 28a56dd3722..446f093a483 100644 --- a/zig/src/server/pdf_worker_test.zig +++ b/zig/src/server/pdf_worker_test.zig @@ -37,7 +37,7 @@ test "tensorfold-native PDF render worker: renders a text PDF in its spawned wor }; try std.testing.expectEqual(@as(usize, 1), jpegs.items.len); try std.testing.expectEqual(pdf.Budget{ .images_used = 1, .pdf_pages_used = 1, .encoded_bytes = jpegs.items[0].len, .pixels_used = 400 * 400 }, budget); - try std.testing.expectEqual(@as(?u64, 400 * 400), pdf.jpegPixels(jpegs.items[0])); + try std.testing.expectEqual([2]u32{ 400, 400 }, pdf.jpegSize(jpegs.items[0]).?); } test "tensorfold-native PDF render worker: libpdfium renders in the spawned worker and enforces request-wide page, byte and pixel caps" { @@ -59,7 +59,7 @@ test "tensorfold-native PDF render worker: libpdfium renders in the spawned work for (parts.items) |p| { const img_bytes = try images.dataBytes(&cx, p.get("image_url").?.get("url").?.string); try std.testing.expectEqual(images.Format.jpeg, images.sniff(img_bytes)); - try std.testing.expectEqual(@as(?u64, 400 * 400), pdf.jpegPixels(img_bytes)); + try std.testing.expectEqual([2]u32{ 400, 400 }, pdf.jpegSize(img_bytes).?); } var second_cx: Cx = .{ .a = a }; try std.testing.expectError(error.Refused, pdf.expandBytes(&second_cx, pdf.minimal_two_page_pdf, "image_url", 4, 3, &budget, &parts)); diff --git a/zig/src/server/root.zig b/zig/src/server/root.zig index 66d0da36073..d3e4954a94f 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -112,6 +112,7 @@ test { _ = @import("messages.zig"); _ = @import("pdf.zig"); _ = @import("pdf_flags_test.zig"); + _ = @import("jpeg_out.zig"); _ = @import("images.zig"); _ = @import("images_video_marker_test.zig"); _ = @import("image_key.zig"); From 066babe699f9bbb2c83f24339d7e5c4796ad987f Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Sun, 8 Nov 2026 12:29:11 +0000 Subject: [PATCH 63/78] feat(core): add OCR reader types and the image frame size An engine now declares an optional OCR reader in Info.ocr: the prompt that follows the page image, the geometry it gives and a parse from generated text to Markdown blocks with quads in original pixels. PreparedImage gains the tower-input frame so readers can map boxes back, the HIP offer carries the reader through, and a bounded reader for the HTML table subset dots.mocr writes turns tables into GFM pipes, clean spanned HTML or text rows, escaping literal markup so Markdown never reinterprets recognized text. --- zig/src/core/engine_api.zig | 3 + zig/src/core/engine_api_test.zig | 1 + zig/src/core/engine_ocr.zig | 54 ++++ zig/src/core/engine_vision.zig | 5 + zig/src/core/ocr_html.zig | 437 +++++++++++++++++++++++++++++++ zig/src/core/ocr_html_test.zig | 269 +++++++++++++++++++ zig/src/hip_serve/offer.zig | 3 + zig/src/native/hip.zig | 1 + 8 files changed, 773 insertions(+) create mode 100644 zig/src/core/engine_ocr.zig create mode 100644 zig/src/core/ocr_html.zig create mode 100644 zig/src/core/ocr_html_test.zig diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index 920af8cd006..7bb5f6fafa4 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -116,6 +116,8 @@ pub const FrameSink = vision_types.FrameSink; pub const PreparedVideo = vision_types.PreparedVideo; pub const VideoOffer = vision_types.VideoOffer; pub const VisionOpen = vision_types.VisionOpen; +pub const Size = vision_types.Size; +pub const ocr = @import("engine_ocr.zig"); pub const SpeechRequest = @import("speech_api.zig").SpeechRequest; pub const SpeechResult = @import("speech_api.zig").SpeechResult; @@ -223,6 +225,7 @@ pub const Info = struct { pictures: bool = false, /// Image input (--vision); null: text only. vision: ?Vision = null, + ocr: ?ocr.Ocr = null, // OCR reader for /v1/ocr and /ocr; null: no OCR routes /// The engine decodes greedily only: the server refuses a request with a temperature before admitting it. greedy_only: bool = false, /// The learner reads conversation prompts at any length through taps. diff --git a/zig/src/core/engine_api_test.zig b/zig/src/core/engine_api_test.zig index 74a4aadc51d..efadf7dc452 100644 --- a/zig/src/core/engine_api_test.zig +++ b/zig/src/core/engine_api_test.zig @@ -15,4 +15,5 @@ test { _ = @import("keepalive.zig"); _ = @import("originals_test.zig"); _ = @import("image_api_test.zig"); + _ = @import("engine_ocr.zig"); } diff --git a/zig/src/core/engine_ocr.zig b/zig/src/core/engine_ocr.zig new file mode 100644 index 00000000000..07d1c507a57 --- /dev/null +++ b/zig/src/core/engine_ocr.zig @@ -0,0 +1,54 @@ +//! An engine's OCR reader (``Info.ocr``): its prompt and the parse from generated text to Markdown blocks. +const std = @import("std"); +const Allocator = std.mem.Allocator; +const Size = @import("engine_vision.zig").Size; + +/// The HTML table subset readers and the server share (ocr_html.zig). +pub const html = @import("ocr_html.zig"); + +/// TL, TR, BR, BL of the text in its upright reading frame, in original-image pixels (rotated text keeps its order). +pub const Quad = [4][2]f32; + +/// What a block holds; ``page`` only for a ``.none`` reader (the whole page's Markdown in one block). +pub const Kind = enum { + text, + title, + section_header, + list, + table, + formula, + code, + caption, + footnote, + page_header, + page_footer, + picture, + page, +}; + +/// The geometry a reader gives: none, one quad a block, or one quad a line. +pub const Geometry = enum { none, block, line }; + +/// One block of a page in reading order. +pub const OcrBlock = struct { + content: []const u8, // normalized Markdown (docs/api.md OCR: Reader Markdown); empty only for .picture + quad: ?Quad, // null only for a reader whose geometry is .none + kind: Kind, +}; + +/// The generated text does not follow the model's grammar. +pub const Error = error{Unreadable}; + +/// OCR an engine serves: the prompt after the page image, and the reader of the generated text. +pub const Ocr = struct { + /// Text parts that follow the image in the single user turn (may be empty). + prompt: []const u8, + /// What geometry `parse` gives. + geometry: Geometry, + /// Text to blocks in pixels of ``original``; ``frame`` is the tower input size the prepare step used. + parse: *const fn (a: Allocator, text: []const u8, original: Size, frame: Size) anyerror![]OcrBlock, +}; + +test { + _ = @import("ocr_html_test.zig"); +} diff --git a/zig/src/core/engine_vision.zig b/zig/src/core/engine_vision.zig index a4d4e819da5..285b697bbe0 100644 --- a/zig/src/core/engine_vision.zig +++ b/zig/src/core/engine_vision.zig @@ -13,8 +13,13 @@ pub const PreparedImage = struct { hash: u64, /// The image as tower inputs in order, each image tokens between `before` and `after`; empty: `pixels` alone. pieces: []const Piece = &.{}, + /// Pixel size of the tower input the image was resized to; set by every family with an OCR reader. + frame: ?Size = null, }; +/// An image's pixel size. +pub const Size = struct { width: u32, height: u32 }; + /// One tower input of an image, and the tokens the placeholder expands to around its rows. pub const Piece = struct { pixels: []f32, gh: u32, gw: u32, tokens: u32, before: []const u32 = &.{}, after: []const u32 = &.{} }; diff --git a/zig/src/core/ocr_html.zig b/zig/src/core/ocr_html.zig new file mode 100644 index 00000000000..83bacb8397d --- /dev/null +++ b/zig/src/core/ocr_html.zig @@ -0,0 +1,437 @@ +//! The HTML table subset OCR models write: cells and spans, and the table as GFM pipes, clean HTML or text rows. +const std = @import("std"); +const Allocator = std.mem.Allocator; + +/// Most HTML bytes, grid rows and grid columns ``parse`` reads; a table past them is ``error.Malformed``. +pub const max_bytes = 64 * 1024; +pub const max_rows = 512; +pub const max_cols = 64; + +/// One cell at its top-left grid position, covering ``rows`` x ``cols``; ``text`` holds '\n' where the HTML had
. +pub const Cell = struct { text: []const u8, row: u32, col: u32, rows: u32, cols: u32, head: bool }; + +/// A table's grid and its cells in row-major order; ``spanned``: some cell covers more than one position. +pub const Table = struct { + rows: u32, + cols: u32, + cells: []Cell, + spanned: bool, + + /// Frees what ``parse`` allocated in ``a``. + pub fn deinit(t: Table, a: Allocator) void { + for (t.cells) |c| a.free(c.text); + a.free(t.cells); + } +}; + +pub const ParseError = error{ Malformed, OutOfMemory }; + +/// The one in ``html``, with only whitespace and comments around it; cells and texts allocated in ``a``. +pub fn parse(a: Allocator, html: []const u8) ParseError!Table { + if (html.len > max_bytes) return error.Malformed; + var p: Parser = .{ .a = a, .src = html }; + defer p.text.deinit(a); + errdefer { + for (p.cells.items) |c| a.free(c.text); + p.cells.deinit(a); + } + try p.run(); + var cols: u32 = 0; + var spanned = false; + for (p.cells.items) |*c| { + c.rows = @min(c.rows, p.rows - c.row); // a rowspan past the last row ends there, as in HTML + cols = @max(cols, c.col + c.cols); + spanned = spanned or c.rows > 1 or c.cols > 1; + } + if (cols == 0) return error.Malformed; + return .{ .rows = p.rows, .cols = cols, .cells = try p.cells.toOwnedSlice(a), .spanned = spanned }; +} + +/// The table as a GFM pipe table: first row the header, cell text escaped to read literally (``cellText``); spans lost. +pub fn pipe(a: Allocator, t: Table) Allocator.Error![]const u8 { + var out: std.ArrayList(u8) = .empty; + errdefer out.deinit(a); + var k: usize = 0; + for (0..t.rows) |r| { + for (0..t.cols) |c| { + try out.appendSlice(a, if (c == 0) "| " else " | "); + try cellText(a, &out, if (topLeft(t, &k, r, c)) |cell| cell.text else ""); + } + try out.appendSlice(a, " |"); + if (r == 0) { + try out.append(a, '\n'); + for (0..t.cols) |_| try out.appendSlice(a, "| --- "); + try out.append(a, '|'); + } + if (r + 1 < t.rows) try out.append(a, '\n'); + } + return out.toOwnedSlice(a); +} + +/// The table as one clean
: rows of
/, rowspan and colspan only above 1, text entity-escaped. +pub fn normalized(a: Allocator, t: Table) Allocator.Error![]const u8 { + var out: std.ArrayList(u8) = .empty; + errdefer out.deinit(a); + try out.appendSlice(a, ""); + var k: usize = 0; + for (0..t.rows) |r| { + try out.appendSlice(a, ""); + while (k < t.cells.len and t.cells[k].row == r) : (k += 1) { + const cell = t.cells[k]; + const name = if (cell.head) "th" else "td"; + try out.print(a, "<{s}", .{name}); + if (cell.rows > 1) try out.print(a, " rowspan=\"{d}\"", .{cell.rows}); + if (cell.cols > 1) try out.print(a, " colspan=\"{d}\"", .{cell.cols}); + try out.append(a, '>'); + for (cell.text) |ch| switch (ch) { + '&' => try out.appendSlice(a, "&"), + '<' => try out.appendSlice(a, "<"), + '>' => try out.appendSlice(a, ">"), + '"' => try out.appendSlice(a, """), + '\n' => try out.appendSlice(a, "
"), + else => try out.append(a, ch), + }; + try out.print(a, "", .{name}); + } + try out.appendSlice(a, ""); + } + try out.appendSlice(a, "
"); + return out.toOwnedSlice(a); +} + +/// One line a grid row, its cells joined by '\t'; a spanned cell's text at its top-left only, trailing cells kept. +pub fn rowsText(a: Allocator, t: Table) Allocator.Error![][]const u8 { + const lines = try a.alloc([]const u8, t.rows); + var done: usize = 0; + errdefer { + for (lines[0..done]) |l| a.free(l); + a.free(lines); + } + var line: std.ArrayList(u8) = .empty; + defer line.deinit(a); + var k: usize = 0; + for (lines, 0..) |*l, r| { + line.clearRetainingCapacity(); + for (0..t.cols) |c| { + if (c > 0) try line.append(a, '\t'); + const text = if (topLeft(t, &k, r, c)) |cell| cell.text else ""; + for (text) |ch| try line.append(a, if (ch == '\n') ' ' else ch); + } + l.* = try a.dupe(u8, line.items); + done += 1; + } + return lines; +} + +/// A cell's text in a pipe row: math as written (only `|` escaped), markup escaped elsewhere, to read back literally. +fn cellText(a: Allocator, out: *std.ArrayList(u8), text: []const u8) Allocator.Error!void { + var no_close = false; // a `$` found no closer, so no later `$` in the cell can (markdown.zig's rule) + var i: usize = 0; + while (i < text.len) { + const ch = text[i]; + if (ch == '$') { + if (mathEnd(text, i, &no_close)) |end| { + try mathText(a, out, text[i..end]); + i = end; + continue; + } + // a `$` that opens no math is escaped, so a reader never pairs it with another + const lone: usize = if (i + 1 < text.len and text[i + 1] == '$') 2 else 1; + for (0..lone) |_| try out.appendSlice(a, "\\$"); + i += lone; + continue; + } + switch (ch) { + '\n' => try out.append(a, ' '), + '|', '`', '*', '[', ']' => try out.appendSlice(a, &[_]u8{ '\\', ch }), + // CommonMark reads a backslash as an escape only before ASCII punctuation; elsewhere it is literal + '\\' => try out.appendSlice(a, if (i + 1 < text.len and isPunct(text[i + 1])) "\\\\" else "\\"), + // `_` between two word characters can neither open nor close emphasis, so snake_case stays readable + '_' => try out.appendSlice(a, if (i > 0 and i + 1 < text.len and isWord(text[i - 1]) and isWord(text[i + 1])) "_" else "\\_"), + '<' => try out.appendSlice(a, if (opensTag(text[i + 1 ..])) "<" else "<"), + '&' => try out.appendSlice(a, if (opensEntity(text[i + 1 ..])) "&" else "&"), + else => try out.append(a, ch), + } + i += 1; + } +} + +/// Math as written, each `|` as `\|`: GFM's table step removes exactly that backslash, whatever precedes it. +fn mathText(a: Allocator, out: *std.ArrayList(u8), m: []const u8) Allocator.Error!void { + for (m) |ch| switch (ch) { + '|' => try out.appendSlice(a, "\\|"), + '\n' => try out.append(a, ' '), + else => try out.append(a, ch), + }; +} + +/// The end of the `$$…$$` or `$…$` math opening at ``s[i]``, by markdown.zig's ``mathSpan`` rule; null when none. +fn mathEnd(s: []const u8, i: usize, no_close: *bool) ?usize { + if (i + 1 < s.len and s[i + 1] == '$') return if (std.mem.indexOfPos(u8, s, i + 2, "$$")) |j| j + 2 else null; + if (i + 1 >= s.len or isSpace(s[i + 1]) or no_close.*) return null; + var j = i + 2; + while (j < s.len) : (j += 1) { + if (s[j] != '$' or isSpace(s[j - 1]) or escaped(s, j)) continue; + if (j + 1 < s.len and std.ascii.isDigit(s[j + 1])) continue; + return j + 1; + } + no_close.* = true; + return null; +} + +/// Whether ``s[i]`` follows an odd run of backslashes. +fn escaped(s: []const u8, i: usize) bool { + var n: usize = 0; + while (n < i and s[i - 1 - n] == '\\') n += 1; + return n % 2 == 1; +} + +fn isSpace(c: u8) bool { + return c == ' ' or c == '\t' or c == '\n' or c == '\r'; +} + +fn isPunct(c: u8) bool { + return c < 0x80 and std.ascii.isPrint(c) and !std.ascii.isAlphanumeric(c) and c != ' '; +} + +/// A byte of a word for emphasis rules: neither whitespace nor ASCII punctuation (UTF-8 letters count). +fn isWord(c: u8) bool { + return !isSpace(c) and !isPunct(c); +} + +/// Whether text after a literal `<` would start raw HTML in Markdown: a tag, an end tag, a comment or a declaration. +fn opensTag(rest: []const u8) bool { + return rest.len > 0 and (std.ascii.isAlphabetic(rest[0]) or std.mem.indexOfScalar(u8, "/!?", rest[0]) != null); +} + +/// Whether text after a literal `&` reads as an entity reference in Markdown (`name;`, `#digits;` or `#xhex;`). +fn opensEntity(rest: []const u8) bool { + var n: usize = 0; + while (n < rest.len and n < 33 and (std.ascii.isAlphanumeric(rest[n]) or (n == 0 and rest[n] == '#'))) n += 1; + return n > 0 and n < rest.len and rest[n] == ';'; +} + +/// The cell whose top-left is (r, c), walking ``k`` through the row-major cells; null: covered or empty. +fn topLeft(t: Table, k: *usize, r: usize, c: usize) ?Cell { + if (k.* >= t.cells.len or t.cells[k.*].row != r or t.cells[k.*].col != c) return null; + k.* += 1; + return t.cells[k.* - 1]; +} + +const Tag = enum { table, tr, td, th, br, other }; + +fn tagKind(name: []const u8) Tag { + inline for (.{ Tag.table, Tag.tr, Tag.td, Tag.th, Tag.br }) |t| { + if (std.ascii.eqlIgnoreCase(name, @tagName(t))) return t; + } + return .other; +} + +/// A span attribute's value as HTML reads it: leading digits, and 1 when there are none or they read 0. +fn spanValue(v: []const u8) u32 { + var i: usize = 0; + while (i < v.len and std.ascii.isWhitespace(v[i])) i += 1; + var n: u32 = 0; + while (i < v.len and std.ascii.isDigit(v[i])) : (i += 1) n = @min(n * 10 + (v[i] - '0'), 1_000_000); + return @max(n, 1); +} + +/// The code point a named or numeric entity body (between '&' and ';') stands for; null: not one this reader knows. +fn entityPoint(name: []const u8) ?u21 { + const named = .{ .{ "amp", '&' }, .{ "lt", '<' }, .{ "gt", '>' }, .{ "quot", '"' }, .{ "apos", '\'' } }; + if (std.mem.eql(u8, name, "nbsp")) return 0xa0; + inline for (named) |e| if (std.mem.eql(u8, name, e[0])) return e[1]; + if (name.len < 2 or name[0] != '#') return null; + const hex = name[1] == 'x' or name[1] == 'X'; + const base: u8 = if (hex) 16 else 10; + const digits = name[if (hex) 2 else 1..]; + if (digits.len == 0 or digits.len > 7) return null; + var v: u32 = 0; + for (digits) |d| v = v * base + (std.fmt.charToDigit(d, base) catch return null); + return if (v == 0 or v > 0x10ffff or (v >= 0xd800 and v <= 0xdfff)) 0xfffd else @intCast(v); +} + +/// Whitespace or one of ``ends``: where an attribute's name or unquoted value stops. +fn stops(ch: u8, ends: []const u8) bool { + return std.ascii.isWhitespace(ch) or std.mem.indexOfScalar(u8, ends, ch) != null; +} + +const Parser = struct { + a: Allocator, + src: []const u8, + i: usize = 0, + where: enum { before, inside, after } = .before, + cells: std.ArrayList(Cell) = .empty, + cell: ?Cell = null, // the open cell; its text gathers in ``text`` until it closes + text: std.ArrayList(u8) = .empty, + space: bool = false, // whitespace met in the open cell since its last character + row_open: bool = false, + rows: u32 = 0, // rows opened so far; the open row is the last + col: u32 = 0, // the open row's next column to try + until: [max_cols]u32 = @splat(0), // each column is covered by earlier cells through row until - 1 + + fn run(p: *Parser) ParseError!void { + while (p.i < p.src.len) { + const ch = p.src[p.i]; + const next = if (p.i + 1 < p.src.len) p.src[p.i + 1] else 0; + const after = if (p.i + 2 < p.src.len) p.src[p.i + 2] else 0; + if (ch == '<' and (next == '!' or next == '?')) { + try p.skipMarkup(); + } else if (ch == '<' and std.ascii.isAlphabetic(if (next == '/') after else next)) { + try p.tag(); + } else if (ch == '&') { + try p.entity(); + } else { + try p.char(ch); + p.i += 1; + } + } + if (p.where != .after) return error.Malformed; + } + + /// A comment, a doctype or a processing instruction: skipped whole. + fn skipMarkup(p: *Parser) ParseError!void { + if (std.mem.startsWith(u8, p.src[p.i..], "") orelse return error.Malformed; + p.i = end + 3; + } else { + const end = std.mem.indexOfScalarPos(u8, p.src, p.i, '>') orelse return error.Malformed; + p.i = end + 1; + } + } + + fn entity(p: *Parser) ParseError!void { + const rest = p.src[p.i..]; + const semi = std.mem.indexOfScalar(u8, rest[0..@min(rest.len, 12)], ';'); + const point = if (semi) |s| entityPoint(rest[1..s]) else null; + const cp = point orelse { + p.i += 1; + return p.char('&'); + }; + p.i += semi.? + 1; + if (cp == 0xa0) return p.char(' '); + var buf: [4]u8 = undefined; + const n = std.unicode.utf8Encode(cp, &buf) catch unreachable; + for (buf[0..n]) |b| try p.char(b); + } + + /// One text byte: whitespace runs fold to one space between words; outside a cell only whitespace may stand. + fn char(p: *Parser, ch: u8) ParseError!void { + const blank = std.ascii.isWhitespace(ch); + if (p.cell == null) { + if (!blank) return error.Malformed; + return; + } + if (blank) { + p.space = true; + return; + } + if (p.space and p.text.items.len > 0 and p.text.getLast() != '\n') try p.text.append(p.a, ' '); + p.space = false; + try p.text.append(p.a, ch); + } + + /// A start or end tag at ``i``: its name, its span attributes in any quoting, then its effect. + fn tag(p: *Parser) ParseError!void { + const src = p.src; + var i = p.i + 1; + const close = src[i] == '/'; + if (close) i += 1; + const start = i; + while (i < src.len and (std.ascii.isAlphanumeric(src[i]) or std.mem.indexOfScalar(u8, "-_:.", src[i]) != null)) i += 1; + const kind = tagKind(src[start..i]); + var span = [2]u32{ 1, 1 }; // rowspan, colspan + while (true) { + while (i < src.len and (std.ascii.isWhitespace(src[i]) or src[i] == '/')) i += 1; + if (i >= src.len) return error.Malformed; + if (src[i] == '>') break; + const name_start = i; + while (i < src.len and !stops(src[i], "=>/")) i += 1; + const name = src[name_start..i]; + while (i < src.len and std.ascii.isWhitespace(src[i])) i += 1; + if (i >= src.len or src[i] != '=') continue; + i += 1; + while (i < src.len and std.ascii.isWhitespace(src[i])) i += 1; + if (i >= src.len) return error.Malformed; + var value: []const u8 = undefined; + if (src[i] == '"' or src[i] == '\'') { + const end = std.mem.indexOfScalarPos(u8, src, i + 1, src[i]) orelse return error.Malformed; + value = src[i + 1 .. end]; + i = end + 1; + } else { + const value_start = i; + while (i < src.len and !stops(src[i], ">")) i += 1; + value = src[value_start..i]; + } + if (std.ascii.eqlIgnoreCase(name, "rowspan")) span[0] = spanValue(value); + if (std.ascii.eqlIgnoreCase(name, "colspan")) span[1] = spanValue(value); + } + p.i = i + 1; + try p.apply(kind, close, span); + } + + fn apply(p: *Parser, kind: Tag, close: bool, span: [2]u32) ParseError!void { + switch (p.where) { + .before => { + if (kind != .table or close) return error.Malformed; + p.where = .inside; + return; + }, + .after => return error.Malformed, + .inside => {}, + } + switch (kind) { + .table => { + if (!close) return error.Malformed; // a table inside a cell + try p.endCell(); + p.row_open = false; + p.where = .after; + }, + .tr => { + try p.endCell(); + p.row_open = false; + if (!close) try p.openRow(); + }, + .td, .th => { + try p.endCell(); + if (!close) try p.openCell(kind == .th, span); + }, + .br => if (p.cell != null) { + p.space = false; + try p.text.append(p.a, '\n'); + }, + .other => {}, + } + } + + fn openRow(p: *Parser) ParseError!void { + if (p.rows == max_rows) return error.Malformed; + p.rows += 1; + p.row_open = true; + p.col = 0; + } + + /// A cell at the open row's first free column; its spans clipped to the grid and to columns earlier rows cover. + fn openCell(p: *Parser, head: bool, span: [2]u32) ParseError!void { + if (!p.row_open) return error.Malformed; + const row = p.rows - 1; + while (p.col < max_cols and p.until[p.col] > row) p.col += 1; + if (p.col == max_cols) return error.Malformed; + var cols: u32 = 0; + while (cols < span[1] and p.col + cols < max_cols and p.until[p.col + cols] <= row) cols += 1; + const rows = @min(span[0], max_rows - row); + for (p.until[p.col..][0..cols]) |*u| u.* = row + rows; + p.cell = .{ .text = "", .row = row, .col = p.col, .rows = rows, .cols = cols, .head = head }; + p.col += cols; + } + + fn endCell(p: *Parser) ParseError!void { + var cell = p.cell orelse return; + p.cell = null; + p.space = false; + defer p.text.clearRetainingCapacity(); + try p.cells.ensureUnusedCapacity(p.a, 1); + cell.text = try p.a.dupe(u8, std.mem.trim(u8, p.text.items, " \n")); + p.cells.appendAssumeCapacity(cell); + } +}; diff --git a/zig/src/core/ocr_html_test.zig b/zig/src/core/ocr_html_test.zig new file mode 100644 index 00000000000..dafaf11bf34 --- /dev/null +++ b/zig/src/core/ocr_html_test.zig @@ -0,0 +1,269 @@ +//! ocr_html.zig against tables dots.mocr and GLM-OCR wrote, and the subset's edges. +const std = @import("std"); +const html = @import("ocr_html.zig"); +const t = std.testing; + +// dots.mocr, TABLE II of the layout oracle (run008), first body rows kept. +const dots_table = "" ++ + "" ++ + "" ++ + "" ++ + "" ++ + "" ++ + "
ControlsHodgkin lymphomaNon-Hodgkin lymphoma
CasesOR95% CICasesOR95% CI
Never Exposed (reference group)2,2733151.00—1,8231.00—
>1 year9070.840.37-1.91861.170.87-1.59
"; + +// GLM-OCR, oracle case table-007: no header row. +const glm_table = "" ++ + "
ItemQuantity
Apples3
Pears5
"; + +fn expectCell(got: html.Cell, text: []const u8, row: u32, col: u32, rows: u32, cols: u32, head: bool) !void { + const want: html.Cell = .{ .text = text, .row = row, .col = col, .rows = rows, .cols = cols, .head = head }; + try t.expectEqualDeep(want, got); +} + +fn expectRows(a: std.mem.Allocator, table: html.Table, want: []const []const u8) !void { + const got = try html.rowsText(a, table); + try t.expectEqual(want.len, got.len); + for (want, got) |w, g| try t.expectEqualStrings(w, g); +} + +test "ocr html: the dots.mocr table keeps its spans" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const table = try html.parse(a, dots_table); + try t.expectEqual(@as(u32, 4), table.rows); + try t.expectEqual(@as(u32, 8), table.cols); + try t.expect(table.spanned); + try expectCell(table.cells[0], "", 0, 0, 2, 1, true); + const nhl = html.Cell{ .text = "Non-Hodgkin lymphoma", .row = 0, .col = 5, .rows = 1, .cols = 3, .head = true }; + try t.expectEqualDeep(nhl, table.cells[3]); + try expectCell(table.cells[4], "Cases", 1, 2, 1, 1, true); + try t.expect(!table.cells[10].head); + try expectRows(a, table, &.{ + "\tControls\tHodgkin lymphoma\t\t\tNon-Hodgkin lymphoma\t\t", + "\t\tCases\tOR\t95% CI\tCases\tOR\t95% CI", + "Never Exposed (reference group)\t2,273\t315\t1.00\t—\t1,823\t1.00\t—", + ">1 year\t90\t7\t0.84\t0.37-1.91\t86\t1.17\t0.87-1.59", + }); +} + +test "ocr html: the GLM-OCR table has no header row and becomes a pipe table with its first row the header" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const table = try html.parse(a, glm_table); + try t.expect(!table.spanned); + for (table.cells) |c| try t.expect(!c.head); + const want = "| Item | Quantity |\n| --- | --- |\n| Apples | 3 |\n| Pears | 5 |"; + try t.expectEqualStrings(want, try html.pipe(a, table)); + try t.expectEqualStrings("" ++ + "
ItemQuantity
Apples3
Pears5
", try html.normalized(a, table)); +} + +test "ocr html: entities decode and unknown ones stay as written" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const src = "
a & b <c> "q" ABC x y z" ++ + " &bogus; & � � � � &
"; + const table = try html.parse(a, src); + const want = "a & b \"q\" ABC x y z &bogus; & \u{fffd} \u{fffd} \u{fffd} � &"; + try t.expectEqualStrings(want, table.cells[0].text); +} + +test "ocr html: other tags drop and keep their text; names and attributes in any case and quoting" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const src = "\n " ++ + "
" ++ + "Bold and it2!
xy
\n"; + const table = try html.parse(a, src); + try t.expectEqual(@as(u32, 3), table.cols); + try expectCell(table.cells[0], "Bold and it2!", 0, 0, 1, 2, true); + try expectCell(table.cells[2], "y", 1, 1, 1, 2, false); +} + +test "ocr html:
is a line break in the cell and a space in pipe cells and text rows" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const src = "

line one
line two
\n three
b
"; + const table = try html.parse(a, src); + try t.expectEqualStrings("line one\nline two\nthree", table.cells[0].text); + try t.expectEqualStrings("| line one line two three | b |\n| --- | --- |", try html.pipe(a, table)); + try expectRows(a, table, &.{"line one line two three\tb"}); + const clean = "
line one
line two
three		b
"; + try t.expectEqualStrings(clean, try html.normalized(a, table)); +} + +test "ocr html: pipe cells escape |" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const table = try html.parse(a, "
a|b
|c
"); + try t.expectEqualStrings("| a\\|b | |\n| --- | --- |\n| \\| | c |", try html.pipe(a, table)); +} + +test "ocr html: pipe cells escape a literal < that would start a tag, and keep other < as written" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const table = try html.parse(a, "
<em>x</em>a < b
"); + try t.expectEqualStrings("| <em>x</em> | a < b |\n| --- | --- |", try html.pipe(a, table)); +} + +test "ocr html: pipe cells escape a literal & that would start an entity, and keep other & as written" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const table = try html.parse(a, "
&lt; &#38; &#x26;AT&T & b; &;
"); + try t.expectEqualStrings("| &lt; &#38; &#x26; | AT&T & b; &; |\n| --- | --- |", try html.pipe(a, table)); +} + +test "ocr html: an unknown tag whose name starts like a table tag is dropped" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const table = try html.parse(a, "
xy
"); + try t.expectEqual(@as(usize, 1), table.cells.len); + try t.expectEqualStrings("xy", table.cells[0].text); +} + +test "ocr html: malformed tables and tables past the limits" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const bad = [_][]const u8{ + "", // unclosed table + "
a
", // unclosed quote + "
a
a