From 0f84952a9912932099dff3dac670f5af56f2290e Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Mon, 2 Nov 2026 21:44:42 +0000 Subject: [PATCH 01/24] feat(quant): keep dense weights in their stored 16-bit dtype A checkpoint published in BF16 widened to fp32 by quant.read takes twice its memory and needs an fp32 kernel; Qwen3-Embedding-8B would not fit a 32 GiB card that way. quant.readStored reads like read but keeps a BF16 or FP16 dense tensor as stored, dense.Matrix carries the weight's dtype for the kernels, and dense.reference reads each element by its dtype. Callers of quant.read keep fp32 weights and their bits; the host tests cover both reads, views and the halves. --- zig/src/core/quant/dense.zig | 104 +++++++++++++++++++++++++-- zig/src/core/quant/quant.zig | 8 +++ zig/src/core/quant_geometry_host.zig | 3 +- 3 files changed, 107 insertions(+), 8 deletions(-) diff --git a/zig/src/core/quant/dense.zig b/zig/src/core/quant/dense.zig index a36f114e6f..14cd691551 100644 --- a/zig/src/core/quant/dense.zig +++ b/zig/src/core/quant/dense.zig @@ -1,4 +1,5 @@ -//! Weights kept in float: fp32 [N, K], the unquantized projections a conversion leaves (an MTP head's fc). +//! Weights kept in float [N, K]: fp32 (the unquantized projections a conversion leaves, an MTP head's fc), or BF16 / +//! FP16 as stored when the loader asks (`readStored`: a checkpoint published in 16-bit floats). const std = @import("std"); const types = @import("types.zig"); @@ -36,6 +37,13 @@ pub fn read(a: Allocator, t: anytype, key: []const u8) !Host { return .{ .weight = try convert.float32(a, try t.get(join(&buf, &.{ key, ".weight" }))) }; } +/// The float weight in its stored dtype when that is BF16 or FP16 (shared, not copied), else widened as `read` does. +pub fn readStored(a: Allocator, t: anytype, key: []const u8) !Host { + var buf: [256]u8 = undefined; + const w = try t.get(join(&buf, &.{ key, ".weight" })); + return .{ .weight = if (w.dtype == .bf16 or w.dtype == .f16) w else try convert.float32(a, w) }; +} + pub fn rows(h: Host) usize { return h.weight.shape[0]; } @@ -72,24 +80,26 @@ pub fn upload(u: types.Uploader, h: Host) !Device { return .{ .weight = try u.tensor(h.weight) }; } -/// What the kernels read: the fp32 weight's address. -pub const Matrix = struct { weight: u64 }; +/// What the kernels read: the weight's address and its element type (fp32, or BF16 / FP16 from `readStored`). +pub const Matrix = struct { weight: u64, dtype: types.DType = .f32 }; pub const View = struct { n: u32, k: u32, matrix: Matrix }; pub fn view(d: Device, stacked: bool) error{UnsupportedTables}!View { const lead: usize = @intFromBool(stacked); - return .{ .n = @intCast(d.weight.dim(lead)), .k = @intCast(d.weight.dim(lead + 1)), .matrix = .{ .weight = d.weight.ptr } }; + return .{ .n = @intCast(d.weight.dim(lead)), .k = @intCast(d.weight.dim(lead + 1)), .matrix = .{ .weight = d.weight.ptr, .dtype = d.weight.dtype } }; } pub fn elements(h: Host) usize { return h.weight.numel(); } -/// The weight as fp64, row by row. +/// The weight as fp64, row by row, read in its own dtype. pub fn reference(h: Host, out: []f64) error{UnexpectedTensor}!void { - if (out.len != h.weight.numel()) return error.UnexpectedTensor; - for (out, 0..) |*o, i| o.* = convert.load(.f32, h.weight.bytes, i); + const dtype = h.weight.dtype; + if (dtype != .f32 and dtype != .bf16 and dtype != .f16) return error.UnexpectedTensor; + if (out.len != h.weight.numel() or h.weight.bytes.len != out.len * dtype.size()) return error.UnexpectedTensor; + for (out, 0..) |*o, i| o.* = convert.load(dtype, h.weight.bytes, i); } test "a float weight halves along K and cuts its rows" { @@ -105,3 +115,83 @@ test "a float weight halves along K and cuts its rows" { try reference(low, &out); try std.testing.expectEqualSlices(f64, &.{ 5, 6, 7, 8 }, &out); } + +/// A one-tensor table for the tests: `get` and `has` over a single `.weight`. +const OneTable = struct { + weight: Tensor, + pub fn get(t: *const OneTable, name: []const u8) !Tensor { + if (!std.mem.eql(u8, name, "p.weight")) return error.MissingTensor; + return t.weight; + } + pub fn has(_: *const OneTable, name: []const u8) bool { + return std.mem.eql(u8, name, "p.weight"); + } + /// What the affine reader asks of a table; this one declares no quantization. + pub fn width(_: *const OneTable, _: []const u8) error{UnsupportedQuantization}!@import("mlx.zig").Width { + return error.UnsupportedQuantization; + } +}; + +test "readStored keeps a BF16 or FP16 weight as stored; read still widens it to fp32" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + // 1.5, -2, 0.25, 3 in BF16 and in FP16 + const bf = [_]u16{ 0x3FC0, 0xC000, 0x3E80, 0x4040 }; + const hf = [_]u16{ 0x3E00, 0xC000, 0x3400, 0x4200 }; + for ([_]struct { dtype: types.DType, bits: []const u16 }{ .{ .dtype = .bf16, .bits = &bf }, .{ .dtype = .f16, .bits = &hf } }) |c| { + const t: OneTable = .{ .weight = .{ .dtype = c.dtype, .rank = 2, .shape = .{ 2, 2, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(c.bits) } }; + const stored = try readStored(a, &t, "p"); + try std.testing.expectEqual(c.dtype, stored.weight.dtype); + try std.testing.expectEqual(@as(usize, 8), stored.weight.bytes.len); + var got: [4]f64 = undefined; + try reference(stored, &got); + try std.testing.expectEqualSlices(f64, &.{ 1.5, -2, 0.25, 3 }, &got); + const wide = try read(a, &t, "p"); + try std.testing.expectEqual(types.DType.f32, wide.weight.dtype); + try reference(wide, &got); + try std.testing.expectEqualSlices(f64, &.{ 1.5, -2, 0.25, 3 }, &got); + // a BF16 host halves along K and cuts rows by its own element size + const pair = try halves(a, stored); + try std.testing.expectEqual(@as(usize, 4), pair[1].weight.bytes.len); + var right: [2]f64 = undefined; + try reference(pair[1], &right); + try std.testing.expectEqualSlices(f64, &.{ -2, 3 }, &right); + } +} + +test "the view of an uploaded 16-bit weight reports its dtype" { + var b: types.Buf = undefined; + b.ptr = 0x1000; + b.dtype = .bf16; + b.rank = 2; + b.shape = .{ 3, 128, 1, 1, 1 }; + b.len = 3 * 128 * 2; + const v = try view(.{ .weight = b }, false); + try std.testing.expectEqual(types.DType.bf16, v.matrix.dtype); + try std.testing.expectEqual(@as(u32, 3), v.n); + try std.testing.expectEqual(@as(u32, 128), v.k); +} + +test "a reference refuses bytes that do not match the dtype's size" { + const bits = [_]u16{ 0x3FC0, 0xC000 }; + const h: Host = .{ .weight = .{ .dtype = .f32, .rank = 2, .shape = .{ 1, 1, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(&bits)[0..2] } }; + var out: [1]f64 = undefined; + try std.testing.expectError(error.UnexpectedTensor, reference(h, &out)); +} + +test "quant.readStored dispatches a float tensor to the dense format in its stored dtype; quant.read widens it" { + const quant = @import("quant.zig"); + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const bf = [_]u16{ 0x3FC0, 0xC000, 0x3E80, 0x4040 }; + const t: OneTable = .{ .weight = .{ .dtype = .bf16, .rank = 2, .shape = .{ 2, 2, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(&bf) } }; + const stored = try quant.readStored(a, &t, "p"); + try std.testing.expectEqual(types.DType.bf16, stored.dense.weight.dtype); + const wide = try quant.read(a, &t, "p"); + try std.testing.expectEqual(types.DType.f32, wide.dense.weight.dtype); + const f32s = [_]f32{ 1, 2 }; + const t32: OneTable = .{ .weight = .{ .dtype = .f32, .rank = 2, .shape = .{ 1, 2, 1, 1, 1 }, .bytes = std.mem.sliceAsBytes(&f32s) } }; + try std.testing.expectEqual(types.DType.f32, (try quant.readStored(a, &t32, "p")).dense.weight.dtype); +} diff --git a/zig/src/core/quant/quant.zig b/zig/src/core/quant/quant.zig index bf46e04dd5..26581a9980 100644 --- a/zig/src/core/quant/quant.zig +++ b/zig/src/core/quant/quant.zig @@ -110,6 +110,14 @@ pub fn read(a: Allocator, t: anytype, key: []const u8) !Host { return types.refuse(key, "no weight format reads this tensor"); } +/// `read`, keeping a dense tensor stored in BF16 or FP16 in that dtype (every other format reads as `read` does). +pub fn readStored(a: Allocator, t: anytype, key: []const u8) !Host { + inline for (read_order) |f| { + if (Module(f).matches(t, key)) return wrap(f, try if (f == .dense) dense.readStored(a, t, key) else Module(f).read(a, t, key)); + } + return types.refuse(key, "no weight format reads this tensor"); +} + /// A fused [embedding | hidden] projection as its two halves along K. pub fn halves(a: Allocator, h: Host) ![2]Host { switch (h) { diff --git a/zig/src/core/quant_geometry_host.zig b/zig/src/core/quant_geometry_host.zig index 41aa2ab82d..f96f791889 100644 --- a/zig/src/core/quant_geometry_host.zig +++ b/zig/src/core/quant_geometry_host.zig @@ -1,6 +1,7 @@ -//! Standalone host entry for the actual affine reader, expert stack and bit-preserving slicing regressions. +//! Standalone host entry for the actual affine reader, dense stored dtypes, expert stack and bit-exact slicing. test { _ = @import("quant/mlx.zig"); _ = @import("quant/experts.zig"); _ = @import("quant/exact_test.zig"); + _ = @import("quant/dense.zig"); } From 43aa5e3d4ccefb2bd9e1a33c369d578b91fec7e2 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Mon, 2 Nov 2026 22:28:25 +0000 Subject: [PATCH 02/24] style(quant): keep the dense module's doc comment on one line The repository's lean checker allows only single-line comments of at most 120 columns, and the stored-dtype change gave the dense format's module comment a second line. The comment now says the same in one line, so the checker adds no problem over main. --- zig/src/core/quant/dense.zig | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/zig/src/core/quant/dense.zig b/zig/src/core/quant/dense.zig index 14cd691551..c989a48d83 100644 --- a/zig/src/core/quant/dense.zig +++ b/zig/src/core/quant/dense.zig @@ -1,5 +1,4 @@ -//! Weights kept in float [N, K]: fp32 (the unquantized projections a conversion leaves, an MTP head's fc), or BF16 / -//! FP16 as stored when the loader asks (`readStored`: a checkpoint published in 16-bit floats). +//! Float weights [N, K]: fp32 (an MTP head's fc a conversion leaves), or BF16 / FP16 as stored (`readStored`). const std = @import("std"); const types = @import("types.zig"); From 8c0c1d04b0ae18be08fc0afde3580f0f086639fe Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Mon, 2 Nov 2026 22:38:47 +0000 Subject: [PATCH 03/24] feat(hip): run BF16 dense weights through the dot2 GEMM tile A BF16 checkpoint kept as stored needs a prompt GEMM that reads BF16 weights; HIP had only the fp32 one-row dense kernel. A BF16 weight decoder plugs into the existing dot2 GEMM tile with unit scales and no bias path, which the tile now leaves out for such a decoder; the MLX tiles compile to the same code. project dispatches a dense weight by its dtype, and the dense embedding gather reads BF16 tables. GPU tests check it against float64 on gfx1201. --- zig/build/hip.zig | 18 ++--- zig/kernels/hip/ops/elementwise.hpp | 6 +- zig/kernels/hip/quant/bf16_decoder.hpp | 65 +++++++++++++++ zig/kernels/hip/quant/dense_tiles.hpp | 25 ++++++ zig/kernels/hip/quant/mlx_decoder.hpp | 2 + zig/kernels/hip/tiles/dot2_tiles.hip | 4 + zig/kernels/hip/tiles/gemm.hpp | 15 +++- zig/src/hip/kernel_tests.zig | 1 + zig/src/hip/kernel_tests/dense.zig | 107 +++++++++++++++++++++++++ zig/src/hip/launches.zig | 7 +- zig/src/hip/launches/dense.zig | 84 +++++++++++++++++++ zig/src/hip/launches/elementwise.zig | 3 +- zig/src/hip/ops/ops.zig | 1 + zig/src/hip/ops/project.zig | 40 +++++++-- 14 files changed, 355 insertions(+), 23 deletions(-) create mode 100644 zig/kernels/hip/quant/bf16_decoder.hpp create mode 100644 zig/kernels/hip/quant/dense_tiles.hpp create mode 100644 zig/src/hip/kernel_tests/dense.zig create mode 100644 zig/src/hip/launches/dense.zig diff --git a/zig/build/hip.zig b/zig/build/hip.zig index 9c98b050d6..22cb6d002a 100644 --- a/zig/build/hip.zig +++ b/zig/build/hip.zig @@ -335,15 +335,15 @@ const groups = [_][]const u8{ /// What the group sources include, so an edit to one rebuilds them. const headers = [_][]const u8{ - "attention/attention_fa.hip", "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", + "attention/attention_fa.hip", "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", }; /// The flags of the kernels' first build: no contraction, wave32 on RDNA, C++20. diff --git a/zig/kernels/hip/ops/elementwise.hpp b/zig/kernels/hip/ops/elementwise.hpp index 0b1d5bb7ae..13037942b0 100644 --- a/zig/kernels/hip/ops/elementwise.hpp +++ b/zig/kernels/hip/ops/elementwise.hpp @@ -32,11 +32,13 @@ extern "C" __global__ void tf_embed_rows(const uint32_t* words, const void* scal } // The float gather: rows of an fp32 table (an embedding kept unquantized), written in the activation type. -extern "C" __global__ void tf_embed_dense(const float* table, const int* ids, int n, int k, void* out, int out_kind) { +// Rows `ids` of a float table stored as `table_kind` (fp32, or BF16 / FP16 as a checkpoint stores them). +extern "C" __global__ void tf_embed_dense(const void* table, int table_kind, const int* ids, int n, int k, void* out, + int out_kind) { long long i = gid(); if (i >= static_cast(n) * k) return; int r = static_cast(i / k), c = static_cast(i % k); - store(out, out_kind, i, table[static_cast(ids[r]) * k + c]); + store(out, out_kind, i, load(table, table_kind, static_cast(ids[r]) * k + c)); } extern "C" __global__ void tf_cast(const void* src, int skind, void* dst, int dkind, long long n) { diff --git a/zig/kernels/hip/quant/bf16_decoder.hpp b/zig/kernels/hip/quant/bf16_decoder.hpp new file mode 100644 index 0000000000..77fe35e686 --- /dev/null +++ b/zig/kernels/hip/quant/bf16_decoder.hpp @@ -0,0 +1,65 @@ +#pragma once + +// The BF16 dense WeightDecoder: a column's K weights stored as BF16, two a 32-bit word; no scales, no biases. + +#include + +#include "common/vec.hpp" +#include "quant/mlx.hpp" + +namespace tf { +namespace rocm { + +// The GEMM tile's decoder over raw BF16 weights (Affine args with bits 16; its tables are unused). Every group of +// `a.group` K folds into the accumulator with scale 1, so the sums are the fp32 dot2 chain, group after group. +struct Bf16Decoder { + using Args = Affine; + static constexpr int kBits = 16; + // Words of a 32-weight chunk of a column. + static constexpr int kWords = 16; + // No bias term: the tile neither sums x nor folds a bias. + static constexpr bool kBias = false; + + // A plain product only: a routed plan's items are not this decoder's. + __device__ static bool take_item(Args& a, int) { return a.route.items == nullptr; } + + __device__ static const uint32_t* words(const Args& a) { return a.words; } + + template + __device__ static Idx row_words(const Args& a) { + return static_cast(a.k) / 2; + } + + // A chunk is 64 bytes: four 16-byte loads (the host checks the alignment). + template + __device__ static void load(const uint32_t* src, uint32_t (&w)[kWords]) { +#pragma unroll + for (int i = 0; i < kWords / 4; ++i) { + const u32x4 v = reinterpret_cast(src)[i]; + w[4 * i] = v.x; + w[4 * i + 1] = v.y; + w[4 * i + 2] = v.z; + w[4 * i + 3] = v.w; + } + } + + // Weights t and t + 1 of a chunk (t even): the word that holds them, as stored. + template + __device__ static typename T::pair pair(const uint32_t (&w)[kWords], int t) { + static_assert(std::is_same_v, "BF16 weights take BF16 activations"); + return __builtin_bit_cast(typename T::pair, w[t >> 1]); + } + + struct Term {}; + + __device__ static Term term(const Args&, long long) { return {}; } + __device__ static float scale_value(const Args&, const Term&) { return 1.f; } + __device__ static float bias_value(const Args&, const Term&) { return 0.f; } + + // A group's dot added as it is: the scale is 1. + __device__ static float fold_scale(float acc, float dot, float) { return acc + dot; } + __device__ static float fold_bias(float acc, float, float) { return acc; } +}; + +} // namespace rocm +} // namespace tf diff --git a/zig/kernels/hip/quant/dense_tiles.hpp b/zig/kernels/hip/quant/dense_tiles.hpp new file mode 100644 index 0000000000..dccb4266fe --- /dev/null +++ b/zig/kernels/hip/quant/dense_tiles.hpp @@ -0,0 +1,25 @@ +#pragma once + +// BF16 dense kernels: the shared prefill GEMM tile over the BF16 decoder, named as the launchers look up. + +#include "quant/act.hpp" +#include "quant/bf16_decoder.hpp" +#include "tiles/epilogue.hpp" +#include "tiles/gemm.hpp" + +namespace tf { +namespace rocm { + +// y = x . W^T for BF16 x and BF16 W at any row count; a build without the BF16 dot2 keeps the symbol and traps. +template +__global__ void __launch_bounds__(256) dense_gemm_block(Affine a) { +#if TF_DEVICE_BF16_DOT2 + gemm_tile, T, GemmOut, RT>(a); +#else + (void)a; + __builtin_trap(); +#endif +} + +} // namespace rocm +} // namespace tf diff --git a/zig/kernels/hip/quant/mlx_decoder.hpp b/zig/kernels/hip/quant/mlx_decoder.hpp index 6f94a58974..f35b257654 100644 --- a/zig/kernels/hip/quant/mlx_decoder.hpp +++ b/zig/kernels/hip/quant/mlx_decoder.hpp @@ -21,6 +21,8 @@ struct MlxDecoder { static constexpr int kBits = BITS; // Words of a 32-code chunk of a column. static constexpr int kWords = BITS; + // A bias a group: the tiles sum x and fold the bias with it. + static constexpr bool kBias = true; // ---- the plan: a routed item's expert and rows ---- diff --git a/zig/kernels/hip/tiles/dot2_tiles.hip b/zig/kernels/hip/tiles/dot2_tiles.hip index 65ef2d4e85..13fd252479 100644 --- a/zig/kernels/hip/tiles/dot2_tiles.hip +++ b/zig/kernels/hip/tiles/dot2_tiles.hip @@ -6,6 +6,7 @@ #include "tiles/stream.hpp" #include "tiles/matrix_gemm.hpp" #include "quant/mlx_tiles.hpp" +#include "quant/dense_tiles.hpp" namespace tf { namespace rocm { @@ -376,5 +377,8 @@ TF_TILES(5) TF_TILES(6) TF_TILES(8) +// The BF16 dense GEMM on every build (its body traps where there is no BF16 dot2), so every image has the symbol. +template __global__ void dense_gemm_block::rt>(Affine); + } // namespace rocm } // namespace tf diff --git a/zig/kernels/hip/tiles/gemm.hpp b/zig/kernels/hip/tiles/gemm.hpp index 6a95340979..b67bedd55e 100644 --- a/zig/kernels/hip/tiles/gemm.hpp +++ b/zig/kernels/hip/tiles/gemm.hpp @@ -41,6 +41,7 @@ struct GemmShape { }; // One 128 x 128 block of y = x . W^T: Dec reads weights and group terms, Act the rows of x, T multiplies, Epi stores. +// A decoder without a bias term (Dec::kBias false) leaves out the x sums and the bias folds; the MLX one keeps them. template __device__ __forceinline__ void gemm_tile(typename Dec::Args& a) { if (!Dec::take_item(a, blockIdx.z)) return; @@ -122,7 +123,7 @@ __device__ __forceinline__ void gemm_tile(typename Dec::Args& a) { } if (side == 0 && st % per == per - 1) { sc[(st / per) % 3][line] = Dec::scale_value(a, term[J]); - bi[(st / per) % 3][line] = Dec::bias_value(a, term[J]); + if constexpr (Dec::kBias) bi[(st / per) % 3][line] = Dec::bias_value(a, term[J]); } }; // Side 1 sums a row's x of the stage in LDS, lo then hi of each pair in order, on through the group. @@ -180,7 +181,9 @@ __device__ __forceinline__ void gemm_tile(typename Dec::Args& a) { // Two stages a trip so a stage's buffer and parity are known when the code is built. auto step = [&](std::integral_constant, int s) { constexpr int cur = P; - if (live && s > 0 && (s - 1) % per == per - 1) apply_bias(((s - 1) / per) % 3, ((s - 1) / per) & 1); + if constexpr (Dec::kBias) { + if (live && s > 0 && (s - 1) % per == per - 1) apply_bias(((s - 1) / per) % 3, ((s - 1) / per) & 1); + } if constexpr (P == 1) fetch(s + 1); if (live) { constexpr int kXRegs = RT / 4; // x words (4 pairs each) a lane loads a chunk @@ -235,7 +238,9 @@ __device__ __forceinline__ void gemm_tile(typename Dec::Args& a) { } } } - if (side == 1 && sxlive) sumx(s, cur); + if constexpr (Dec::kBias) { + if (side == 1 && sxlive) sumx(s, cur); + } stash(std::integral_constant{}, s + 1, cur ^ 1); __syncthreads(); }; @@ -243,7 +248,9 @@ __device__ __forceinline__ void gemm_tile(typename Dec::Args& a) { step(std::integral_constant{}, s); if (s + 1 < stages) step(std::integral_constant{}, s + 1); } - if (live && (stages - 1) % per == per - 1) apply_bias(((stages - 1) / per) % 3, ((stages - 1) / per) & 1); + if constexpr (Dec::kBias) { + if (live && (stages - 1) % per == per - 1) apply_bias(((stages - 1) / per) % 3, ((stages - 1) / per) & 1); + } #pragma unroll for (int r = 0; r < RT; ++r) { const int row = m0 + rg * 16 + hrow + r; diff --git a/zig/src/hip/kernel_tests.zig b/zig/src/hip/kernel_tests.zig index ff98d00f33..440b1062df 100644 --- a/zig/src/hip/kernel_tests.zig +++ b/zig/src/hip/kernel_tests.zig @@ -26,6 +26,7 @@ test { _ = @import("kernel_tests/attention_pin.zig"); _ = @import("kernel_tests/stats.zig"); _ = @import("kernel_tests/ops.zig"); + _ = @import("kernel_tests/dense.zig"); _ = @import("kernel_tests/argmax.zig"); _ = @import("kernel_tests/tp.zig"); _ = @import("kernel_tests/affine/affine.zig"); diff --git a/zig/src/hip/kernel_tests/dense.zig b/zig/src/hip/kernel_tests/dense.zig new file mode 100644 index 0000000000..2d86400116 --- /dev/null +++ b/zig/src/hip/kernel_tests/dense.zig @@ -0,0 +1,107 @@ +//! The BF16 dense GEMM against float64: within the dot2 chain's bound, BF16 out rounded once, rows batch-independent. +const std = @import("std"); +const rig = @import("rig.zig"); +const Rig = rig.Rig; +const gpa = rig.gpa; +const at = rig.at; +const ops = @import("../ops/ops.zig"); +const Tensor = ops.Tensor; + +const Case = struct { m: usize, n: usize, k: usize }; + +/// Edges of both tile dimensions, one and several groups of K, and Qwen3's widest K. +const cases = [_]Case{ + .{ .m = 1, .n = 2, .k = 128 }, + .{ .m = 7, .n = 130, .k = 4096 }, + .{ .m = 128, .n = 256, .k = 128 }, + .{ .m = 129, .n = 300, .k = 4096 }, + .{ .m = 300, .n = 130, .k = 12288 }, +}; + +/// The fp32 value as BF16, rounded to nearest even (the tile's epilogue; the rig's bf16Bits truncates). +fn nearest(v: f32) u16 { + const u: u32 = @bitCast(v); + if (std.math.isNan(v)) return 0x7fc0; + return @intCast((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +fn bf16Of(rng: *rig.Rng) u16 { + return rig.bf16Bits(@floatCast(rng.normal())); +} + +test "dense GEMM: BF16 weights against float64, fp32 within the dot2 chain's bound, BF16 out rounded once" { + const t = try Rig.open(64 << 20); + defer t.close(); + if (!t.bf16) return error.SkipZigTest; + const o = t.ops(&t.on); + var rng: rig.Rng = .{ .state = 0xDE45E }; + for (cases) |c| { + const hx = try gpa.alloc(u16, c.m * c.k); + defer gpa.free(hx); + const hw = try gpa.alloc(u16, c.n * c.k); + defer gpa.free(hw); + for (hx) |*v| v.* = bf16Of(&rng); + for (hw) |*v| v.* = bf16Of(&rng); + var x = try t.upload(hx); + defer x.free(); + var w = try t.upload(hw); + defer w.free(); + var y32 = try t.alloc(c.m * c.n * 4); + defer y32.free(); + var y16 = try t.alloc(c.m * c.n * 2); + defer y16.free(); + const xt: Tensor = .{ .ptr = at(x), .kind = .bf16 }; + try o.denseGemm(xt, at(w), .{ .ptr = at(y32), .kind = .f32 }, c.m, c.n, c.k); + try o.denseGemm(xt, at(w), .{ .ptr = at(y16), .kind = .bf16 }, c.m, c.n, c.k); + try t.stream.synchronize(); + const g32 = try rig.download(f32, y32); + defer gpa.free(g32); + const g16 = try rig.download(u16, y16); + defer gpa.free(g16); + // a group's 64 dot2s (two roundings each at most), then one add a group: the longest rounding chain of a sum + const depth: f64 = 128 + @as(f64, @floatFromInt(c.k / 128)) + 2; + var worst: f64 = 0; + for (0..c.m) |r| for (0..c.n) |j| { + var want: f64 = 0; + var mag: f64 = 0; + for (0..c.k) |i| { + const p = rig.bf16Value(hx[r * c.k + i]) * rig.bf16Value(hw[j * c.k + i]); + want += p; + mag += @abs(p); + } + const got = g32[r * c.n + j]; + const err = @abs(@as(f64, got) - want); + worst = @max(worst, err / @max(mag, 1e-30)); + try rig.expect(err <= depth * std.math.pow(f64, 2, -24) * mag, "dense m {d} n {d} k {d} at ({d}, {d}): {d} vs {d} (mag {d})", .{ c.m, c.n, c.k, r, j, got, want, mag }); + try rig.expect(g16[r * c.n + j] == nearest(got), "dense bf16 out at ({d}, {d}) is {x}, not {x}: the fp32 sum {d} rounded once", .{ r, j, g16[r * c.n + j], nearest(got), got }); + }; + std.log.info("dense m {d} n {d} k {d}: worst error {e:.2} of the magnitude", .{ c.m, c.n, c.k, worst }); + // the last row alone has the bits it has in the batch + var lone = try t.alloc(c.n * 4); + defer lone.free(); + const last: Tensor = .{ .ptr = at(x) + (c.m - 1) * c.k * 2, .kind = .bf16 }; + try o.denseGemm(last, at(w), .{ .ptr = at(lone), .kind = .f32 }, 1, c.n, c.k); + try t.stream.synchronize(); + const g1 = try rig.download(f32, lone); + defer gpa.free(g1); + try rig.expect(std.mem.eql(u8, std.mem.sliceAsBytes(g1), std.mem.sliceAsBytes(g32[(c.m - 1) * c.n ..][0..c.n])), "dense m {d}: the last row alone differs from its batch bits", .{c.m}); + } +} + +test "dense GEMM: misaligned x, K off a group and FP16 operands are refused before launch" { + const t = try Rig.open(1 << 20); + defer t.close(); + if (!t.bf16) return error.SkipZigTest; + const o = t.ops(&t.on); + var x = try t.alloc(4 * 256 * 2 + 16); + defer x.free(); + var w = try t.alloc(4 * 256 * 2); + defer w.free(); + var y = try t.alloc(4 * 4 * 4); + defer y.free(); + const out: Tensor = .{ .ptr = at(y), .kind = .f32 }; + try std.testing.expectError(error.BadShape, o.denseGemm(.{ .ptr = at(x) + 2, .kind = .bf16 }, at(w), out, 1, 4, 256)); + try std.testing.expectError(error.BadShape, o.denseGemm(.{ .ptr = at(x), .kind = .bf16 }, at(w), out, 1, 4, 200)); + try std.testing.expectError(error.BadShape, o.denseGemm(.{ .ptr = at(x), .kind = .f16 }, at(w), out, 1, 4, 256)); + try std.testing.expectError(error.BadShape, o.denseGemm(.{ .ptr = at(x), .kind = .bf16 }, at(w), .{ .ptr = at(y), .kind = .f16 }, 1, 4, 256)); +} diff --git a/zig/src/hip/launches.zig b/zig/src/hip/launches.zig index a2aae7237b..54543599cc 100644 --- a/zig/src/hip/launches.zig +++ b/zig/src/hip/launches.zig @@ -9,6 +9,7 @@ const Function = @import("module.zig").Function; const Stream = @import("stream.zig").Stream; const caps = @import("caps.zig"); const affine = @import("launches/affine.zig"); +const dense = @import("launches/dense.zig"); const util = @import("launches/util.zig"); @@ -86,6 +87,8 @@ pub const Launcher = struct { wide: bool, /// The MLX affine products and the registry that picks their tiles. affine: affine.Kernels, + /// The BF16 dense GEMM (raw BF16 weights through the dot2 GEMM tile). + dense: dense.Kernels, /// Flash Next's kernels: its residual streams, n-gram rows, indexer and sparse attention. fnx: flashnext.Fns, /// Nemotron's Mamba-2 mixer kernels. @@ -291,7 +294,9 @@ pub const Launcher = struct { l.mamba = try mamba.Fns.load(l.mods[@backingInt(kernels.Group.mamba)]); const tiles = l.mods[@backingInt(kernels.Group.affine_tiles)]; const dot2 = l.mods[@backingInt(kernels.Group.affine_dot2)]; - l.affine = try affine.Kernels.load(tiles, dot2, caps.Caps.of(kernels.arch) orelse return error.Invalid, choices.products); + const chip = caps.Caps.of(kernels.arch) orelse return error.Invalid; + l.affine = try affine.Kernels.load(tiles, dot2, chip, choices.products); + l.dense = try dense.Kernels.load(tiles, chip.act == .bf16); try l.affine.fillByteLut(r, dot2); l.fnx = try flashnext.Fns.load(l.mods[@backingInt(kernels.Group.flashnext)]); return l; diff --git a/zig/src/hip/launches/dense.zig b/zig/src/hip/launches/dense.zig new file mode 100644 index 0000000000..ec771a0fd0 --- /dev/null +++ b/zig/src/hip/launches/dense.zig @@ -0,0 +1,84 @@ +//! The BF16 dense GEMM of quant/dense_tiles.hpp (the dot2 GEMM tile over raw BF16 weights) launched from Zig. + +const std = @import("std"); +const abi = @import("../abi.zig"); +const runtime = @import("../runtime.zig"); +const hl = @import("../launch.zig"); +const Module = @import("../module.zig").Module; +const Function = @import("../module.zig").Function; +const affine = @import("affine.zig"); +const util = @import("util.zig"); + +const Error = runtime.Error; + +/// K a group folds into the accumulator (the tile's group; the weights carry no scales). +pub const group: c_int = 128; +/// The tile's rows and columns a block. +const side: c_int = 128; + +const symbol = "_ZN2tf4rocm16dense_gemm_blockINS0_7DotBF16ELi8EEEvNS0_6AffineE"; + +pub const Kernels = struct { + tile: Function, + /// The build's activations are BF16 (the tile's dot2); on an FP16 build the symbol only traps. + bf16: bool, + + /// `tiles` is dot2_tiles.hip's module, where every build keeps the symbol. + pub fn load(tiles: Module, bf16: bool) Error!Kernels { + return .{ .tile = try tiles.function(symbol), .bf16 = bf16 }; + } + + /// The product's refusal, or null when the tile takes it: BF16 build, K in whole groups, 16-byte rows and bases. + pub fn refusal(k: *const Kernels, x: u64, w: u64, m: c_int, n: c_int, kk: c_int) ?[]const u8 { + if (!k.bf16) return "the BF16 dense GEMM needs a BF16 build"; + if (m <= 0 or n <= 0 or kk <= 0) return "empty product"; + if (@rem(kk, group) != 0) return "K is not a multiple of 128"; + if (x % 16 != 0 or w % 16 != 0) return "x or the weights are not 16-byte aligned"; + return null; + } + + /// y = x . W^T for `m` BF16 rows of x and the [n, k] BF16 weights: fp32 into `out`, or BF16 into `out16`. + pub fn gemm(k: *const Kernels, r: *const runtime.Runtime, x: u64, w: u64, out: u64, out16: u64, m: c_int, n: c_int, kk: c_int, s: abi.Stream) Error!void { + if (k.refusal(x, w, m, n, kk)) |why| { + std.log.err("dense gemm: {s} (m {d}, n {d}, k {d})", .{ why, m, n, kk }); + return error.Invalid; + } + if ((out == 0) == (out16 == 0)) { + std.log.err("dense gemm: exactly one of out and out16", .{}); + return error.Invalid; + } + const a: affine.Arg = .{ + .x = x, + .words = w, + .scale = .{ .p = 0, .kind = 0 }, + .bias = .{ .p = 0, .kind = 0 }, + .out = out, + .m = m, + .n = n, + .k = kk, + .bits = 16, + .group = group, + .fp16 = 0, + .out16 = out16, + }; + if (s == util.counting) return; + var args: hl.Args = .{}; + try args.add(a); + try hl.launch(k.tile, .{ .grid = .{ .x = cdiv(n), .y = cdiv(m) }, .block = .{ .x = 256 } }, .{ .r = r, .handle = s }, &args); + } + + fn cdiv(v: c_int) c_uint { + return @intCast(@divTrunc(v + side - 1, side)); + } +}; + +test "the dense GEMM refuses what its tile does not take" { + const k: Kernels = .{ .tile = undefined, .bf16 = true }; + try std.testing.expect(k.refusal(0, 0, 3, 5, 256) == null); + try std.testing.expect(k.refusal(0, 0, 3, 5, 200) != null); + try std.testing.expect(k.refusal(8, 0, 3, 5, 256) != null); + try std.testing.expect(k.refusal(0, 8, 3, 5, 256) != null); + try std.testing.expect(k.refusal(0, 0, 0, 5, 256) != null); + const half: Kernels = .{ .tile = undefined, .bf16 = false }; + try std.testing.expect(half.refusal(0, 0, 3, 5, 256) != null); +} diff --git a/zig/src/hip/launches/elementwise.zig b/zig/src/hip/launches/elementwise.zig index 4f767f37fd..4dcdb77a22 100644 --- a/zig/src/hip/launches/elementwise.zig +++ b/zig/src/hip/launches/elementwise.zig @@ -28,9 +28,10 @@ pub fn tf_embed_rows(l: *const Launcher, words: C, scale: C, bias: C, scale_kind try l.flat(l.op.embed_rows, @as(i64, n) * k, s, &a); } -pub fn tf_embed_dense(l: *const Launcher, table: CF, ids: CI, n: c_int, k: c_int, out: P, out_kind: c_int, s: S) Error!void { +pub fn tf_embed_dense(l: *const Launcher, table: C, table_kind: c_int, ids: CI, n: c_int, k: c_int, out: P, out_kind: c_int, s: S) Error!void { var a: Args = .{}; try a.add(ad(table)); + try a.add(table_kind); try a.add(ad(ids)); try a.add(n); try a.add(k); diff --git a/zig/src/hip/ops/ops.zig b/zig/src/hip/ops/ops.zig index 8c842f4e2f..a7c32470d1 100644 --- a/zig/src/hip/ops/ops.zig +++ b/zig/src/hip/ops/ops.zig @@ -48,6 +48,7 @@ pub const Ops = struct { pub const affineRouted = products.affineRouted; pub const embedAffine = products.embedAffine; pub const denseRows = products.denseRows; + pub const denseGemm = products.denseGemm; /// Whether decode.hip's merged launches are on. pub fn fused(o: Ops) bool { diff --git a/zig/src/hip/ops/project.zig b/zig/src/hip/ops/project.zig index a4d7e38914..7638a75337 100644 --- a/zig/src/hip/ops/project.zig +++ b/zig/src/hip/ops/project.zig @@ -30,11 +30,21 @@ fn affineOf(proj: Projection, h: mlx.Matrix) Affine { pub fn project(o: Ops, x: Tensor, proj: Projection, m: usize, f32_out: bool) Error!Tensor { switch (proj.handle) { .mlx => |h| return affine(o, x, affineOf(proj, h), m, f32_out), - .dense => |h| { - if (f32_out or m == 0 or x.kind == .f32) return error.BadShape; - const out: Tensor = .{ .ptr = try o.arena.take(m * proj.n * x.kind.size()), .kind = x.kind }; - try denseRows(o, x, h.weight, out, m, proj.n, proj.k); - return out; + .dense => |h| switch (h.dtype) { + .f32 => { + if (f32_out or m == 0 or x.kind == .f32) return error.BadShape; + const out: Tensor = .{ .ptr = try o.arena.take(m * proj.n * x.kind.size()), .kind = x.kind }; + try denseRows(o, x, h.weight, out, m, proj.n, proj.k); + return out; + }, + .bf16 => { + if (m == 0 or x.kind != .bf16) return error.BadShape; + const kind: t.Kind = if (f32_out) .f32 else .bf16; + const out: Tensor = .{ .ptr = try o.arena.take(m * proj.n * kind.size()), .kind = kind }; + try denseGemm(o, x, h.weight, out, m, proj.n, proj.k); + return out; + }, + else => return error.BadShape, }, } } @@ -200,7 +210,7 @@ pub fn projectRouted(o: Ops, x: Tensor, proj: Projection, items: u64, count: usi pub fn embedRows(o: Ops, table: Projection, ids: u64, n: usize, out: Tensor) Error!void { switch (table.handle) { .mlx => |h| return embedAffine(o, affineOf(table, h), ids, n, out), - .dense => |h| try o.l.tf_embed_dense(f(h.weight), @ptrFromInt(ids), int(n), int(table.k), p(out.ptr), @backingInt(out.kind), o.stream), + .dense => |h| try o.l.tf_embed_dense(p(h.weight), @backingInt(try kindOf(h.dtype)), @ptrFromInt(ids), int(n), int(table.k), p(out.ptr), @backingInt(out.kind), o.stream), } } @@ -351,3 +361,21 @@ pub fn embedAffine(o: Ops, table: Affine, ids: u64, n: usize, out: Tensor) Error pub fn denseRows(o: Ops, x: Tensor, w: u64, out: Tensor, rows: usize, n: usize, k: usize) Error!void { try o.l.tf_dense_rows(p(x.ptr), @backingInt(x.kind), f(w), p(out.ptr), int(rows), int(n), int(k), o.stream); } + +/// x (rows, k) BF16 times BF16 weights [n, k] on the dot2 GEMM tile: `out` fp32, or BF16 rounded once. +pub fn denseGemm(o: Ops, x: Tensor, w: u64, out: Tensor, rows: usize, n: usize, k: usize) Error!void { + if (x.kind != .bf16 or out.kind == .f16) return error.BadShape; + const wide = out.kind == .f32; + if (o.l.dense.refusal(x.ptr, w, int(rows), int(n), int(k)) != null) return error.BadShape; + try o.l.dense.gemm(o.l.r, x.ptr, w, if (wide) out.ptr else 0, if (wide) 0 else out.ptr, int(rows), int(n), int(k), o.stream); +} + +/// The kernels' element kind of a dense weight's dtype. +fn kindOf(dtype: quant.DType) Error!t.Kind { + return switch (dtype) { + .f32 => .f32, + .bf16 => .bf16, + .f16 => .f16, + else => error.BadShape, + }; +} From cdc4ca48c6225fb776907a6bf8a859d12b6afff4 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Mon, 2 Nov 2026 21:53:20 +0000 Subject: [PATCH 04/24] feat(qwen3): parse the dense Qwen3 config for HIP The HIP qwen3 forward runs one shape of Qwen3 dense model, so its config reader takes the geometry and constants from config.json, defaults head_dim to 128 as Qwen3Config does, and refuses by name the features the forward does not run (projection biases, sliding window, RoPE scaling, head_dim above 256, a non-SiLU MLP, uneven GQA groups), with host tests on the 8B embedding and reranker configs. --- zig/src/families/qwen3/config.zig | 153 ++++++++++++++++++++++++++++++ 1 file changed, 153 insertions(+) create mode 100644 zig/src/families/qwen3/config.zig diff --git a/zig/src/families/qwen3/config.zig b/zig/src/families/qwen3/config.zig new file mode 100644 index 0000000000..35e6e15136 --- /dev/null +++ b/zig/src/families/qwen3/config.zig @@ -0,0 +1,153 @@ +//! A Qwen3 dense checkpoint's config.json: the geometry the HIP forward runs, and refusals of what it does not. +const std = @import("std"); +const Allocator = std.mem.Allocator; + +pub const Error = error{ InvalidConfig, UnsupportedConfig } || Allocator.Error; + +/// The forward's geometry and constants, as config.json names them. +pub const Spec = struct { + layers: usize, + hidden: usize, + heads: usize, + kv_heads: usize, + head_dim: usize, + inter: usize, + vocab: usize, + eps: f64, + theta: f64, + max_positions: usize, + /// `tie_word_embeddings`: the output head is the embedding table. + tied: bool, +}; + +/// The spec of config.json's `bytes`; a missing or malformed key is InvalidConfig, an unrun feature UnsupportedConfig. +pub fn parse(gpa: Allocator, bytes: []const u8) Error!Spec { + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const doc = std.json.parseFromSliceLeaky(std.json.Value, arena.allocator(), bytes, .{}) catch |err| switch (err) { + error.OutOfMemory => return error.OutOfMemory, + else => { + std.log.warn("qwen3 config.json is not JSON ({s})", .{@errorName(err)}); + return error.InvalidConfig; + }, + }; + if (doc != .object) return invalid("the document", "is not an object"); + const o = doc.object; + const model_type = o.get("model_type") orelse return invalid("model_type", "is missing"); + if (model_type != .string or !std.mem.eql(u8, model_type.string, "qwen3")) return unsupported("model_type", "is not qwen3"); + const s: Spec = .{ + .layers = try count(o, "num_hidden_layers"), + .hidden = try count(o, "hidden_size"), + .heads = try count(o, "num_attention_heads"), + .kv_heads = try count(o, "num_key_value_heads"), + // Qwen3Config's own default when an older config.json leaves it out + .head_dim = if (o.get("head_dim")) |v| (if (v == .null) 128 else try count(o, "head_dim")) else 128, + .inter = try count(o, "intermediate_size"), + .vocab = try count(o, "vocab_size"), + .eps = try number(o, "rms_norm_eps"), + .theta = try number(o, "rope_theta"), + .max_positions = try count(o, "max_position_embeddings"), + .tied = try flag(o, "tie_word_embeddings", false), + }; + if (try flag(o, "attention_bias", false)) return unsupported("attention_bias", "is true: the forward has no projection biases"); + if (try flag(o, "use_sliding_window", false)) return unsupported("use_sliding_window", "is true: the forward attends to the whole input"); + if (o.get("rope_scaling")) |r| if (r != .null) return unsupported("rope_scaling", "is set: the forward runs plain RoPE"); + if (s.head_dim > 256) return unsupported("head_dim", "is above 256, the attention kernel's limit"); + if (s.heads % s.kv_heads != 0) return unsupported("num_attention_heads", "is not a multiple of num_key_value_heads"); + const act = o.get("hidden_act") orelse return invalid("hidden_act", "is missing"); + if (act != .string or !std.mem.eql(u8, act.string, "silu")) return unsupported("hidden_act", "is not silu"); + if (!(s.eps > 0) or !(s.theta > 0)) return invalid(if (s.eps > 0) "rope_theta" else "rms_norm_eps", "is not positive"); + return s; +} + +fn invalid(key: []const u8, why: []const u8) error{InvalidConfig} { + std.log.warn("qwen3 config.json: {s} {s}", .{ key, why }); + return error.InvalidConfig; +} + +fn unsupported(key: []const u8, why: []const u8) error{UnsupportedConfig} { + std.log.warn("qwen3 config.json: {s} {s}", .{ key, why }); + return error.UnsupportedConfig; +} + +/// A positive integer key. +fn count(o: std.json.ObjectMap, key: []const u8) Error!usize { + const v = o.get(key) orelse return invalid(key, "is missing"); + if (v != .integer or v.integer < 1) return invalid(key, "is not a positive integer"); + return @intCast(v.integer); +} + +/// A number key, integer or float. +fn number(o: std.json.ObjectMap, key: []const u8) Error!f64 { + const v = o.get(key) orelse return invalid(key, "is missing"); + return switch (v) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => invalid(key, "is not a number"), + }; +} + +/// A boolean key, `default` when absent or null. +fn flag(o: std.json.ObjectMap, key: []const u8, default: bool) Error!bool { + const v = o.get(key) orelse return default; + return switch (v) { + .null => default, + .bool => |b| b, + else => invalid(key, "is not a boolean"), + }; +} + +/// Qwen3-Embedding-8B's config.json (the reranker's differs only in vocab_size 151669). +const embedding = + \\{"architectures": ["Qwen3ForCausalLM"], "attention_bias": false, "attention_dropout": 0.0, + \\ "bos_token_id": 151643, "eos_token_id": 151645, "head_dim": 128, "hidden_act": "silu", + \\ "hidden_size": 4096, "initializer_range": 0.02, "intermediate_size": 12288, + \\ "max_position_embeddings": 40960, "max_window_layers": 36, "model_type": "qwen3", + \\ "num_attention_heads": 32, "num_hidden_layers": 36, "num_key_value_heads": 8, "rms_norm_eps": 1e-06, + \\ "rope_scaling": null, "rope_theta": 1000000, "sliding_window": null, "tie_word_embeddings": false, + \\ "torch_dtype": "bfloat16", "transformers_version": "4.51.2", "use_cache": true, + \\ "use_sliding_window": false, "vocab_size": 151665} +; + +test "the 8B embedding and reranker configs give their geometry" { + const a = std.testing.allocator; + const s = try parse(a, embedding); + try std.testing.expectEqual(Spec{ .layers = 36, .hidden = 4096, .heads = 32, .kv_heads = 8, .head_dim = 128, .inter = 12288, .vocab = 151665, .eps = 1e-6, .theta = 1e6, .max_positions = 40960, .tied = false }, s); + const reranker = try std.mem.replaceOwned(u8, a, embedding, "\"vocab_size\": 151665", "\"vocab_size\": 151669"); + defer a.free(reranker); + try std.testing.expectEqual(@as(usize, 151669), (try parse(a, reranker)).vocab); + const older = try std.mem.replaceOwned(u8, a, embedding, "\"head_dim\": 128, ", ""); + defer a.free(older); + try std.testing.expect(!std.mem.eql(u8, older, embedding)); + try std.testing.expectEqual(@as(usize, 128), (try parse(a, older)).head_dim); +} + +test "each feature the forward does not run is refused by name, and each malformed key is invalid" { + const a = std.testing.allocator; + const Case = struct { []const u8, []const u8, Error }; + const cases = [_]Case{ + .{ "\"model_type\": \"qwen3\"", "\"model_type\": \"qwen3_moe\"", error.UnsupportedConfig }, + .{ "\"attention_bias\": false", "\"attention_bias\": true", error.UnsupportedConfig }, + .{ "\"use_sliding_window\": false", "\"use_sliding_window\": true", error.UnsupportedConfig }, + .{ "\"rope_scaling\": null", "\"rope_scaling\": {\"type\": \"yarn\", \"factor\": 4}", error.UnsupportedConfig }, + .{ "\"head_dim\": 128", "\"head_dim\": 512", error.UnsupportedConfig }, + .{ "\"hidden_act\": \"silu\"", "\"hidden_act\": \"gelu\"", error.UnsupportedConfig }, + .{ "\"num_key_value_heads\": 8", "\"num_key_value_heads\": 7", error.UnsupportedConfig }, + .{ "\"model_type\": \"qwen3\",", "", error.InvalidConfig }, + .{ "\"hidden_size\": 4096", "\"hidden_size\": 0", error.InvalidConfig }, + .{ "\"num_hidden_layers\": 36", "\"num_hidden_layers\": \"36\"", error.InvalidConfig }, + .{ "\"intermediate_size\": 12288,", "", error.InvalidConfig }, + .{ "\"rms_norm_eps\": 1e-06", "\"rms_norm_eps\": 0", error.InvalidConfig }, + .{ "\"rope_theta\": 1000000", "\"rope_theta\": \"1e6\"", error.InvalidConfig }, + .{ "\"tie_word_embeddings\": false", "\"tie_word_embeddings\": 0", error.InvalidConfig }, + .{ "\"hidden_act\": \"silu\",", "", error.InvalidConfig }, + }; + for (cases) |case| { + const bytes = try std.mem.replaceOwned(u8, a, embedding, case[0], case[1]); + defer a.free(bytes); + try std.testing.expect(!std.mem.eql(u8, bytes, embedding)); + try std.testing.expectError(case[2], parse(a, bytes)); + } + try std.testing.expectError(error.InvalidConfig, parse(a, "[1]")); + try std.testing.expectError(error.InvalidConfig, parse(a, "{not json")); +} From f1fc2b38bdcce88d8e17a8a526e85971c4eca207 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Mon, 2 Nov 2026 22:20:32 +0000 Subject: [PATCH 05/24] feat(hip): head-major attention rows for ungated attention Qwen3's attention has no output gate, so nothing turned the causal kernel's fp32 (heads, len, d) rows into o_proj's row-major input; tf_heads_rows does that transpose with one rounding to the output kind, reached through Ops.headsRows, with an exact kernel test, a kernel contract (copy and round, served Qwen3-8B shapes) and its compiled-usage record on all four RDNA chips. --- zig/kernels/hip/ops/attention.hpp | 10 +++ zig/src/core/contract/hip_records.zig | 9 +++ zig/src/hip/kernel_tests/ops.zig | 32 ++++++++++ zig/src/hip/launches.zig | 3 + zig/src/hip/launches/attention.zig | 11 ++++ zig/src/hip/ops/attention.zig | 5 ++ zig/src/hip/ops/ops.zig | 1 + zig/tests/hip/contract_heads_rows.zig | 90 +++++++++++++++++++++++++++ zig/tests/hip/contract_main.zig | 4 ++ 9 files changed, 165 insertions(+) create mode 100644 zig/tests/hip/contract_heads_rows.zig diff --git a/zig/kernels/hip/ops/attention.hpp b/zig/kernels/hip/ops/attention.hpp index c97bcf67f5..6e305c42d2 100644 --- a/zig/kernels/hip/ops/attention.hpp +++ b/zig/kernels/hip/ops/attention.hpp @@ -19,6 +19,16 @@ extern "C" __global__ void tf_attn_gate(const float* att, const void* qg, void* store(out, kind, i, a * sigmoid(g)); } +// out[r, h * d + j] = att[h, r, j] rounded once to kind: the causal kernel's head-major rows, row-major for o_proj. +extern "C" __global__ void tf_heads_rows(const float* att, void* out, int kind, int len, int heads, int d) { + long long i = gid(); + if (i >= static_cast(len) * heads * d) return; + int j = static_cast(i % d); + int h = static_cast((i / d) % heads); + int r = static_cast(i / (static_cast(d) * heads)); + store(out, kind, i, att[(static_cast(h) * len + r) * d + j]); +} + // keys or values (len, kv_heads, d) of kind into the cache (kv_heads, total, d) at slot pos0. extern "C" __global__ void tf_kv_write(const void* src, void* cache, int kind, int len, int kv_heads, int d, int total, int pos0) { diff --git a/zig/src/core/contract/hip_records.zig b/zig/src/core/contract/hip_records.zig index 899a403a9d..8784eafcd1 100644 --- a/zig/src/core/contract/hip_records.zig +++ b/zig/src/core/contract/hip_records.zig @@ -25,6 +25,15 @@ fn relu2Tile(comptime rows: u32, comptime cb: u32) limits.HipLaunch { /// HIP kernels as launched; tf-hip-kernel-contract limits DIR re-derives the usage records from the code objects. const own_records = [_]Record{ + .{ + .launch = .{ .name = "attention tf_heads_rows", .symbol = "tf_heads_rows", .image = "ops", .threads = 256 }, + .usage = .{ + .{ .vgprs = 17, .sgprs = 28, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 17, .sgprs = 28, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 17, .sgprs = 28, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 17, .sgprs = 30, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + }, + }, .{ .launch = .{ .name = "RMS norm tf_rms_rows", .symbol = "tf_rms_rows", .image = "ops", .threads = 256 }, .usage = .{ diff --git a/zig/src/hip/kernel_tests/ops.zig b/zig/src/hip/kernel_tests/ops.zig index 1c427945f3..c9358329a9 100644 --- a/zig/src/hip/kernel_tests/ops.zig +++ b/zig/src/hip/kernel_tests/ops.zig @@ -104,6 +104,38 @@ test "casts, sums, column copies and cache writes give torch's bits" { } } +test "head-major attention rows become row-major o_proj input, rounded once to the activation or kept in fp32" { + const t = try Rig.open(1 << 20); + defer t.close(); + const o = t.ops(&t.on); + var rng: rig.Rng = .{ .state = 0x4EAD }; + const len = 37; + const heads = 5; + const d = 24; + const n = len * heads * d; + const ha = try gpa.alloc(f32, n); + defer gpa.free(ha); + for (ha) |*x| x.* = @floatCast(rng.normal() * 9.1); + var att = try t.upload(ha); + defer att.free(); + var half = try t.alloc(n * 2); + defer half.free(); + var wide = try t.alloc(n * 4); + defer wide.free(); + try o.headsRows(at(att), tensor(half, act(t)), len, heads, d); + try o.headsRows(at(att), tensor(wide, .f32), len, heads, d); + try t.stream.synchronize(); + const gh = try rig.download(u16, half); + defer gpa.free(gh); + const gw = try rig.download(f32, wide); + defer gpa.free(gw); + for (0..len) |r| for (0..heads) |h| for (0..d) |j| { + const x = ha[(h * len + r) * d + j]; + try std.testing.expectEqual(round(t, x), gh[(r * heads + h) * d + j]); + try std.testing.expectEqual(x, gw[(r * heads + h) * d + j]); + }; +} + test "argmax and top-k follow torch's order: the larger value, then the lower index" { const t = try Rig.open(1 << 20); defer t.close(); diff --git a/zig/src/hip/launches.zig b/zig/src/hip/launches.zig index 54543599cc..06608d5731 100644 --- a/zig/src/hip/launches.zig +++ b/zig/src/hip/launches.zig @@ -132,6 +132,7 @@ pub const Launcher = struct { silu_mul: Function, add: Function, attn_gate: Function, + heads_rows: Function, gnorm_silu: Function, copy_cols: Function, rope_prefill: Function, @@ -206,6 +207,7 @@ pub const Launcher = struct { .silu_mul = try ops.function("tf_silu_mul"), .add = try ops.function("tf_add"), .attn_gate = try ops.function("tf_attn_gate"), + .heads_rows = try ops.function("tf_heads_rows"), .gnorm_silu = try ops.function("tf_gnorm_silu"), .copy_cols = try ops.function("tf_copy_cols"), .rope_prefill = try ops.function("tf_rope_prefill"), @@ -330,6 +332,7 @@ pub const Launcher = struct { pub const tf_moe_route = moe.tf_moe_route; pub const tf_causal = attention.tf_causal; pub const tf_attn_gate = attention.tf_attn_gate; + pub const tf_heads_rows = attention.tf_heads_rows; pub const tf_kv_write = attention.tf_kv_write; pub const tf_kv_write_at = attention.tf_kv_write_at; pub const tf_argmax_rows = draw.tf_argmax_rows; diff --git a/zig/src/hip/launches/attention.zig b/zig/src/hip/launches/attention.zig index ea31be1c9f..00e6ea71ba 100644 --- a/zig/src/hip/launches/attention.zig +++ b/zig/src/hip/launches/attention.zig @@ -134,6 +134,17 @@ pub fn tf_attn_gate(l: *const Launcher, att: CF, qg: C, out: P, kind: c_int, len try l.flat(l.op.attn_gate, @as(i64, len) * heads * d, s, &a); } +pub fn tf_heads_rows(l: *const Launcher, att: CF, out: P, kind: c_int, len: c_int, heads: c_int, d: c_int, s: S) Error!void { + var a: Args = .{}; + try a.add(ad(att)); + try a.add(ad(out)); + try a.add(kind); + try a.add(len); + try a.add(heads); + try a.add(d); + try l.flat(l.op.heads_rows, @as(i64, len) * heads * d, s, &a); +} + pub fn tf_kv_write(l: *const Launcher, src: C, cache: P, kind: c_int, len: c_int, kv_heads: c_int, d: c_int, total: c_int, pos0: c_int, s: S) Error!void { var a: Args = .{}; try a.add(ad(src)); diff --git a/zig/src/hip/ops/attention.zig b/zig/src/hip/ops/attention.zig index 7e356c4e4c..132e2d29e3 100644 --- a/zig/src/hip/ops/attention.zig +++ b/zig/src/hip/ops/attention.zig @@ -14,6 +14,11 @@ pub fn attnGate(o: Ops, att: u64, qg: Tensor, out: Tensor, len: usize, heads: us try o.l.tf_attn_gate(f(att), p(qg.ptr), p(out.ptr), @backingInt(out.kind), int(len), int(heads), int(d), @intFromBool(rows_major), o.stream); } +/// The causal kernel's fp32 (heads, len, d) rows as (len, heads * d) of `out`'s kind, rounded once: o_proj's input. +pub fn headsRows(o: Ops, att: u64, out: Tensor, len: usize, heads: usize, d: usize) Error!void { + try o.l.tf_heads_rows(f(att), p(out.ptr), @backingInt(out.kind), int(len), int(heads), int(d), o.stream); +} + /// Prefill RoPE of x (len, heads, d), rounded to x's kind, stored as `out` kind at out[h * s_head + r * s_row + j]. pub fn ropePrefill(o: Ops, x: Tensor, out: Tensor, s_head: usize, s_row: usize, len: usize, heads: usize, d: usize, rotary: usize, pos0: usize, theta: f32) Error!void { try o.l.tf_rope_prefill(p(x.ptr), @backingInt(x.kind), p(out.ptr), @backingInt(out.kind), @intCast(s_head), @intCast(s_row), int(len), int(heads), int(d), int(rotary), int(pos0), theta, o.stream); diff --git a/zig/src/hip/ops/ops.zig b/zig/src/hip/ops/ops.zig index a7c32470d1..98907749f1 100644 --- a/zig/src/hip/ops/ops.zig +++ b/zig/src/hip/ops/ops.zig @@ -56,6 +56,7 @@ pub const Ops = struct { } pub const attnGate = attention.attnGate; + pub const headsRows = attention.headsRows; pub const ropePrefill = attention.ropePrefill; pub const qkRope = attention.qkRope; pub const ropeDecode = attention.ropeDecode; diff --git a/zig/tests/hip/contract_heads_rows.zig b/zig/tests/hip/contract_heads_rows.zig new file mode 100644 index 0000000000..0f319f43cc --- /dev/null +++ b/zig/tests/hip/contract_heads_rows.zig @@ -0,0 +1,90 @@ +//! tf_heads_rows under the contract on HIP: the causal kernel's fp32 (heads, len, d) rows to row-major o_proj input. +const std = @import("std"); +const hip = @import("hip"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); +const kit = @import("contract_kit.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; + +fn none(_: c.Shape) usize { + return 0; +} + +pub fn spec(k: kit.Kind) c.Spec { + return .{ + .name = if (k == .bf16) "hip tf_heads_rows (bf16)" else "hip tf_heads_rows (fp16)", + .accum = .none, + .depth = none, + .out = kit.elem(k), + // rows: a prompt block (1, a mid one, a whole 2048); n: query heads; k: head width (Qwen3-8B: 32 of 128) + .served = &.{ .{ .rows = 1, .n = 32, .k = 128 }, .{ .rows = 300, .n = 32, .k = 128 }, .{ .rows = 2048, .n = 32, .k = 128 } }, + }; +} + +/// shape.n heads of shape.rows rows of shape.k fp32 values, each to (row, head, j) of the activation kind. +pub const HeadsRows = struct { + rig: rig_.Rig, + l: *const hip.Launcher, + k: kit.Kind, + arena: std.heap.ArenaAllocator, + shape: c.Shape = undefined, + att: []f64 = &.{}, + in: u64 = 0, + out: u64 = 0, + + pub fn init(a: std.mem.Allocator, rt: *const hip.Runtime, s: hip.Stream, l: *const hip.Launcher) HeadsRows { + return .{ .rig = .{ .a = a, .r = rt, .s = s }, .l = l, .k = kit.act(), .arena = .init(a) }; + } + + pub fn deinit(self: *HeadsRows) void { + self.arena.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *HeadsRows) g.Kernel { + return .{ .spec = spec(self.k), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(HeadsRows); + + pub fn prepare(self: *HeadsRows, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.k == 0) return false; + self.rig.reset(); + _ = self.arena.reset(.retain_capacity); + self.shape = shape; + const a = self.arena.allocator(); + var prng = std.Random.DefaultPrng.init(seed); + const n = shape.rows * shape.n * shape.k; + // fp32 attention outputs: spread past the activation's step so the rounding is exercised everywhere + self.att = try kit.floats(a, prng.random(), n, -40, 40); + if (fill == .nonfinite) for ([_]f64{ std.math.nan(f64), std.math.inf(f64), -std.math.inf(f64) }, 0..) |v, i| { + self.att[(i * 97 + 3) % n] = v; + }; + self.in = try self.rig.input(try kit.f32Bytes(a, self.att)); + self.out = try self.rig.output(n * kit.size(self.k)); + return true; + } + + pub fn run(self: *HeadsRows) !void { + const s = self.shape; + try self.l.tf_heads_rows(kit.ptr(?[*]const f32, self.in), kit.ptr(?*anyopaque, self.out), @backingInt(self.k), @intCast(s.rows), @intCast(s.n), @intCast(s.k), self.rig.s.handle); + try self.rig.s.synchronize(); + } + + pub fn count(self: *HeadsRows) usize { + return self.shape.rows * self.shape.n * self.shape.k; + } + + /// Output (r, h, j) is input (h, r, j), exactly; the bound adds the output's one rounding. + pub fn reference(self: *HeadsRows, want: []f64, mag: []f64) !void { + const s = self.shape; + for (0..s.rows) |r| for (0..s.n) |h| for (0..s.k) |j| { + const v = self.att[(h * s.rows + r) * s.k + j]; + want[(r * s.n + h) * s.k + j] = v; + mag[(r * s.n + h) * s.k + j] = @abs(v); + }; + } +}; diff --git a/zig/tests/hip/contract_main.zig b/zig/tests/hip/contract_main.zig index 32657d7a93..fce0236f1f 100644 --- a/zig/tests/hip/contract_main.zig +++ b/zig/tests/hip/contract_main.zig @@ -18,6 +18,7 @@ const relu2_k = @import("contract_relu2.zig"); const kv_k = @import("contract_kv_write.zig"); const argmax_k = @import("contract_argmax.zig"); const pair_k = @import("contract_relu2_pair.zig"); +const heads_rows_k = @import("contract_heads_rows.zig"); const gate_k = @import("contract_moe_gate.zig"); const router_k = @import("contract_router.zig"); const kit = @import("contract_kit.zig"); @@ -152,6 +153,9 @@ fn run(a: std.mem.Allocator, io: std.Io) !usize { var pair = pair_k.Pair.init(a, &r, s, &l); defer pair.deinit(); try kernels.append(a, pair.kernel()); + var heads_rows = heads_rows_k.HeadsRows.init(a, &r, s, &l); + defer heads_rows.deinit(); + try kernels.append(a, heads_rows.kernel()); var gates: [3]gate_k.Gate = undefined; for (&gates, std.enums.values(gate_k.Which)) |*m, w| m.* = .init(a, &r, s, &l, w); defer for (&gates) |*m| m.deinit(); From 2532bb87dee2e218ee32d161e0b78942d7ff0ee6 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Mon, 2 Nov 2026 22:33:34 +0000 Subject: [PATCH 06/24] feat(qwen3): load Qwen3 dense weights and run batched prompts The embedding and reranker checkpoints share one Qwen3 dense body, so the qwen3 family now uploads it as stored (BF16 projections, fp32 norms, an optional model. prefix), reads lm_head rows on the host for the reranker, and runs lastRows: inputs batched layer by layer, each attending only to itself from position 0, final-normed last rows out. It adds the module, host tests and an env-gated GPU check of both checkpoints against transformers. --- zig/build/hip.zig | 29 ++++ zig/src/core/root.zig | 1 + zig/src/families/qwen3/forward.zig | 134 +++++++++++++++ zig/src/families/qwen3/forward_test.zig | 211 ++++++++++++++++++++++++ zig/src/families/qwen3/model.zig | 209 +++++++++++++++++++++++ zig/src/families/qwen3/qwen3.zig | 15 ++ 6 files changed, 599 insertions(+) create mode 100644 zig/src/families/qwen3/forward.zig create mode 100644 zig/src/families/qwen3/forward_test.zig create mode 100644 zig/src/families/qwen3/model.zig create mode 100644 zig/src/families/qwen3/qwen3.zig diff --git a/zig/build/hip.zig b/zig/build/hip.zig index 22cb6d002a..1fb6ae2eaf 100644 --- a/zig/build/hip.zig +++ b/zig/build/hip.zig @@ -141,6 +141,9 @@ pub fn steps(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.buil test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = host.flashnext })).step); b.step("test-flashnext-proposal-hip-host", "Original Flash Next HIP IDs and head-row spans without device runtime").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = host.flashnext, .filters = &.{"Flash Next proposal:"} })).step); test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = host.nemotron })).step); + const qwen3_host = b.addRunArtifact(b.addTest(.{ .root_module = host.qwen3 })); + test_step.dependOn(&qwen3_host.step); + b.step("test-qwen3-hip-host", "Qwen3 dense config, checkpoint prefix and lm_head rows on the CPU").dependOn(&qwen3_host.step); test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = host.serve })).step); test_step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = host.native })).step); // the GPU programs compile on any host against stub kernels, so `test` catches what only they reach @@ -150,6 +153,17 @@ pub fn steps(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.buil const check = engineModules(b, check_target, .Debug, stub_kernels, null); b.step("check-native-hip-host", "Compile the real native HIP server with inert kernels, without GPU execution").dependOn(&server(b, check_target, check, build_options).step); const gpu = engineModules(b, target, optimize, objects, hip_include); + const qwen3_gpu = b.addTest(.{ .root_module = b.createModule(.{ + .root_source_file = b.path("zig/src/families/qwen3/forward_test.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ .{ .name = "hip", .module = gpu.hip }, .{ .name = "qwen3", .module = gpu.qwen3 } }, + }) }); + const qwen3_run = b.addRunArtifact(qwen3_gpu); + qwen3_run.has_side_effects = true; // the checkpoints and oracle files named by TF_QWEN3_* are not cached inputs + b.step("qwen3-hip-gpu-build", "Compile the Qwen3 dense GPU checks without running them").dependOn(&qwen3_gpu.step); + b.step("qwen3-hip-gpu-test", "Qwen3 lastRows on the GPU against transformers (TF_QWEN3_EMBED_*, TF_QWEN3_RERANK_*)").dependOn(&qwen3_run.step); b.step("tf-qwen35-test", "Qwen3.5 / 3.6 model checks on the GPU (-Dhip-arch)").dependOn(&b.addInstallArtifact(checksExe(b, target, optimize, gpu), .{}).step); const fx_test = b.addExecutable(.{ .name = "tf-flashnext-hip-test", .root_module = b.createModule(.{ .root_source_file = b.path("zig/tests/flashnext_hip/main.zig"), @@ -212,6 +226,7 @@ pub const Engine = struct { family: *std.Build.Module, flashnext: *std.Build.Module, nemotron: *std.Build.Module, + qwen3: *std.Build.Module, native: *std.Build.Module, }; @@ -280,6 +295,18 @@ pub fn engineModules(b: *std.Build, target: std.Build.ResolvedTarget, optimize: }); const ocr_kernels = b.createModule(.{ .root_source_file = b.addWriteFiles().add("glm_ocr_kernels.zig", "const hip = @import(\"hip\"); pub const arch = hip.kernels.arch; pub const bytes = hip.kernels.images[@intFromEnum(hip.kernels.Group.glm_ocr)];\n"), .imports = &.{.{ .name = "hip", .module = hip }} }); const glm_ocr = b.createModule(.{ .root_source_file = b.path("zig/src/families/glm_ocr/hip.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "hip", .module = hip }, .{ .name = "core", .module = core }, .{ .name = "lanes", .module = lanes }, .{ .name = "engine_api", .module = api }, .{ .name = "hip_serve", .module = serve }, .{ .name = "glm_ocr_kernels", .module = ocr_kernels } } }); + const qwen3 = b.createModule(.{ + .root_source_file = b.path("zig/src/families/qwen3/qwen3.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{ + .{ .name = "hip", .module = hip }, + .{ .name = "core", .module = core }, + .{ .name = "lanes", .module = lanes }, + .{ .name = "engine_api", .module = api }, + }, + }); const native = b.createModule(.{ .root_source_file = b.path("zig/src/native/hip.zig"), .target = target, @@ -294,6 +321,7 @@ pub fn engineModules(b: *std.Build, target: std.Build.ResolvedTarget, optimize: .{ .name = "nemotron", .module = nemotron }, .{ .name = "glm_ocr", .module = glm_ocr }, .{ .name = "hip_serve", .module = serve }, + .{ .name = "qwen3", .module = qwen3 }, }, }); return .{ @@ -304,6 +332,7 @@ pub fn engineModules(b: *std.Build, target: std.Build.ResolvedTarget, optimize: .family = family, .flashnext = flashnext, .nemotron = nemotron, + .qwen3 = qwen3, .native = native, }; } diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 9e537a0e20..5fd3441d3e 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -12,6 +12,7 @@ pub const eaq_container = @import("eaq_container.zig"); pub const digest_cache = @import("digest_cache.zig"); pub const checkpoint_identity = @import("checkpoint_identity.zig"); pub const checkpoint = @import("checkpoint.zig"); +pub const checkpoint_fixture = @import("checkpoint_fixture.zig"); pub const exl3_load = @import("exl3_load.zig"); pub const direct_io = @import("direct_io.zig"); pub const direct_gather = @import("direct_gather.zig"); diff --git a/zig/src/families/qwen3/forward.zig b/zig/src/families/qwen3/forward.zig new file mode 100644 index 0000000000..8e62817640 --- /dev/null +++ b/zig/src/families/qwen3/forward.zig @@ -0,0 +1,134 @@ +//! A batch of independent prompts through every layer; each input's final-normed last row comes out. +const std = @import("std"); +const hip = @import("hip"); +const model = @import("model.zig"); + +const Ops = hip.ops.Ops; +const Tensor = hip.ops.Tensor; +const Kind = hip.ops.Kind; + +/// Rows of one block of row-wise work, and of one attention query block. +pub const SPAN = 2048; + +pub const Error = hip.ops.Error; + +fn take(o: Ops, kind: Kind, n: usize) Error!Tensor { + return .{ .ptr = try o.arena.take(n * kind.size()), .kind = kind }; +} + +fn at(t: Tensor, n: usize) Tensor { + return .{ .ptr = t.ptr + n * t.kind.size(), .kind = t.kind }; +} + +/// head_dim ** -0.5 in double, as Python has it, passed as fp32. +pub fn scaleOf(head_dim: usize) f32 { + return @floatCast(std.math.pow(f64, @floatFromInt(head_dim), -0.5)); +} + +/// Each input's final-normed last row into `out` (ends.len, hidden); inputs end at `ends`, start at position 0. +pub fn lastRows(ops: Ops, m: *const model.Model, ids: u64, ends: []const u32, out: Tensor) Error!void { + var o = ops; + o.prefill = true; // prefill's kernels at any row count: an input's rows have the same bits alone or batched + const s = m.spec; + const total: usize = if (ends.len == 0) return else ends[ends.len - 1]; + var longest: usize = 0; + var begin: usize = 0; + for (ends) |e| { + if (e <= begin) return error.BadShape; + longest = @max(longest, e - begin); + begin = e; + } + const x = try take(o, m.act, total * s.hidden); + var start: usize = 0; + while (start < total) : (start += SPAN) { + const rows = @min(SPAN, total - start); + try o.embedRows(m.embed, ids + start * 4, rows, at(x, start * s.hidden)); + } + for (m.layers) |*layer| { + const mark = o.arena.mark(); + defer o.arena.release(mark); + try layerRows(o, m, layer, x, ends, total, longest); + } + for (ends, 0..) |e, i| try o.rms(at(x, (e - 1) * s.hidden), m.final_norm, at(out, i * s.hidden), 1, s.hidden, @floatCast(s.eps)); +} + +/// One decoder layer over every row: norms, projections and MLP in SPAN blocks, attention input by input. +fn layerRows(o: Ops, m: *const model.Model, l: *const model.Layer, x: Tensor, ends: []const u32, total: usize, longest: usize) Error!void { + const s = m.spec; + const hd = s.head_dim; + const qw = s.heads * hd; + const kvw = s.kv_heads * hd; + const q = try take(o, m.act, total * qw); + const k = try take(o, m.act, total * kvw); + const v = try take(o, m.act, total * kvw); + var start: usize = 0; + while (start < total) : (start += SPAN) { + const rows = @min(SPAN, total - start); + const mark = o.arena.mark(); + defer o.arena.release(mark); + const normed = try take(o, m.act, rows * s.hidden); + try o.rms(at(x, start * s.hidden), l.input_norm, normed, rows, s.hidden, @floatCast(s.eps)); + try copyRows(o, try o.project(normed, l.q, rows, false), at(q, start * qw), rows * qw); + try copyRows(o, try o.project(normed, l.k, rows, false), at(k, start * kvw), rows * kvw); + try copyRows(o, try o.project(normed, l.v, rows, false), at(v, start * kvw), rows * kvw); + } + const attn = try take(o, m.act, total * qw); + // one input's flat cache, reused by each input in turn: (kv_heads, longest, d) + const cache: Ops.Cache = .{ .k = try o.arena.take(kvw * longest * m.act.size()), .v = try o.arena.take(kvw * longest * m.act.size()), .kind = m.act, .kv_heads = s.kv_heads, .total = longest, .d = hd }; + var a: usize = 0; + for (ends) |b| { + defer a = b; + try attention(o, m, l, q, k, v, attn, cache, a, b); + } + start = 0; + while (start < total) : (start += SPAN) { + const rows = @min(SPAN, total - start); + const mark = o.arena.mark(); + defer o.arena.release(mark); + const xs = at(x, start * s.hidden); + const y = try o.project(at(attn, start * qw), l.o, rows, false); + try o.add(xs, y, xs, rows * s.hidden); + const normed = try take(o, m.act, rows * s.hidden); + try o.rms(xs, l.post_norm, normed, rows, s.hidden, @floatCast(s.eps)); + const gate = try o.project(normed, l.gate, rows, false); + const up = try o.project(normed, l.up, rows, false); + const act = try take(o, m.act, rows * s.inter); + try o.siluMul(gate, up, act, rows * s.inter); + const z = try o.project(act, l.down, rows, false); + try o.add(xs, z, xs, rows * s.hidden); + } +} + +/// Input rows [a, b): head norms, RoPE from position 0, its cache, causal attention, head-major rows back. +fn attention(o: Ops, m: *const model.Model, l: *const model.Layer, q: Tensor, k: Tensor, v: Tensor, attn: Tensor, c: Ops.Cache, a: usize, b: usize) Error!void { + const s = m.spec; + const hd = s.head_dim; + const qw = s.heads * hd; + const kvw = s.kv_heads * hd; + const theta: f32 = @floatCast(s.theta); + const len = b - a; + var p: usize = 0; + while (p < len) : (p += SPAN) { + const rows = @min(SPAN, len - p); + const mark = o.arena.mark(); + defer o.arena.release(mark); + const qn = try take(o, m.act, rows * qw); + try o.rms(at(q, (a + p) * qw), l.q_norm, qn, rows * s.heads, hd, @floatCast(s.eps)); + const kn = try take(o, m.act, rows * kvw); + try o.rms(at(k, (a + p) * kvw), l.k_norm, kn, rows * s.kv_heads, hd, @floatCast(s.eps)); + // q rotated, rounded to the activation dtype, then widened: (heads, rows, d) fp32 for the kernel + const q32 = try o.arena.of(f32, rows * qw); + try o.ropePrefill(qn, .{ .ptr = q32, .kind = .f32 }, rows * hd, hd, rows, s.heads, hd, hd, p, theta); + // k rotated into the cache (kv_heads, total, d) from slot p (the kernel writes row r at r * s_row); v as it is + try o.ropePrefill(kn, .{ .ptr = c.k + p * hd * c.kind.size(), .kind = c.kind }, c.total * hd, hd, rows, s.kv_heads, hd, hd, p, theta); + try o.kvWrite(at(v, (a + p) * kvw), c.v, rows, s.kv_heads, hd, c.total, p); + const att = try o.arena.of(f32, rows * qw); + try o.causalPrefill(q32, c, att, rows, p + rows, s.heads, scaleOf(hd), p); + try o.headsRows(att, at(attn, (a + p) * qw), rows, s.heads, hd); + } +} + +/// `n` values of `src` into `dst` (same kind), as a one-row column copy. +fn copyRows(o: Ops, src: Tensor, dst: Tensor, n: usize) Error!void { + try o.copyCols(src, n, 0, dst.ptr, 1, n); +} diff --git a/zig/src/families/qwen3/forward_test.zig b/zig/src/families/qwen3/forward_test.zig new file mode 100644 index 0000000000..3cc2af8584 --- /dev/null +++ b/zig/src/families/qwen3/forward_test.zig @@ -0,0 +1,211 @@ +//! lastRows on the GPU against transformers: the embedding's pooled rows and the reranker's yes/no logits. + +// Skipped without TF_QWEN3_EMBED_DIR / _REF / _BAR (fp32, BF16 oracles) or TF_QWEN3_RERANK_DIR / _REF. +const std = @import("std"); +const hip = @import("hip"); +const qwen3 = @import("qwen3"); + +const t = std.testing; +const gpa = t.allocator; + +/// The device, its kernels and an arena: what the engine holds around `lastRows`. +const Gpu = struct { + r: hip.Runtime, + ctx: hip.Context, + stream: hip.Stream, + l: hip.Launcher, + arena: hip.Arena, + + fn open(g: *Gpu, arena_bytes: usize) !void { + g.r = try hip.Runtime.open(); + errdefer g.r.close(); + g.ctx = try hip.Context.init(&g.r, 0); + errdefer g.ctx.deinit(); + g.stream = try hip.Stream.init(&g.r); + errdefer g.stream.deinit(); + g.l = try hip.Launcher.load(&g.r, .{}, hip.kernels.images); + errdefer g.l.unload(); + g.arena = try hip.Arena.init(&g.r, arena_bytes); + } + + fn close(g: *Gpu) void { + g.arena.deinit(); + g.l.unload(); + g.stream.deinit(); + g.ctx.deinit(); + g.r.close(); + } + + fn ops(g: *Gpu) hip.ops.Ops { + return .{ .l = &g.l, .bf16 = true, .stream = g.stream.handle, .arena = &g.arena }; + } + + /// Each input's final-normed last row, as BF16 bits (`ends.len` rows of `hidden`). + fn lastRows(g: *Gpu, m: *const qwen3.Model, ids: []const i32, ends: []const u32) ![]u16 { + const mark = g.arena.mark(); + defer g.arena.release(mark); + var dev = try hip.DeviceBuffer.fromHost(&g.r, std.mem.sliceAsBytes(ids)); + defer dev.free(); + const hidden = m.spec.hidden; + var out = try hip.DeviceBuffer.alloc(&g.r, ends.len * hidden * 2); + defer out.free(); + try qwen3.lastRows(g.ops(), m, @intFromPtr(dev.ptr), ends, .{ .ptr = @intFromPtr(out.ptr), .kind = .bf16 }); + try g.stream.synchronize(); + const rows = try gpa.alloc(u16, ends.len * hidden); + errdefer gpa.free(rows); + try out.download(0, std.mem.sliceAsBytes(rows)); + return rows; + } +}; + +fn env(a: std.mem.Allocator, key: []const u8) !?[]u8 { + return t.environ.getAlloc(a, key) catch |err| switch (err) { + error.EnvironmentVariableMissing => null, + else => err, + }; +} + +fn widen(bits: u16) f64 { + return @as(f32, @bitCast(@as(u32, bits) << 16)); +} + +fn readJson(a: std.mem.Allocator, path: []const u8) !std.json.Value { + const bytes = try std.Io.Dir.cwd().readFileAlloc(t.io, path, a, .limited(256 << 20)); + return std.json.parseFromSliceLeaky(std.json.Value, a, bytes, .{}); +} + +fn idsOf(a: std.mem.Allocator, v: std.json.Value) ![]i32 { + const out = try a.alloc(i32, v.array.items.len); + for (v.array.items, out) |x, *o| o.* = @intCast(x.integer); + return out; +} + +fn number(v: std.json.Value) f64 { + return switch (v) { + .integer => |i| @floatFromInt(i), + .float => |f| f, + else => std.math.nan(f64), + }; +} + +fn cosine(x: []const f64, y: []const f64) f64 { + var xy: f64 = 0; + var xx: f64 = 0; + var yy: f64 = 0; + for (x, y) |p, q| { + xy += p * q; + xx += p * p; + yy += q * q; + } + return xy / @sqrt(xx * yy); +} + +test "Qwen3-Embedding: the card's inputs batched equal them alone, and sit within the BF16 bar of the fp32 vectors" { + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const dir = try env(a, "TF_QWEN3_EMBED_DIR") orelse return error.SkipZigTest; + const ref = try readJson(a, try env(a, "TF_QWEN3_EMBED_REF") orelse return error.SkipZigTest); + const bar = try readJson(a, try env(a, "TF_QWEN3_EMBED_BAR") orelse return error.SkipZigTest); + var g: Gpu = undefined; + try g.open(4 << 30); + defer g.close(); + // the card's 2 queries and 2 documents in both references, each vector full width, before any GPU work + const inputs = ref.object.get("ids").?.array.items; + const refs = ref.object.get("vectors").?.array.items; + const bars = bar.object.get("vectors").?.array.items; + try t.expectEqual(@as(usize, 4), inputs.len); + try t.expectEqual(@as(usize, 4), refs.len); + try t.expectEqual(@as(usize, 4), bars.len); + for (inputs) |one| try t.expect(one.array.items.len > 0); + var m = try qwen3.Model.load(gpa, t.io, &g.r, dir); + defer m.deinit(); + for (refs, bars) |r, b| { + try t.expectEqual(m.spec.hidden, r.array.items.len); + try t.expectEqual(m.spec.hidden, b.array.items.len); + } + var ids: std.ArrayList(i32) = .empty; + const ends = try a.alloc(u32, inputs.len); + for (inputs, ends) |one, *e| { + try ids.appendSlice(a, try idsOf(a, one)); + e.* = @intCast(ids.items.len); + } + const batch = try g.lastRows(&m, ids.items, ends); + defer gpa.free(batch); + const hidden = m.spec.hidden; + var bar_worst: f64 = 0; + for (bars, refs) |b, r| { + const bv = try a.alloc(f64, hidden); + const rv = try a.alloc(f64, hidden); + for (b.array.items, r.array.items, bv, rv) |x, y, *p, *q| { + p.* = number(x); + q.* = number(y); + } + bar_worst = @max(bar_worst, 1 - cosine(bv, rv)); + } + const limit = @min(0.001, 2 * bar_worst); + var begin: u32 = 0; + for (ends, 0..) |e, i| { + const alone = try g.lastRows(&m, ids.items[begin..e], &.{e - begin}); + defer gpa.free(alone); + try t.expectEqualSlices(u16, batch[i * hidden ..][0..hidden], alone); + const got = try a.alloc(f64, hidden); + const want = try a.alloc(f64, hidden); + for (batch[i * hidden ..][0..hidden], refs[i].array.items, got, want) |bits, r, *p, *q| { + p.* = widen(bits); + q.* = number(r); + } + const distance = 1 - cosine(got, want); + std.debug.print("embedding input {d}: cosine distance {e:.3} (limit {e:.3})\n", .{ i, distance, limit }); + try t.expect(distance <= limit); + begin = e; + } +} + +test "Qwen3-Reranker: each case's yes - no logit from lastRows and lm_head rows matches transformers" { + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const dir = try env(a, "TF_QWEN3_RERANK_DIR") orelse return error.SkipZigTest; + const ref = try readJson(a, try env(a, "TF_QWEN3_RERANK_REF") orelse return error.SkipZigTest); + var g: Gpu = undefined; + try g.open(8 << 30); + defer g.close(); + var m = try qwen3.Model.load(gpa, t.io, &g.r, dir); + defer m.deinit(); + const yes: u32 = @intCast(ref.object.get("yes").?.integer); + const no: u32 = @intCast(ref.object.get("no").?.integer); + const head = try qwen3.readRows(gpa, t.io, dir, "lm_head.weight", &.{ yes, no }); + defer gpa.free(head); + const hidden = m.spec.hidden; + var worst_diff: f64 = 0; + var worst_score: f64 = 0; + var cases: usize = 0; + for (ref.object.get("cases").?.array.items) |c| { + const want_diff = c.object.get("diff") orelse continue; // ids-only cases carry no score + const ids = try idsOf(a, c.object.get("ids").?); + const row = try g.lastRows(&m, ids, &.{@intCast(ids.len)}); + defer gpa.free(row); + var ly: f64 = 0; + var ln: f64 = 0; + for (row, head[0..hidden], head[hidden..][0..hidden]) |bits, wy, wn| { + ly += widen(bits) * wy; + ln += widen(bits) * wn; + } + const diff = ly - ln; + const score = 1 / (1 + @exp(-diff)); + const ref_score = number(c.object.get("score").?); + // a NaN would slip through @max: every value must be finite before it counts + for ([_]f64{ diff, score, number(want_diff), ref_score }) |v| try t.expect(std.math.isFinite(v)); + const d = @abs(diff - number(want_diff)); + const ds = @abs(score - ref_score); + std.debug.print("rerank {s} {d} tokens: diff {d:.4} ref {d:.4} score {d:.6} ref {d:.6}\n", .{ c.object.get("kind").?.string, ids.len, diff, number(want_diff), score, number(c.object.get("score").?) }); + worst_diff = @max(worst_diff, d); + worst_score = @max(worst_score, ds); + cases += 1; + } + std.debug.print("rerank: {d} cases, worst |diff| {d:.4} (gate 0.25), worst |score| {d:.5} (gate 0.02)\n", .{ cases, worst_diff, worst_score }); + try t.expect(cases >= 32); + try t.expect(worst_diff <= 0.25); + try t.expect(worst_score <= 0.02); +} diff --git a/zig/src/families/qwen3/model.zig b/zig/src/families/qwen3/model.zig new file mode 100644 index 0000000000..c57b620fa1 --- /dev/null +++ b/zig/src/families/qwen3/model.zig @@ -0,0 +1,209 @@ +//! A Qwen3 dense checkpoint on the GPU: BF16 projections as stored, fp32 norms, and host rows of lm_head on request. +const std = @import("std"); +const core = @import("core"); +const hip = @import("hip"); +const config = @import("config.zig"); + +const quant = core.quant; +const Checkpoint = core.checkpoint.Checkpoint; +const Allocator = std.mem.Allocator; + +pub const Projection = quant.Projection; + +/// One decoder layer: norms are fp32 device addresses, projections BF16 (N, K) as the checkpoint stores them. +pub const Layer = struct { + input_norm: u64, + post_norm: u64, + q_norm: u64, + k_norm: u64, + q: Projection, + k: Projection, + v: Projection, + o: Projection, + gate: Projection, + up: Projection, + down: Projection, +}; + +pub const Model = struct { + gpa: Allocator, + spec: config.Spec, + /// The activation dtype: BF16, as the checkpoint's weights. + act: hip.ops.Kind = .bf16, + embed: Projection, + final_norm: u64, + layers: []Layer, + buffers: std.ArrayList(hip.memory.DeviceBuffer) = .empty, + + /// The checkpoint in `dir` on the device of `rt`; its mapped pages are dropped layer by layer. + pub fn load(gpa: Allocator, io: std.Io, rt: *const hip.Runtime, dir: []const u8) !Model { + const spec = try readSpec(gpa, io, dir); + var ck = try Checkpoint.openModel(gpa, io, dir); + defer ck.close(); + ck.scope = scopeOf(&ck); + var m: Model = .{ .gpa = gpa, .spec = spec, .embed = undefined, .final_norm = 0, .layers = &.{} }; + errdefer m.deinit(); + const up: hip.Upload = .{ .r = rt, .buffers = &m.buffers, .gpa = gpa }; + var scratch: std.heap.ArenaAllocator = .init(gpa); + defer scratch.deinit(); + var t: Table = .{ .ck = &ck }; + m.embed = try projection(scratch.allocator(), &t, &up, "embed_tokens", spec.vocab, spec.hidden); + m.final_norm = try norm(scratch.allocator(), &t, &up, "norm.weight", spec.hidden); + ck.release(); + m.layers = try gpa.alloc(Layer, spec.layers); + for (m.layers, 0..) |*layer, i| { + _ = scratch.reset(.retain_capacity); + layer.* = try readLayer(scratch.allocator(), &t, &up, spec, i); + ck.release(); + } + return m; + } + + pub fn deinit(m: *Model) void { + for (m.buffers.items) |*b| b.free(); + m.buffers.deinit(m.gpa); + m.gpa.free(m.layers); + m.* = undefined; + } +}; + +/// The reranker's save keeps the base model under `model.` (beside `lm_head`); the embedding's has no prefix. +pub fn scopeOf(ck: *const Checkpoint) []const u8 { + return if (ck.has("model.embed_tokens.weight")) "model." else ""; +} + +/// `quant.read`'s table over a checkpoint: dense tensors only, so no format asks for an affine width. +const Table = struct { + ck: *Checkpoint, + + pub fn get(t: *const Table, name: []const u8) !quant.Tensor { + return t.ck.get(name); + } + + pub fn has(t: *const Table, name: []const u8) bool { + return t.ck.has(name); + } + + pub fn width(_: *const Table, _: []const u8) error{UnsupportedQuantization}!quant.mlx.Width { + return error.UnsupportedQuantization; + } +}; + +fn readSpec(gpa: Allocator, io: std.Io, dir: []const u8) !config.Spec { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 20)); + defer gpa.free(bytes); + return config.parse(gpa, bytes); +} + +/// A BF16 projection `key` of shape (n, k), refused in any other dtype or shape. +fn projection(a: Allocator, t: *const Table, up: *const hip.Upload, key: []const u8, n: usize, k: usize) !Projection { + const host = try quant.readStored(a, t, key); + const w = switch (host) { + .dense => |d| d.weight, + else => return refuse(key, "is not a dense tensor"), + }; + if (w.dtype != .bf16 or w.rank != 2 or w.shape[0] != n or w.shape[1] != k) { + std.log.err("{s}: {t} {any}, expected bf16 [{d}, {d}]", .{ key, w.dtype, w.shape[0..w.rank], n, k }); + return error.UnexpectedTensor; + } + return quant.view(try quant.upload(up.uploader(), host), false); +} + +/// An RMS norm weight `key` of `n` values, widened to fp32 on the device. +fn norm(a: Allocator, t: *const Table, up: *const hip.Upload, key: []const u8, n: usize) !u64 { + const src = try t.get(key); + if (src.rank != 1 or src.shape[0] != n) { + std.log.err("{s}: shape {any}, expected [{d}]", .{ key, src.shape[0..src.rank], n }); + return error.UnexpectedTensor; + } + const wide = try quant.convert.float32(a, src); + return (try up.uploader().tensor(wide)).ptr; +} + +fn readLayer(a: Allocator, t: *const Table, up: *const hip.Upload, s: config.Spec, i: usize) !Layer { + var buf: [128]u8 = undefined; + const q_width = s.heads * s.head_dim; + const kv_width = s.kv_heads * s.head_dim; + return .{ + .input_norm = try norm(a, t, up, try layerKey(&buf, i, "input_layernorm.weight"), s.hidden), + .post_norm = try norm(a, t, up, try layerKey(&buf, i, "post_attention_layernorm.weight"), s.hidden), + .q_norm = try norm(a, t, up, try layerKey(&buf, i, "self_attn.q_norm.weight"), s.head_dim), + .k_norm = try norm(a, t, up, try layerKey(&buf, i, "self_attn.k_norm.weight"), s.head_dim), + .q = try projection(a, t, up, try layerKey(&buf, i, "self_attn.q_proj"), q_width, s.hidden), + .k = try projection(a, t, up, try layerKey(&buf, i, "self_attn.k_proj"), kv_width, s.hidden), + .v = try projection(a, t, up, try layerKey(&buf, i, "self_attn.v_proj"), kv_width, s.hidden), + .o = try projection(a, t, up, try layerKey(&buf, i, "self_attn.o_proj"), s.hidden, q_width), + .gate = try projection(a, t, up, try layerKey(&buf, i, "mlp.gate_proj"), s.inter, s.hidden), + .up = try projection(a, t, up, try layerKey(&buf, i, "mlp.up_proj"), s.inter, s.hidden), + .down = try projection(a, t, up, try layerKey(&buf, i, "mlp.down_proj"), s.hidden, s.inter), + }; +} + +fn layerKey(buf: *[128]u8, i: usize, rest: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "layers.{d}.{s}", .{ i, rest }); +} + +fn refuse(key: []const u8, why: []const u8) error{UnexpectedTensor} { + std.log.err("{s} {s}", .{ key, why }); + return error.UnexpectedTensor; +} + +/// Rows `rows` of the 2-D tensor `key` (named as stored, no prefix) as fp32 host values, row after row. +pub fn readRows(gpa: Allocator, io: std.Io, dir: []const u8, key: []const u8, rows: []const u32) ![]f32 { + var ck = try Checkpoint.openModel(gpa, io, dir); + defer ck.close(); + const t = try ck.get(key); + if (t.rank != 2) return refuse(key, "is not a matrix"); + const width = t.shape[1]; + const out = try gpa.alloc(f32, rows.len * width); + errdefer gpa.free(out); + const size = t.dtype.size(); + for (rows, 0..) |r, i| { + if (r >= t.shape[0]) { + std.log.warn("{s} has no row {d} of {d}", .{ key, r, t.shape[0] }); + return error.UnexpectedTensor; + } + var shape = t.shape; + shape[0] = width; + const one: quant.Tensor = .{ .bytes = t.bytes[r * width * size ..][0 .. width * size], .dtype = t.dtype, .rank = 1, .shape = shape }; + var scratch: std.heap.ArenaAllocator = .init(gpa); + defer scratch.deinit(); + const wide = try quant.convert.float32(scratch.allocator(), one); + // little-endian fp32 words, read one by one: a mapped fp32 tensor need not be 4-byte aligned + for (out[i * width ..][0..width], 0..) |*x, j| x.* = @bitCast(std.mem.readInt(u32, wide.bytes[4 * j ..][0..4], .little)); + } + return out; +} + +fn writeFixture(dir: std.Io.Dir, prefixed: bool) !void { + const fixture = core.checkpoint_fixture; + const a = std.testing.allocator; + // lm_head: 3 rows of 2 bf16 values: row r is (r + 1, -(r + 1)) as 1.0, 2.0, 3.0 and their negatives + const head = [_]u16{ 0x3f80, 0xbf80, 0x4000, 0xc000, 0x4040, 0xc040 }; + const embed = [_]u16{ 0, 0, 0, 0, 0, 0 }; + try fixture.write(a, std.testing.io, dir, "model.safetensors", &.{ + .{ .name = if (prefixed) "model.embed_tokens.weight" else "embed_tokens.weight", .dtype = "BF16", .shape = &.{ 3, 2 }, .bytes = std.mem.sliceAsBytes(&embed) }, + .{ .name = "lm_head.weight", .dtype = "BF16", .shape = &.{ 3, 2 }, .bytes = std.mem.sliceAsBytes(&head) }, + }); +} + +test "the reranker's model. prefix is found, the embedding's bare names are not, and lm_head rows read as fp32" { + const a = std.testing.allocator; + const io = std.testing.io; + for ([_]bool{ true, false }) |prefixed| { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try writeFixture(tmp.dir, prefixed); + const dir = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}", .{tmp.sub_path}); + defer a.free(dir); + var ck = try Checkpoint.openModel(a, io, dir); + defer ck.close(); + try std.testing.expectEqualStrings(if (prefixed) "model." else "", scopeOf(&ck)); + const rows = try readRows(a, io, dir, "lm_head.weight", &.{ 2, 0 }); + defer a.free(rows); + try std.testing.expectEqualSlices(f32, &.{ 3, -3, 1, -1 }, rows); + try std.testing.expectError(error.UnexpectedTensor, readRows(a, io, dir, "lm_head.weight", &.{3})); + } +} diff --git a/zig/src/families/qwen3/qwen3.zig b/zig/src/families/qwen3/qwen3.zig new file mode 100644 index 0000000000..68c56adede --- /dev/null +++ b/zig/src/families/qwen3/qwen3.zig @@ -0,0 +1,15 @@ +//! Qwen3 dense checkpoints on the native HIP engine: config, BF16 weights on the device and the batched prompt pass. + +pub const config = @import("config.zig"); +pub const model = @import("model.zig"); +pub const forward = @import("forward.zig"); +pub const Spec = config.Spec; +pub const Model = model.Model; +pub const readRows = model.readRows; +pub const lastRows = forward.lastRows; + +test { + _ = config; + _ = model; + _ = forward; +} From a7ea69737bdb6cb92cc14fdffc384184e3d966b2 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Mon, 2 Nov 2026 22:55:44 +0000 Subject: [PATCH 07/24] feat(embed): accept any MRL width from a declared least one Qwen3-Embedding-8B is trained for any output width from 32 to 4096, and the embedding route only took widths from a declared list, so a request for 384 or 2049 would be refused. The lane's embedding Info gains mrl_min: with it set, every width from that one to the full dimension crops; the list rule and the refusal without either stay. Tests cover both bounds, widths between them and the old rules. --- zig/src/core/lanes/embedding.zig | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/zig/src/core/lanes/embedding.zig b/zig/src/core/lanes/embedding.zig index 0490730668..ef4573e41a 100644 --- a/zig/src/core/lanes/embedding.zig +++ b/zig/src/core/lanes/embedding.zig @@ -7,12 +7,15 @@ pub const Info = struct { max_batch_tokens: u32, max_inputs: u32 = 128, mrl_dimensions: []const u32 = &.{}, + /// Every width from this one to `dimensions` crops (a model trained for any MRL width); 0: only the list's. + mrl_min: u32 = 0, pub fn dimension(i: Info, requested: ?u32) !u32 { const n = requested orelse return i.dimensions; if (n == 0 or n > i.dimensions) return error.UnsupportedEmbeddingDimension; - if (i.mrl_dimensions.len == 0) return error.MrlNotDeclared; + if (i.mrl_dimensions.len == 0 and i.mrl_min == 0) return error.MrlNotDeclared; if (n == i.dimensions or std.mem.indexOfScalar(u32, i.mrl_dimensions, n) != null) return n; + if (i.mrl_min > 0 and n >= i.mrl_min) return n; return error.UnsupportedEmbeddingDimension; } }; @@ -47,3 +50,16 @@ fn overlap(a: []const u8, b: []const u8) bool { const bp = @intFromPtr(b.ptr); return if (ap <= bp) bp - ap < a.len else ap - bp < b.len; } + +test "an MRL range takes every width from its least to the full one; a list takes its own widths" { + const range: Info = .{ .dimensions = 4096, .vocabulary = 10, .max_batch_tokens = 8, .mrl_min = 32 }; + for ([_]u32{ 32, 384, 2049, 4096 }) |n| try std.testing.expectEqual(n, try range.dimension(n)); + for ([_]u32{ 0, 31, 4097 }) |n| try std.testing.expectError(error.UnsupportedEmbeddingDimension, range.dimension(n)); + try std.testing.expectEqual(@as(u32, 4096), try range.dimension(null)); + const list: Info = .{ .dimensions = 8, .vocabulary = 10, .max_batch_tokens = 8, .mrl_dimensions = &.{ 2, 4 } }; + try std.testing.expectEqual(@as(u32, 4), try list.dimension(4)); + try std.testing.expectError(error.UnsupportedEmbeddingDimension, list.dimension(3)); + const none: Info = .{ .dimensions = 8, .vocabulary = 10, .max_batch_tokens = 8 }; + try std.testing.expectError(error.MrlNotDeclared, none.dimension(4)); + try std.testing.expectEqual(@as(u32, 8), try none.dimension(null)); +} From ab81d7275908f03b4fa3b3990334eb92b6c0c215 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Mon, 2 Nov 2026 22:48:21 +0000 Subject: [PATCH 08/24] fix(qwen3): read a tied model's head rows from its embedding Qwen3-Reranker-4B ties its output head to the embedding and stores no lm_head.weight, so readRows of the head failed for it; the model now records the stored name of its head rows (lm_head.weight, or the embedding as the checkpoint names it when tie_word_embeddings is set) and the GPU check reads that name, with host tests for both layouts. --- zig/src/families/qwen3/forward_test.zig | 2 +- zig/src/families/qwen3/model.zig | 15 ++++++++++++++- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/zig/src/families/qwen3/forward_test.zig b/zig/src/families/qwen3/forward_test.zig index 3cc2af8584..edcec15fab 100644 --- a/zig/src/families/qwen3/forward_test.zig +++ b/zig/src/families/qwen3/forward_test.zig @@ -175,7 +175,7 @@ test "Qwen3-Reranker: each case's yes - no logit from lastRows and lm_head rows defer m.deinit(); const yes: u32 = @intCast(ref.object.get("yes").?.integer); const no: u32 = @intCast(ref.object.get("no").?.integer); - const head = try qwen3.readRows(gpa, t.io, dir, "lm_head.weight", &.{ yes, no }); + const head = try qwen3.readRows(gpa, t.io, dir, m.head, &.{ yes, no }); defer gpa.free(head); const hidden = m.spec.hidden; var worst_diff: f64 = 0; diff --git a/zig/src/families/qwen3/model.zig b/zig/src/families/qwen3/model.zig index c57b620fa1..f504071ee6 100644 --- a/zig/src/families/qwen3/model.zig +++ b/zig/src/families/qwen3/model.zig @@ -31,6 +31,8 @@ pub const Model = struct { /// The activation dtype: BF16, as the checkpoint's weights. act: hip.ops.Kind = .bf16, embed: Projection, + /// The stored name of the output head's rows for `readRows`: lm_head, or the embedding when the model ties them. + head: []const u8, final_norm: u64, layers: []Layer, buffers: std.ArrayList(hip.memory.DeviceBuffer) = .empty, @@ -41,7 +43,7 @@ pub const Model = struct { var ck = try Checkpoint.openModel(gpa, io, dir); defer ck.close(); ck.scope = scopeOf(&ck); - var m: Model = .{ .gpa = gpa, .spec = spec, .embed = undefined, .final_norm = 0, .layers = &.{} }; + var m: Model = .{ .gpa = gpa, .spec = spec, .embed = undefined, .head = headKey(spec.tied, ck.scope), .final_norm = 0, .layers = &.{} }; errdefer m.deinit(); const up: hip.Upload = .{ .r = rt, .buffers = &m.buffers, .gpa = gpa }; var scratch: std.heap.ArenaAllocator = .init(gpa); @@ -72,6 +74,12 @@ pub fn scopeOf(ck: *const Checkpoint) []const u8 { return if (ck.has("model.embed_tokens.weight")) "model." else ""; } +/// The head rows' stored name: `lm_head.weight`, or the embedding's (prefixed as the checkpoint stores it) when tied. +pub fn headKey(tied: bool, scope: []const u8) []const u8 { + if (!tied) return "lm_head.weight"; + return if (scope.len > 0) "model.embed_tokens.weight" else "embed_tokens.weight"; +} + /// `quant.read`'s table over a checkpoint: dense tensors only, so no format asks for an affine width. const Table = struct { ck: *Checkpoint, @@ -205,5 +213,10 @@ test "the reranker's model. prefix is found, the embedding's bare names are not, defer a.free(rows); try std.testing.expectEqualSlices(f32, &.{ 3, -3, 1, -1 }, rows); try std.testing.expectError(error.UnexpectedTensor, readRows(a, io, dir, "lm_head.weight", &.{3})); + // a tied model reads its head rows from the embedding, as stored + const tied = try readRows(a, io, dir, headKey(true, scopeOf(&ck)), &.{1}); + defer a.free(tied); + try std.testing.expectEqualSlices(f32, &.{ 0, 0 }, tied); + try std.testing.expectEqualStrings("lm_head.weight", headKey(false, scopeOf(&ck))); } } From a4b8e0857ce59941484575457e42c46e95b65051 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Mon, 2 Nov 2026 23:37:39 +0000 Subject: [PATCH 09/24] feat(qwen3): serve Qwen3 embedding checkpoints on the HIP engine The embedding route and lane contract had no device backend, and the HIP registry no qwen3 family. A prompt-only engine sizes its scratch to the memory budget, shrinking the window to fit, and a lane backend runs each batch's whole prompt in one step and returns the pooled last rows, unit-normed when modules.json lists Normalize; the head is read from the sentence-transformers folder. The registry gains qwen3, a family's Loaded can say it serves embeddings in its own window, and one-GPU families refuse --tp before any group forms. --- zig/src/families/qwen3/engine.zig | 127 +++++++++++++++++ zig/src/families/qwen3/native.zig | 230 ++++++++++++++++++++++++++++++ zig/src/families/qwen3/prompt.zig | 106 ++++++++++++++ zig/src/families/qwen3/qwen3.zig | 6 + zig/src/native/hip.zig | 47 +++++- 5 files changed, 514 insertions(+), 2 deletions(-) create mode 100644 zig/src/families/qwen3/engine.zig create mode 100644 zig/src/families/qwen3/native.zig create mode 100644 zig/src/families/qwen3/prompt.zig diff --git a/zig/src/families/qwen3/engine.zig b/zig/src/families/qwen3/engine.zig new file mode 100644 index 0000000000..197968f27a --- /dev/null +++ b/zig/src/families/qwen3/engine.zig @@ -0,0 +1,127 @@ +//! One Qwen3 dense model on one HIP device for prompt-only heads: weights, a stream and scratch sized to a window. +const std = @import("std"); +const hip = @import("hip"); +const model = @import("model.zig"); +const forward = @import("forward.zig"); + +const Allocator = std.mem.Allocator; + +/// Inputs one pass pools at most (the embedding lane's default batch). +pub const max_inputs: usize = 128; +/// The least window the budget may shrink a load to. +pub const least_window: usize = 2048; + +pub const Engine = struct { + gpa: Allocator, + driver: hip.Runtime, + ctx: hip.Context, + lib: hip.Launcher, + stream: hip.Stream, + m: model.Model, + /// Device bytes the weights took. + weights: usize = 0, + /// Tokens a pass takes at most, over all its inputs; set by `size`. + window: usize = 0, + arena: hip.Arena = undefined, + ids_host: hip.HostBuffer = undefined, + ids_dev: hip.DeviceBuffer = undefined, + rows_dev: hip.DeviceBuffer = undefined, + rows_host: hip.HostBuffer = undefined, + sized: bool = false, + + /// The model on device `index`, its scratch not yet sized. + pub fn load(gpa: Allocator, io: std.Io, dir: []const u8, index: c_int, choices: hip.Launcher.Choices) !*Engine { + const e = try gpa.create(Engine); + errdefer gpa.destroy(e); + e.* = .{ .gpa = gpa, .driver = undefined, .ctx = undefined, .lib = undefined, .stream = undefined, .m = undefined }; + e.driver = try hip.Runtime.open(); + errdefer e.driver.close(); + e.ctx = try hip.Context.init(&e.driver, index); + errdefer e.ctx.deinit(); + e.lib = try hip.Launcher.load(&e.driver, choices, hip.kernels.images); + errdefer e.lib.unload(); + e.stream = try hip.Stream.init(&e.driver); + errdefer e.stream.deinit(); + const before = hip.usage(false).device; + e.m = try model.Model.load(gpa, io, &e.driver, dir); + e.weights = hip.usage(false).device - before; + return e; + } + + pub fn ops(e: *Engine, arena: *hip.Arena) hip.ops.Ops { + return .{ .l = &e.lib, .bf16 = true, .stream = e.stream.handle, .arena = arena }; + } + + /// The pass's scratch for `window` tokens (one input of them: the longest cache) plus the ids and pooled rows. + pub fn scratchBytes(e: *Engine, window: usize) !usize { + var arena = hip.Arena.counting(); + var o = e.ops(&arena); + o.stream = hip.counting; + const out: hip.ops.Tensor = .{ .ptr = try arena.take(max_inputs * e.m.spec.hidden * 2), .kind = .bf16 }; + try forward.lastRows(o, &e.m, 0, &[_]u32{@intCast(window)}, out); + return arena.peak + Engine.fixedBytes(e, window); + } + + /// The buffers beside the arena: ids for `window` tokens and `max_inputs` pooled BF16 rows. + fn fixedBytes(e: *const Engine, window: usize) usize { + return window * 4 + max_inputs * e.m.spec.hidden * 2; + } + + /// Allocates the scratch for passes of `window` tokens; refuses when the bytes held differ from the count. + pub fn size(e: *Engine, window: usize) !void { + const want = try e.scratchBytes(window); + const arena_bytes = want - e.fixedBytes(window); + const before = hip.usage(false).device; + e.arena = try hip.Arena.init(&e.driver, arena_bytes); + errdefer e.arena.deinit(); + e.ids_dev = try hip.DeviceBuffer.alloc(&e.driver, window * 4); + errdefer e.ids_dev.free(); + e.rows_dev = try hip.DeviceBuffer.alloc(&e.driver, max_inputs * e.m.spec.hidden * 2); + errdefer e.rows_dev.free(); + e.ids_host = try hip.HostBuffer.alloc(&e.driver, window * 4); + errdefer e.ids_host.free(); + e.rows_host = try hip.HostBuffer.alloc(&e.driver, max_inputs * e.m.spec.hidden * 2); + errdefer e.rows_host.free(); + const held = hip.usage(false).device - before; + if (held != want) { + std.log.err("qwen3: sizing held {d} device bytes, the count says {d}", .{ held, want }); + return error.SizeMismatch; + } + e.window = window; + e.sized = true; + } + + /// Each input's final-normed last row as fp32 into `out` (ends.len rows of hidden); inputs end at `ends`. + pub fn run(e: *Engine, tokens: []const u32, ends: []const u32, out: []f32) !void { + const hidden = e.m.spec.hidden; + if (!e.sized or tokens.len == 0 or tokens.len > e.window or ends.len == 0 or ends.len > max_inputs or + ends[ends.len - 1] != tokens.len or out.len != ends.len * hidden) return error.InvalidBatch; + for (tokens) |t| if (t >= e.m.spec.vocab) return error.InvalidBatch; + @memcpy(e.ids_host.slice(u32)[0..tokens.len], tokens); + try e.ids_dev.uploadAsync(0, e.ids_host.view(0, tokens.len * 4), e.stream); + const mark = e.arena.mark(); + defer e.arena.release(mark); + const rows: hip.ops.Tensor = .{ .ptr = e.rows_dev.base(), .kind = .bf16 }; + try forward.lastRows(e.ops(&e.arena), &e.m, e.ids_dev.base(), ends, rows); + try e.rows_dev.downloadAsync(0, e.rows_host.view(0, ends.len * hidden * 2), e.stream); + try e.stream.synchronize(); + for (out, e.rows_host.slice(u16)[0..out.len]) |*o, b| o.* = @bitCast(@as(u32, b) << 16); + } + + pub fn deinit(e: *Engine) void { + e.stream.synchronize() catch {}; + if (e.sized) { + e.rows_host.free(); + e.ids_host.free(); + e.rows_dev.free(); + e.ids_dev.free(); + e.arena.deinit(); + } + e.m.deinit(); + e.stream.deinit(); + e.lib.unload(); + e.ctx.deinit(); + e.driver.close(); + e.gpa.destroy(e); + } +}; diff --git a/zig/src/families/qwen3/native.zig b/zig/src/families/qwen3/native.zig new file mode 100644 index 0000000000..e75de11ccb --- /dev/null +++ b/zig/src/families/qwen3/native.zig @@ -0,0 +1,230 @@ +//! Qwen3 dense checkpoints for the native HIP server: the head the folder declares, sized to the memory budget. +const std = @import("std"); +const hip = @import("hip"); +const lanes = @import("lanes"); +const api = @import("engine_api"); +const engine = @import("engine.zig"); +const prompt = @import("prompt.zig"); + +const Engine = engine.Engine; +const Allocator = std.mem.Allocator; + +/// The registry entry: `native/hip.zig` opens `qwen3` checkpoints through it. +pub const model_type = "qwen3"; +pub const formats: []const []const u8 = &.{"bf16"}; +/// The window when neither --context nor the checkpoint's config names one (the Qwen3 8B serving default). +pub const default_context: i64 = 32768; +/// The engine takes a stream's whole prompt in one step. +pub const prefill_step: u32 = 0; +pub const slide_refusal = "--slide: Qwen3 dense checkpoints do not learn on HIP"; +pub const confidence_refusal = "--mtp-confidence: Qwen3 dense checkpoints have no draft head"; +/// An embedding engine writes no reply text, so no reasoning block (Info.reasoning). +pub const reasoning = false; +/// Checked by native/hip.zig before any rank joins a group. +pub const tp_refusal = "--tp: Qwen3 dense checkpoints run on one HIP GPU; serve with --tp 1"; +pub const open = openEngine; +pub const follow = followRank; + +/// What the server drives, as the other HIP families' `Loaded`, plus what this engine serves. +pub const Loaded = struct { + backend: lanes.backend.Backend, + facts: lanes.Model, + rows: u32, + ctx: *anyopaque, + deinit: *const fn (*anyopaque) void, + streams: usize, + startup: []const u8, + /// No chat or completions: the server refuses the generation routes. + text_generation: bool = false, + embeddings: ?api.EmbeddingInfo = null, + /// The window the engine was sized to (the asked one, or less where the budget shrank it). + window: ?usize = null, +}; + +pub const Refused = error{Refused}; + +fn refuse(a: Allocator, problem: *[]const u8, comptime fmt: []const u8, args: anytype) Refused { + problem.* = std.fmt.allocPrint(a, fmt, args) catch fmt; + return error.Refused; +} + +/// The head a checkpoint folder declares. +pub const Head = union(enum) { + /// sentence-transformers last-token pooling, unit-normed when modules.json lists Normalize. + embedding: struct { normalize: bool }, +}; + +const pooling_modes = [_][]const u8{ + "pooling_mode_cls_token", "pooling_mode_mean_tokens", "pooling_mode_max_tokens", + "pooling_mode_mean_sqrt_len_tokens", "pooling_mode_weightedmean_tokens", "pooling_mode_lasttoken", +}; + +/// The head from `1_Pooling/config.json` and `modules.json`, or a refusal naming what the folder lacks. +pub fn headOf(a: Allocator, io: std.Io, dir: []const u8, problem: *[]const u8) (Refused || Allocator.Error)!Head { + const pooling = try readJson(a, io, dir, "1_Pooling/config.json", problem) orelse return refuse( + a, + problem, + "qwen3 on HIP serves sentence-transformers embedding checkpoints: {s} has no 1_Pooling/config.json", + .{dir}, + ); + if (pooling != .object) return refuse(a, problem, "1_Pooling/config.json is not a JSON object", .{}); + for (pooling_modes) |mode| { + const on = if (pooling.object.get(mode)) |v| v == .bool and v.bool else false; + const want = std.mem.eql(u8, mode, "pooling_mode_lasttoken"); + if (on != want) return refuse(a, problem, "1_Pooling/config.json: {s} is {s}; qwen3 on HIP pools the last token only", .{ mode, if (on) "on" else "off" }); + } + // no modules.json: the pooled rows as they are; one that is there must be a list of modules with string types + var normalize = false; + if (try readJson(a, io, dir, "modules.json", problem)) |mods| { + if (mods != .array) return refuse(a, problem, "modules.json is not a JSON list of modules", .{}); + for (mods.array.items) |m| { + const kind = if (m == .object) m.object.get("type") else null; + if (kind == null or kind.? != .string) return refuse(a, problem, "modules.json: a module without a string \"type\"", .{}); + if (std.mem.eql(u8, kind.?.string, "sentence_transformers.models.Normalize")) normalize = true; + } + } + return .{ .embedding = .{ .normalize = normalize } }; +} + +/// The folder's JSON file `name`, null when it is absent; unreadable or malformed JSON is refused. +fn readJson(a: Allocator, io: std.Io, dir: []const u8, name: []const u8, problem: *[]const u8) (Refused || Allocator.Error)!?std.json.Value { + const path = try std.fs.path.join(a, &.{ dir, name }); + const bytes = std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 20)) catch |err| switch (err) { + error.FileNotFound => return null, + error.OutOfMemory => return error.OutOfMemory, + else => return refuse(a, problem, "{s}: cannot read it ({s})", .{ name, @errorName(err) }), + }; + return std.json.parseFromSliceLeaky(std.json.Value, a, bytes, .{}) catch |err| switch (err) { + error.OutOfMemory => error.OutOfMemory, + else => refuse(a, problem, "{s} is not valid JSON ({s})", .{ name, @errorName(err) }), + }; +} + +fn gibs(bytes: usize) f64 { + return @as(f64, @floatFromInt(bytes)) / (1 << 30); +} + +/// The window the budget holds: `asked`, halved while its scratch passes `room`, down to `engine.least_window`. +pub fn fitWindow(asked: usize, room: usize, ctx: anytype, bytesFor: fn (@TypeOf(ctx), usize) anyerror!usize) !?usize { + var window = asked; + while (true) { + if (try bytesFor(ctx, window) <= room) return window; + if (window <= engine.least_window) return null; + window = @max(window / 2, engine.least_window); + } +} + +fn scratchOf(e: *Engine, window: usize) anyerror!usize { + return e.scratchBytes(window); +} + +const Owned = struct { + gpa: Allocator, + e: *Engine, + p: prompt.Prompt, + startup: []const u8, +}; + +fn release(ptr: *anyopaque) void { + const own: *Owned = @ptrCast(@alignCast(ptr)); + own.p.deinit(); + own.e.deinit(); + own.gpa.free(own.startup); + own.gpa.destroy(own); +} + +/// The engine for `dir` on `dev`, its head from the folder and its window from the memory budget. +fn openEngine(a: Allocator, gpa: Allocator, io: std.Io, dev: hip.Device, dir: []const u8, o: anytype, problem: *[]const u8) anyerror!Loaded { + if (o.world > 1) return refuse(a, problem, "--tp {d}: Qwen3 dense checkpoints run on one HIP GPU", .{o.world}); + const head = try headOf(a, io, dir, problem); + const e = Engine.load(gpa, io, dir, dev.index, dev.policy.choices()) catch |err| + return refuse(a, problem, "the native HIP engine cannot load {s} ({s})", .{ dir, @errorName(err) }); + errdefer e.deinit(); + const pool = hip.admission.budget(gpa, io, e.ctx) catch |err| + return refuse(a, problem, "the native HIP engine cannot read the GPU's memory ({s})", .{@errorName(err)}); + const room = pool.room(e.weights); + const window = (fitWindow(o.window, room, e, scratchOf) catch |err| + return refuse(a, problem, "the native HIP engine cannot count its scratch ({s})", .{@errorName(err)})) orelse { + const least = e.scratchBytes(@min(o.window, engine.least_window)) catch 0; + return refuse(a, problem, "the HIP memory budget fits no {d}-token pass: it takes {d:.2} GiB, and {d:.2} GiB is left after the model's {d:.2} GiB and the {d:.1} GiB reserve; free device memory or lower --context", .{ @min(o.window, engine.least_window), gibs(least), gibs(room), gibs(e.weights), gibs(pool.reserve) }); + }; + e.size(window) catch |err| + return refuse(a, problem, "the native HIP engine cannot allocate its scratch ({s})", .{@errorName(err)}); + const own = try gpa.create(Owned); + errdefer gpa.destroy(own); + own.* = .{ .gpa = gpa, .e = e, .p = prompt.Prompt.init(gpa, e, head.embedding.normalize), .startup = undefined }; + const shrunk: []const u8 = if (window < o.window) " (shrunk from the asked window to fit the budget)" else ""; + own.startup = try std.fmt.allocPrint(gpa, "HIP {s} device {d}: qwen3 embedding, model {d:.2} GiB; {d}-token passes{s}, {d:.2} GiB scratch, of {d:.1} GiB left after a {d:.1} GiB reserve; vectors {s}", .{ + hip.kernels.arch, dev.index, gibs(e.weights), window, shrunk, gibs(e.scratchBytes(window) catch 0), gibs(room), gibs(pool.reserve), if (head.embedding.normalize) "unit-normed" else "as pooled", + }); + const spec = e.m.spec; + return .{ + .backend = own.p.backend(), + .facts = .{}, + .rows = 1, + .ctx = own, + .deinit = release, + .streams = 1, + .startup = own.startup, + .text_generation = false, + .embeddings = .{ .dimensions = @intCast(spec.hidden), .vocabulary = @intCast(spec.vocab), .max_batch_tokens = @intCast(window), .max_inputs = @intCast(engine.max_inputs), .mrl_min = @min(32, @as(u32, @intCast(spec.hidden))) }, + .window = window, + }; +} + +/// Ranks above 0 have nothing to hold: the engine runs on one GPU. +fn followRank(a: Allocator, _: Allocator, _: std.Io, _: hip.Device, _: []const u8, o: anytype, problem: *[]const u8) anyerror!void { + return refuse(a, problem, "--tp {d}: Qwen3 dense checkpoints run on one HIP GPU", .{o.world}); +} + +test "the window halves to fit the room, down to the least window" { + const Count = struct { + per: usize, + fn bytes(c: @This(), window: usize) anyerror!usize { + return window * c.per; + } + }; + const c: Count = .{ .per = 10 }; + try std.testing.expectEqual(@as(?usize, 32768), try fitWindow(32768, 327680, c, Count.bytes)); + try std.testing.expectEqual(@as(?usize, 8192), try fitWindow(32768, 100000, c, Count.bytes)); + try std.testing.expectEqual(@as(?usize, 2048), try fitWindow(32768, 20480, c, Count.bytes)); + try std.testing.expectEqual(@as(?usize, null), try fitWindow(32768, 20479, c, Count.bytes)); + try std.testing.expectEqual(@as(?usize, 1000), try fitWindow(1000, 10000, c, Count.bytes)); + try std.testing.expectEqual(@as(?usize, null), try fitWindow(1000, 9999, c, Count.bytes)); +} + +test "a pooling folder declares the embedding head: last token only, normalized when modules.json says so" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + var problem: []const u8 = ""; + try std.testing.expectError(error.Refused, headOf(a, io, dir, &problem)); + try std.testing.expect(std.mem.indexOf(u8, problem, "1_Pooling") != null); + try tmp.dir.createDirPath(io, "1_Pooling"); + const last = + \\{"word_embedding_dimension": 4096, "pooling_mode_cls_token": false, "pooling_mode_mean_tokens": false, + \\ "pooling_mode_max_tokens": false, "pooling_mode_mean_sqrt_len_tokens": false, + \\ "pooling_mode_weightedmean_tokens": false, "pooling_mode_lasttoken": true, "include_prompt": true} + ; + try tmp.dir.writeFile(io, .{ .sub_path = "1_Pooling/config.json", .data = last }); + try std.testing.expect(!(try headOf(a, io, dir, &problem)).embedding.normalize); + try tmp.dir.writeFile(io, .{ .sub_path = "modules.json", .data = "[{\"type\": \"sentence_transformers.models.Pooling\"}, {\"type\": \"sentence_transformers.models.Normalize\"}]" }); + try std.testing.expect((try headOf(a, io, dir, &problem)).embedding.normalize); + try tmp.dir.writeFile(io, .{ .sub_path = "modules.json", .data = "[{\"type\": \"sentence_transformers.models.Normalize\"" }); + try std.testing.expectError(error.Refused, headOf(a, io, dir, &problem)); + try std.testing.expect(std.mem.indexOf(u8, problem, "not valid JSON") != null); + try tmp.dir.writeFile(io, .{ .sub_path = "modules.json", .data = "{\"type\": \"sentence_transformers.models.Normalize\"}" }); + try std.testing.expectError(error.Refused, headOf(a, io, dir, &problem)); + try tmp.dir.writeFile(io, .{ .sub_path = "modules.json", .data = "[{\"path\": \"1_Pooling\"}]" }); + try std.testing.expectError(error.Refused, headOf(a, io, dir, &problem)); + try tmp.dir.writeFile(io, .{ .sub_path = "modules.json", .data = "[]" }); + try tmp.dir.writeFile(io, .{ .sub_path = "1_Pooling/config.json", .data = "{\"pooling_mode_lasttoken\": tru" }); + try std.testing.expectError(error.Refused, headOf(a, io, dir, &problem)); + try tmp.dir.writeFile(io, .{ .sub_path = "1_Pooling/config.json", .data = "{\"pooling_mode_mean_tokens\": true, \"pooling_mode_lasttoken\": true}" }); + try std.testing.expectError(error.Refused, headOf(a, io, dir, &problem)); + try std.testing.expect(std.mem.indexOf(u8, problem, "pooling_mode_mean_tokens") != null); +} diff --git a/zig/src/families/qwen3/prompt.zig b/zig/src/families/qwen3/prompt.zig new file mode 100644 index 0000000000..2ff8e34887 --- /dev/null +++ b/zig/src/families/qwen3/prompt.zig @@ -0,0 +1,106 @@ +//! The lane backend of a prompt-only Qwen3 engine: a stream's whole prompt in one step, pooled rows until release. +const std = @import("std"); +const lanes = @import("lanes"); +const api = @import("engine_api"); +const Engine = @import("engine.zig").Engine; + +const be = lanes.backend; +const Stream = lanes.Stream; +const Allocator = std.mem.Allocator; + +/// What a prompt-only engine refuses: drawing, verifying or drafting tokens. +pub const Refusal = error{PromptOnly}; + +pub const Prompt = struct { + gpa: Allocator, + engine: *Engine, + /// The checkpoint's Normalize module: unit L2 rows. + normalize: bool, + /// Each stream's pooled fp32 rows, from its prompt pass to its release. + slots: std.AutoHashMapUnmanaged(*Stream, []f32) = .empty, + table: be.Backend.VTable = .{ + .embedding = embedding, + .prefill = prefill, + .prefill_step = prefillStep, + .first = first, + .queue = queue, + .read = read, + .verify = verify, + .keep = keep, + .draft = draft, + .release = release, + }, + + pub fn init(gpa: Allocator, engine: *Engine, normalize: bool) Prompt { + return .{ .gpa = gpa, .engine = engine, .normalize = normalize }; + } + + pub fn deinit(p: *Prompt) void { + var it = p.slots.valueIterator(); + while (it.next()) |rows| p.gpa.free(rows.*); + p.slots.deinit(p.gpa); + } + + /// The backend over this engine; heads add their hooks to `table` before it serves. + pub fn backend(p: *Prompt) be.Backend { + return .{ .ptr = p, .vtable = &p.table, .chunk_invariant = true }; + } + + fn of(ptr: *anyopaque) *Prompt { + return @ptrCast(@alignCast(ptr)); + } + + /// The whole prompt in one step: every input's last row, kept for `embedding`. + fn prefillStep(ptr: *anyopaque, s: *Stream, _: u32) anyerror!bool { + const p = of(ptr); + const batch = s.embedding orelse return error.PromptOnly; + const hidden = p.engine.m.spec.hidden; + const rows = try p.gpa.alloc(f32, batch.ends.len * hidden); + errdefer p.gpa.free(rows); + try p.engine.run(s.prompt(), batch.ends, rows); + const slot = try p.slots.getOrPut(p.gpa, s); + if (slot.found_existing) p.gpa.free(slot.value_ptr.*); + slot.value_ptr.* = rows; + return true; + } + + fn prefill(ptr: *anyopaque, s: *Stream) anyerror!void { + _ = try prefillStep(ptr, s, 0); + } + + /// The stream's pooled rows into `out`, each unit-normed when the checkpoint normalizes. + fn embedding(ptr: *anyopaque, s: *Stream, out: []f32) anyerror!void { + const p = of(ptr); + const rows = p.slots.get(s) orelse return error.EmbeddingNotFilled; + if (rows.len != out.len) return error.EmbeddingShape; + @memcpy(out, rows); + if (!p.normalize) return; + const hidden = p.engine.m.spec.hidden; + var at: usize = 0; + while (at < out.len) : (at += hidden) try api.embedding_vector.normalize(out[at..][0..hidden]); + } + + fn release(ptr: *anyopaque, s: *Stream) void { + const p = of(ptr); + if (p.slots.fetchRemove(s)) |kv| p.gpa.free(kv.value); + } + + fn first(_: *anyopaque, _: *Stream, _: u64) anyerror!u64 { + return error.PromptOnly; + } + fn queue(_: *anyopaque, _: *Stream, _: be.Feed, _: u64) anyerror!u64 { + return error.PromptOnly; + } + fn read(_: *anyopaque, _: u64) anyerror!u32 { + return error.PromptOnly; + } + fn verify(_: *anyopaque, _: []const be.Window, _: []be.Verified) anyerror!void { + return error.PromptOnly; + } + fn keep(_: *anyopaque, _: []const be.Window, _: []const []const u32) anyerror!void { + return error.PromptOnly; + } + fn draft(_: *anyopaque, _: []const be.DraftRequest) anyerror!void { + return error.PromptOnly; + } +}; diff --git a/zig/src/families/qwen3/qwen3.zig b/zig/src/families/qwen3/qwen3.zig index 68c56adede..775e417391 100644 --- a/zig/src/families/qwen3/qwen3.zig +++ b/zig/src/families/qwen3/qwen3.zig @@ -3,6 +3,9 @@ pub const config = @import("config.zig"); pub const model = @import("model.zig"); pub const forward = @import("forward.zig"); +pub const engine = @import("engine.zig"); +pub const prompt = @import("prompt.zig"); +pub const native = @import("native.zig"); pub const Spec = config.Spec; pub const Model = model.Model; pub const readRows = model.readRows; @@ -12,4 +15,7 @@ test { _ = config; _ = model; _ = forward; + _ = engine; + _ = prompt; + _ = native; } diff --git a/zig/src/native/hip.zig b/zig/src/native/hip.zig index 8a922cdcc6..3a74134408 100644 --- a/zig/src/native/hip.zig +++ b/zig/src/native/hip.zig @@ -9,12 +9,13 @@ const flashnext = @import("flashnext"); const nemotron = @import("nemotron"); const serve = @import("hip_serve"); const glm_ocr = @import("glm_ocr"); +const qwen3 = @import("qwen3"); const Allocator = std.mem.Allocator; const admission = hip.admission; const modelContext = admission.modelContext; /// The HIP families: namespaces with `model_type`, `formats`, `default_context`, `prefill_step`, `open` and `follow`. -const registry = .{ qwen35.native, qwen35.native_moe, flashnext.native, nemotron.native, glm_ocr.native }; +const registry = .{ qwen35.native, qwen35.native_moe, flashnext.native, nemotron.native, glm_ocr.native, qwen3.native }; pub const backends: []const []const u8 = &.{"hip"}; /// No per-thread alternate signal stack: its thread-local leaves RCCL's and HIP's threads too little stack to start. @@ -310,7 +311,8 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O h.clock = .{ .io = io }; h.core = lanes.Engine.init(gpa, &h.cfg, loaded.backend, h.clock.clock()); errdefer h.core.deinit(); - h.host = api.LaneHost.init(gpa, io, &h.core, .{ .benchmark_ranks = if (r.dev.group) |g| @intCast(g.world) else 1, .lanes = @intCast(loaded.streams), .context_window = @intCast(r.options.window), .prefill_step = F.prefill_step, .decode_share = o.decode_share, .policy = h.policy_line, .startup = loaded.startup, .structures = true, .vision = visionOffer(loaded), .reasoning = reasoning(F) }); + const what = Serves.of(loaded, r.options.window); + h.host = api.LaneHost.init(gpa, io, &h.core, .{ .benchmark_ranks = if (r.dev.group) |g| @intCast(g.world) else 1, .lanes = @intCast(loaded.streams), .context_window = @intCast(what.window), .prefill_step = F.prefill_step, .decode_share = o.decode_share, .policy = h.policy_line, .startup = loaded.startup, .structures = what.text_generation, .text_generation = what.text_generation, .embeddings = what.embeddings, .vision = visionOffer(loaded), .reasoning = reasoning(F) }); h.host.memory = .{ .read = readMemory }; if (@hasField(@TypeOf(loaded), "cache")) { h.host.cache = loaded.cache; @@ -322,6 +324,24 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; } +/// What a family's `Loaded` says it serves; a family that names none generates text in the window asked. +const Serves = struct { + text_generation: bool = true, + embeddings: ?api.EmbeddingInfo = null, + window: usize, + + fn of(loaded: anytype, asked: usize) Serves { + const L = @TypeOf(loaded); + var s: Serves = .{ .window = asked }; + if (@hasField(L, "text_generation")) s.text_generation = loaded.text_generation; + if (@hasField(L, "embeddings")) s.embeddings = loaded.embeddings; + if (@hasField(L, "window")) if (loaded.window) |w| { + s.window = w; + }; + return s; + } +}; + /// A rank above 0: holds its share and runs rank 0's steps until rank 0 stops; false with `problem` if it cannot start. pub fn follow(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !bool { const report = @import("sys").optional_libraries.atBoundary(); @@ -366,6 +386,7 @@ test "every registered family is listed for capabilities" { try std.testing.expectEqualStrings("qwen4_exp", families[2].model_type); try std.testing.expectEqualStrings("nemotron_h", families[3].model_type); try std.testing.expectEqualStrings("glm_ocr", families[4].model_type); + try std.testing.expectEqualStrings("qwen3", families[5].model_type); } test "GLM-OCR native flags: MTP and tensor ranks refuse before readiness; copy drafts remain" { @@ -386,6 +407,17 @@ test "Nemotron refuses --slide and --mtp-confidence on HIP before it loads, and try std.testing.expect(flagRefused(nemotron.native, .{ .dir = "d", .model_type = "nemotron_h" }) == null); } +test "a one-GPU family refuses --tp before the group forms; the others are not refused for it" { + const OneGpu = struct { + pub const tp_refusal = "--tp: one GPU only"; + }; + try std.testing.expectEqualStrings("--tp: one GPU only", flagRefused(OneGpu, .{ .dir = "d", .model_type = "m", .tp = 2 }).?); + try std.testing.expect(flagRefused(OneGpu, .{ .dir = "d", .model_type = "m" }) == null); + try std.testing.expect(flagRefused(nemotron.native, .{ .dir = "d", .model_type = "nemotron_h", .tp = 2 }) == null); + try std.testing.expectEqualStrings(qwen3.native.tp_refusal, flagRefused(qwen3.native, .{ .dir = "d", .model_type = "qwen3", .tp = 2, .rank = 1 }).?); + try std.testing.expect(!reasoning(qwen3.native)); +} + test "a family that names a --slide or --mtp-confidence refusal is refused before it loads; the others are not" { const Refuses = struct { pub const slide_refusal = "--slide: this family does not learn on HIP"; @@ -482,3 +514,14 @@ test "slide admission: HIP followers refuse before joining ranks" { try std.testing.expectEqualStrings(try std.fmt.allocPrint(a, "--slide: {s} checkpoints do not learn on HIP yet", .{f.model_type}), why); } } + +test "a family's Loaded names what it serves; one without the fields generates text in the asked window" { + const Plain = struct { streams: usize }; + const p = Serves.of(Plain{ .streams = 1 }, 4096); + try std.testing.expect(p.text_generation and p.embeddings == null and p.window == 4096); + const Embed = struct { text_generation: bool, embeddings: ?api.EmbeddingInfo, window: ?usize }; + const e = Serves.of(Embed{ .text_generation = false, .embeddings = .{ .dimensions = 8, .vocabulary = 4, .max_batch_tokens = 2048 }, .window = 2048 }, 32768); + try std.testing.expect(!e.text_generation and e.embeddings.?.dimensions == 8 and e.window == 2048); + const Unsized = struct { window: ?usize }; + try std.testing.expectEqual(@as(usize, 32768), Serves.of(Unsized{ .window = null }, 32768).window); +} From f2c33043bdd756b1ccf70f61d3bf3dd2809ae58c Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Mon, 2 Nov 2026 23:38:17 +0000 Subject: [PATCH 10/24] test(qwen3): pin the model-card oracle for Qwen3-Embedding-8B The embedding gate compares the HIP engine with the model card's Transformers example, so the reference must come from known sources. The oracle pins the checkpoint revision, the SHA256 of every file in the snapshot, the Transformers 5.22.1 package tree and its Qwen3 files, checks them before any model code runs, writes them into each reference, and the server check refuses a reference made otherwise. It also checks distance, norms, batch and token-ID equality. --- tools/zig/oracle_qwen3_embedding.py | 264 +++++++++++++++++++++++ tools/zig/test_oracle_qwen3_embedding.py | 43 ++++ 2 files changed, 307 insertions(+) create mode 100644 tools/zig/oracle_qwen3_embedding.py create mode 100644 tools/zig/test_oracle_qwen3_embedding.py diff --git a/tools/zig/oracle_qwen3_embedding.py b/tools/zig/oracle_qwen3_embedding.py new file mode 100644 index 0000000000..5e4a16f2e4 --- /dev/null +++ b/tools/zig/oracle_qwen3_embedding.py @@ -0,0 +1,264 @@ +"""Pinned offline Qwen3-Embedding-8B oracle: the model card's Transformers example, and a server checked against it.""" + +import argparse +import base64 +import hashlib +import importlib.metadata +import importlib.util +import json +import math +import sys +import urllib.request +from pathlib import Path + +USAGE = """ +Pins checked before any model code runs (the snapshot's revision, every file's SHA256, Transformers 5.22.1's tree): + python tools/zig/oracle_qwen3_embedding.py sources --model DIR +Reference (writes JSON with token ids, vectors and the card's 2 x 2 scores): + python tools/zig/oracle_qwen3_embedding.py ref --model DIR --dtype float32 --device cpu --out ref32.json +Compare a running server's /v1/embeddings (same texts, one batch and one by one) with the fp32 target and the +transformers BF16 bar: + python tools/zig/oracle_qwen3_embedding.py check --ref ref32.json --bar ref16.json --url http://127.0.0.1:8080 +Gates: cosine distance (1 - cos) to fp32 at most 0.001 and at most --bar-factor times the BF16 bar's worst distance; +unit norm; the card's scores within 0.01; batch, alone and token-id vectors equal as fp32 bytes (base64). +""" + +CHECKPOINT = ("Qwen/Qwen3-Embedding-8B", "1d8ad4ca9b3dd8059ad90a75d4983776a23d44af") +# Every non-weight file of the pinned snapshot: card, configs, sentence-transformers modules, tokenizer, licence. +SOURCES = { + ".gitattributes": "34448b82c17d60fec9b65b1f093c115ddbaadc04beb1b0140b6bfed2e012a930", + "1_Pooling/config.json": "2e1da26b3fd65cf7e370d2fabf28a8c59efa7edb525d1b8b50be8e5ca1048ea6", + "LICENSE": "832dd9e00a68dd83b3c3fb9f5588dad7dcf337a0db50f7d9483f310cd292e92e", + "README.md": "1fd6f3583d642a3319efb3eb8c22b6aa6320e6ab92263447bd84fd92615206d4", + "config.json": "dcccf7c7890c8debc4af8dace8d6acd8dd50bd56d50d4d8b60529381ba3de3d7", + "config_sentence_transformers.json": "10667c72ddb772627bf1780cb7f86af8e2ae0032b8c243c731172064105c6961", + "generation_config.json": "28396d421a2108acce96383f6a7de78008f7f1b17f807958f3c14c51dbfb65fb", + "merges.txt": "8831e4f1a044471340f7c0a83d7bd71306a5b867e95fd870f74d0c5308a904d5", + "model.safetensors.index.json": "ceaf1429bfa18eecc75d079737e8d1fe8dc0627f5ec7e6a985514626d1ac6a77", + "modules.json": "84e40c8e006c9b1d6c122e02cba9b02458120b5fb0c87b746c41e0207cf642cf", + "tokenizer.json": "83cdf8c3a34f68862319cb1810ee7b1e2c0a44e0864ae930194ddb76bb7feb8d", + "tokenizer_config.json": "2f58f4bbd7bbce15d683f525954ef3a92cd82f5e06415a9c513859bf8ab72436", + "vocab.json": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910", +} +# The shards' SHA256 (the Hub's LFS object ids). +WEIGHTS = { + "model-00001-of-00004.safetensors": "99b343597fe840706146144699a8b9188dd3387e43eb61faf0231b70b249d451", + "model-00002-of-00004.safetensors": "dff635b0f6dbbaad2a2d633ef037ec0a39bc165cc1806c712fbd6fcbcb4526c0", + "model-00003-of-00004.safetensors": "30b1d4c53d84eb018f642cad7b373f0aabf79699872d8702c1f38577c0a59a2f", + "model-00004-of-00004.safetensors": "36cbc9c60375693629f25743c1e77ebb1724af58e671b2376463193c7fd21ef6", +} +VERSION = "5.22.1" +# tree_sha256 of the whole installed Transformers package, so the shared modeling and tokenizer code is pinned too. +TRANSFORMERS_TREE = "b08cead6d1b841ef9723106a81a92ec37d67aff3e43a198e56b7dc76538d377b" +# The Qwen3 model, configuration and tokenizer code the card's example runs (package-relative). +LIBRARY = { + "models/qwen3/modeling_qwen3.py": "63a6213c15d6fb04e58a1c64fd05a4cdfb49476443fbeef524f0e0486043824b", + "models/qwen3/configuration_qwen3.py": "49d9ccd0f29ccd977b93dba2004f84b867891b3d0509b71fd8ba3ed01054e64d", + "models/qwen2/tokenization_qwen2.py": "fac4e6576bfe2369731be147a4e530f262bdf32f2ac50436f96f0d8bdd2fc628", +} + +TASK = "Given a web search query, retrieve relevant passages that answer the query" +CARD_SCORES = [[0.7493016123771667, 0.0750647559762001], [0.08795969933271408, 0.6318399906158447]] + + +def detailed_instruct(task: str, query: str) -> str: + return f"Instruct: {task}\nQuery:{query}" + + +def card_texts() -> list: + queries = [detailed_instruct(TASK, "What is the capital of China?"), detailed_instruct(TASK, "Explain gravity")] + documents = [ + "The capital of China is Beijing.", + "Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and " + "is responsible for the movement of planets around the sun.", + ] + return queries + documents + + +def digest(path: Path) -> str: + h = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1 << 20), b""): + h.update(chunk) + return h.hexdigest() + + +def hashed(root: Path, pins: dict, what: str) -> dict: + """The pinned files' hashes under root, or a refusal naming every file that differs from its pin.""" + found = {name: digest(root / name) if (root / name).is_file() else None for name in pins} + if wrong := sorted(name for name, value in found.items() if value != pins[name]): + raise SystemExit(f"FAIL: {what} differs from its pin: {', '.join(wrong)}") + return found + + +def tree_sha256(package: Path) -> str: + """Every file of an installed package (bytecode caches excluded) as sorted "path NUL sha256" lines, hashed.""" + if links := sorted(p.relative_to(package).as_posix() for p in package.rglob("*") if p.is_symlink() and p.is_dir()): + raise SystemExit(f"FAIL: {package.name} has linked directories: {', '.join(links)}") + files = sorted(p for p in package.rglob("*") if p.is_file() and "__pycache__" not in p.relative_to(package).parts) + return hashlib.sha256("".join(f"{p.relative_to(package).as_posix()}\0{digest(p)}\n" for p in files).encode()).hexdigest() + + +def pinned(model: Path) -> dict: + """Refuse, before any model code runs, a snapshot or Transformers install whose sources differ from the pins.""" + repo, revision = CHECKPOINT + model = model.resolve() + if model.parent.name == "snapshots" and model.name != revision: + raise SystemExit(f"FAIL: snapshot directory is not the pinned revision {revision}") + if links := sorted(p.relative_to(model).as_posix() for p in model.rglob("*") if p.is_symlink() and p.is_dir()): + raise SystemExit(f"FAIL: snapshot has linked directories: {', '.join(links)}") + files = {p.relative_to(model).as_posix() for p in model.rglob("*") if p.is_file()} + if extra := sorted(files - set(SOURCES) - set(WEIGHTS)): + raise SystemExit(f"FAIL: snapshot has unpinned files: {', '.join(extra)}") + if (version := importlib.metadata.version("transformers")) != VERSION: + raise SystemExit(f"FAIL: expected Transformers {VERSION}, found {version}") + package = Path(importlib.util.find_spec("transformers").origin).parent + if (tree := tree_sha256(package)) != TRANSFORMERS_TREE: + raise SystemExit(f"FAIL: Transformers {VERSION} package tree differs from its pin: {tree}") + return {"repo": repo, "revision": revision, "checkpoint": hashed(model, SOURCES, f"{repo}@{revision}"), + "weights": hashed(model, WEIGHTS, f"{repo}@{revision} weights"), + "transformers": VERSION, "transformers_tree": tree, + "library": hashed(package, LIBRARY, f"Transformers {VERSION}"), + "torch": importlib.metadata.version("torch")} + + +def sources(args) -> None: + print(json.dumps(pinned(args.model), indent=1)) + + +def expected_pins() -> dict: + """The pin record every reference must carry (torch's version is recorded, not pinned).""" + return {"repo": CHECKPOINT[0], "revision": CHECKPOINT[1], "checkpoint": SOURCES, "weights": WEIGHTS, + "transformers": VERSION, "transformers_tree": TRANSFORMERS_TREE, "library": LIBRARY} + + +def check_pins(name: str, made: dict) -> None: + """Refuse a reference whose pin record is missing a field or differs from the pins in any of them.""" + pins = made.get("pins") + if not isinstance(pins, dict): + raise SystemExit(f"FAIL: {name} carries no pin record") + if wrong := sorted(k for k, v in expected_pins().items() if pins.get(k) != v): + raise SystemExit(f"FAIL: {name} was not made under the pins: {', '.join(wrong)}") + + +def reference(args) -> None: + # the pins first: no Transformers or torch code runs before its tree and the snapshot are checked + pins = pinned(args.model) + import torch + import torch.nn.functional as F + from transformers import AutoModel, AutoTokenizer + + def last_token_pool(last_hidden_states, attention_mask): + # the card's pooling, verbatim in effect + left_padding = attention_mask[:, -1].sum() == attention_mask.shape[0] + if left_padding: + return last_hidden_states[:, -1] + lengths = attention_mask.sum(dim=1) - 1 + return last_hidden_states[torch.arange(last_hidden_states.shape[0], device=last_hidden_states.device), lengths] + + dtype = getattr(torch, args.dtype) + texts = card_texts() + tokenizer = AutoTokenizer.from_pretrained(args.model, padding_side="left", local_files_only=True) + model = AutoModel.from_pretrained(args.model, dtype=dtype, local_files_only=True).to(args.device).eval() + batch = tokenizer(texts, padding=True, truncation=True, max_length=8192, return_tensors="pt").to(args.device) + with torch.no_grad(): + out = model(**batch) + emb = F.normalize(last_token_pool(out.last_hidden_state, batch["attention_mask"]).float(), p=2, dim=1) + scores = (emb[:2] @ emb[2:].T).tolist() + ids = [tokenizer(t)["input_ids"] for t in texts] + print("scores", scores) + print("card ", CARD_SCORES) + json.dump({"dtype": args.dtype, "device": args.device, "pins": pins, "texts": texts, "ids": ids, + "vectors": emb.cpu().tolist(), "scores": scores}, open(args.out, "w")) + + +def post(url: str, body: dict) -> dict: + req = urllib.request.Request(url + "/v1/embeddings", data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=600) as r: + return json.load(r) + + +def vectors(resp: dict, count: int, width: int) -> list: + """The response's fp32 vectors, in index order, after count, index, width and finiteness checks.""" + data = resp.get("data", []) + if len(data) != count or [d.get("index") for d in data] != list(range(count)): + raise SystemExit(f"FAIL: expected {count} vectors with indices 0..{count - 1}, got {len(data)}") + out = [] + for d in data: + raw = base64.b64decode(d["embedding"]) + if len(raw) != 4 * width: + raise SystemExit(f"FAIL: vector {d['index']} has {len(raw) // 4} values, expected {width}") + out.append(raw) + for i, raw in enumerate(out): + if not all(math.isfinite(x) for x in floats(raw)): + raise SystemExit(f"FAIL: vector {i} is not finite") + return out + + +def floats(raw: bytes) -> list: + import struct + return list(struct.unpack(f"<{len(raw) // 4}f", raw)) + + +def cosine(a, b) -> float: + dot = sum(x * y for x, y in zip(a, b)) + return dot / math.sqrt(sum(x * x for x in a) * sum(y * y for y in b)) + + +def check(args) -> None: + ref = json.load(open(args.ref)) + bar = json.load(open(args.bar)) + if ref["dtype"] != "float32" or bar["dtype"] != "bfloat16" or ref["texts"] != bar["texts"]: + raise SystemExit("FAIL: --ref must be the float32 reference and --bar the bfloat16 one, of the same texts") + check_pins("--ref", ref) + check_pins("--bar", bar) + texts = ref["texts"] + n, width = len(texts), len(ref["vectors"][0]) + bar_worst = max(1 - cosine(b, r) for b, r in zip(bar["vectors"], ref["vectors"])) + limit = min(args.max_distance, args.bar_factor * bar_worst) + print(f"bar: transformers bf16 worst distance {bar_worst:.2e}; limit {limit:.2e}") + enc = {"encoding_format": "base64"} + batch = vectors(post(args.url, {"input": texts, **enc}), n, width) + alone = [vectors(post(args.url, {"input": t, **enc}), 1, width)[0] for t in texts] + ids = vectors(post(args.url, {"input": ref["ids"], **enc}), n, width) + ok = True + for i in range(n): + v = floats(batch[i]) + dist = 1 - cosine(v, ref["vectors"][i]) + norm = math.sqrt(sum(x * x for x in v)) + same, by_ids = batch[i] == alone[i], batch[i] == ids[i] + print(f"input {i}: distance {dist:.2e}, norm {norm:.6f}, batch == alone {same}, ids == text {by_ids}") + ok &= dist <= limit and same and by_ids and abs(norm - 1) < 1e-4 + vs = [floats(b) for b in batch] + scores = [[sum(x * y for x, y in zip(vs[q], vs[2 + d])) for d in range(2)] for q in range(2)] + worst = max(abs(scores[q][d] - CARD_SCORES[q][d]) for q in range(2) for d in range(2)) + print("scores", scores, "card", CARD_SCORES, f"max diff {worst:.4f}") + ok &= worst <= args.max_score_diff + print("PASS" if ok else "FAIL") + sys.exit(0 if ok else 1) + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__, epilog=USAGE, formatter_class=argparse.RawDescriptionHelpFormatter) + sub = p.add_subparsers(dest="cmd", required=True) + so = sub.add_parser("sources") + so.add_argument("--model", required=True, type=Path) + r = sub.add_parser("ref") + r.add_argument("--model", required=True, type=Path) + r.add_argument("--dtype", default="float32", choices=["float32", "bfloat16"]) + r.add_argument("--device", default="cpu") + r.add_argument("--out", required=True) + c = sub.add_parser("check") + c.add_argument("--ref", required=True, help="float32 reference JSON (the target)") + c.add_argument("--bar", required=True, help="bfloat16 transformers reference JSON (the bar)") + c.add_argument("--bar-factor", type=float, default=2.0) + c.add_argument("--url", default="http://127.0.0.1:8080") + c.add_argument("--max-distance", type=float, default=0.001) + c.add_argument("--max-score-diff", type=float, default=0.01) + args = p.parse_args() + {"sources": sources, "ref": reference, "check": check}[args.cmd](args) + + +if __name__ == "__main__": + main() diff --git a/tools/zig/test_oracle_qwen3_embedding.py b/tools/zig/test_oracle_qwen3_embedding.py new file mode 100644 index 0000000000..a58acb52f0 --- /dev/null +++ b/tools/zig/test_oracle_qwen3_embedding.py @@ -0,0 +1,43 @@ +"""The Qwen3-Embedding oracle refuses references not made under its pins, and checks pins before model code.""" +import copy +import inspect +import unittest + +import oracle_qwen3_embedding as oracle + + +def made() -> dict: + return {"pins": copy.deepcopy(oracle.expected_pins()) | {"torch": "any"}} + + +class PinRecord(unittest.TestCase): + def test_a_reference_made_under_the_pins_passes(self): + oracle.check_pins("--ref", made()) + + def test_a_changed_or_missing_record_is_refused(self): + changes = [ + lambda p: p.update(revision="0" * 40), + lambda p: p["checkpoint"].update({"tokenizer.json": "0" * 64}), + lambda p: p.update(weights={}), + lambda p: p.update(library={}), + lambda p: p.update(transformers="5.22.0"), + lambda p: p.update(transformers_tree="0" * 64), + lambda p: p.pop("checkpoint"), + ] + for change in changes: + ref = made() + change(ref["pins"]) + with self.assertRaises(SystemExit): + oracle.check_pins("--ref", ref) + with self.assertRaises(SystemExit): + oracle.check_pins("--ref", {}) + + def test_pins_are_checked_before_torch_or_transformers_is_imported(self): + body = inspect.getsource(oracle.reference) + self.assertLess(body.index("pinned(args.model)"), body.index("import torch")) + self.assertLess(body.index("pinned(args.model)"), body.index("from transformers")) + self.assertNotIn("transformers", {m.split(".")[0] for m in oracle.__dict__ if not m.startswith("_")}) + + +if __name__ == "__main__": + unittest.main() From b4b3432ab4a8aee7585f9d12ecbd2e57a25709bc Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Mon, 2 Nov 2026 23:54:22 +0000 Subject: [PATCH 11/24] fix(qwen3): keep the checkpoint in the page cache while it loads Loading Qwen3-Embedding-8B took about 400 s, warm cache or not. The loader released the checkpoint after every layer, and that release also drops each whole shard from the page cache, so every later layer's upload faulted its pages back in from disk. The loader now keeps the mapping until the checkpoint closes at the end of the load; the same server loads in 4 s and serves the same vectors. --- zig/src/families/qwen3/model.zig | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/zig/src/families/qwen3/model.zig b/zig/src/families/qwen3/model.zig index f504071ee6..d299b66621 100644 --- a/zig/src/families/qwen3/model.zig +++ b/zig/src/families/qwen3/model.zig @@ -51,12 +51,11 @@ pub const Model = struct { var t: Table = .{ .ck = &ck }; m.embed = try projection(scratch.allocator(), &t, &up, "embed_tokens", spec.vocab, spec.hidden); m.final_norm = try norm(scratch.allocator(), &t, &up, "norm.weight", spec.hidden); - ck.release(); m.layers = try gpa.alloc(Layer, spec.layers); for (m.layers, 0..) |*layer, i| { _ = scratch.reset(.retain_capacity); + // no ck.release() between layers: its fadvise would drop the shards from the page cache mid-load layer.* = try readLayer(scratch.allocator(), &t, &up, spec, i); - ck.release(); } return m; } From 77119bf0f4597dd7dc3c04f632fedd86ce50aaeb Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Mon, 2 Nov 2026 23:58:52 +0000 Subject: [PATCH 12/24] test(hip): register the BF16 dense GEMM with the kernel contract Every kernel must pass the contract's gates before it lands, and the BF16 dense GEMM had only its unit test. It now runs gates 1 to 4 for fp32 and BF16 output at Qwen3-8B's projection shapes, on BF16 builds, and gate 5 records its compiled usage on gfx1100, gfx1151 and gfx1201. A record can mark a chip not served, as gfx1030's FP16 build keeps only the kernel's trap body; the audits skip such a chip for a launch, never for a planted control. --- zig/src/core/contract/hip_records.zig | 15 ++- zig/src/core/contract/hip_records_dense.zig | 29 +++++ zig/src/core/contract/limits.zig | 2 + zig/tests/hip/contract_dense.zig | 114 ++++++++++++++++++++ zig/tests/hip/contract_limits.zig | 10 ++ zig/tests/hip/contract_main.zig | 8 ++ 6 files changed, 176 insertions(+), 2 deletions(-) create mode 100644 zig/src/core/contract/hip_records_dense.zig create mode 100644 zig/tests/hip/contract_dense.zig diff --git a/zig/src/core/contract/hip_records.zig b/zig/src/core/contract/hip_records.zig index 8784eafcd1..1bf4862e11 100644 --- a/zig/src/core/contract/hip_records.zig +++ b/zig/src/core/contract/hip_records.zig @@ -508,7 +508,8 @@ const own_records = [_]Record{ pub const records = own_records ++ @import("hip_records_stream.zig").records ++ @import("hip_records_router.zig").records ++ @import("hip_records_glm_ocr.zig").records ++ - @import("hip_records_item.zig").records; + @import("hip_records_item.zig").records ++ + @import("hip_records_dense.zig").records; /// Gate 5's planted HIP controls: a launch past its __launch_bounds__, and 72 KiB of LDS where a workgroup has 64. pub const controls = [_]Record{ @@ -549,7 +550,9 @@ pub fn check(a: std.mem.Allocator, list: []const Record, quiet: bool) !usize { defer out.deinit(a); var n: usize = 0; for (list) |r| for (limits.rdna, r.usage) |t, u| { - if (u.vgprs == 0) { + if (!u.served) { + if (!quiet) std.debug.print("N/A {s} on {s}: not served there\n", .{ r.launch.name, t.name }); + } else if (u.vgprs == 0) { try out.print(a, "{s} on {s}: no compiled usage recorded\n", .{ r.launch.name, t.name }); n += 1; } else n += try limits.hip(r.launch, u, t, &out, a); @@ -576,6 +579,10 @@ pub fn controlProblems(a: std.mem.Allocator, r: Record, quiet: bool) !usize { return n; } +test { + _ = @import("hip_records_dense.zig"); +} + test "every registered HIP launch fits every RDNA chip we serve, and each planted control breaks one" { const a = std.testing.allocator; try std.testing.expectEqual(@as(usize, 0), try check(a, &records, false)); @@ -599,6 +606,10 @@ test "a control cannot fit or lack compiled usage on any one chip" { var missing = r; missing.usage[i].vgprs = 0; try std.testing.expectEqual(@as(usize, 1), try controlProblems(a, missing, true)); + // a control is never excused as unserved + var unserved = r; + unserved.usage[i] = .{ .vgprs = 0, .served = false }; + try std.testing.expectEqual(@as(usize, 1), try controlProblems(a, unserved, true)); } } } diff --git a/zig/src/core/contract/hip_records_dense.zig b/zig/src/core/contract/hip_records_dense.zig new file mode 100644 index 0000000000..26968332ec --- /dev/null +++ b/zig/src/core/contract/hip_records_dense.zig @@ -0,0 +1,29 @@ +//! Gate 5 for the BF16 dense GEMM: gfx1030's FP16 build keeps only its trap body, so the record does not serve it. +const std = @import("std"); +const limits = @import("limits.zig"); +const hip_records = @import("hip_records.zig"); + +const Record = hip_records.Record; +const check = hip_records.check; + +pub const records = [_]Record{ + .{ + .launch = .{ .name = "dense gemm dense_gemm_block", .symbol = "_ZN2tf4rocm16dense_gemm_blockINS0_7DotBF16ELi8EEEvNS0_6AffineE", .image = "affine_tiles", .threads = 256 }, + .usage = .{ + .{ .vgprs = 0, .served = false }, + .{ .vgprs = 254, .sgprs = 34, .lds = 42496, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 254, .sgprs = 34, .lds = 42496, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 254, .sgprs = 34, .lds = 42496, .private = 0, .bound = 256, .wave = 32 }, + }, + }, +}; + +test "a chip a record does not serve is skipped; a served chip without usage is still a problem" { + const a = std.testing.allocator; + const launch: limits.HipLaunch = .{ .name = "t", .symbol = "t", .image = "ops", .threads = 256 }; + const one: limits.HipUsage = .{ .vgprs = 16, .bound = 256 }; + const skip: Record = .{ .launch = launch, .usage = .{ .{ .vgprs = 0, .served = false }, one, one, one } }; + try std.testing.expectEqual(@as(usize, 0), try check(a, &.{skip}, true)); + const missing: Record = .{ .launch = launch, .usage = .{ .{ .vgprs = 0 }, one, one, one } }; + try std.testing.expectEqual(@as(usize, 1), try check(a, &.{missing}, true)); +} diff --git a/zig/src/core/contract/limits.zig b/zig/src/core/contract/limits.zig index 3089d01d31..b624135929 100644 --- a/zig/src/core/contract/limits.zig +++ b/zig/src/core/contract/limits.zig @@ -200,6 +200,8 @@ pub const HipLaunch = struct { /// What hipcc made of one kernel for one chip, from the code object's metadata (vgprs 0: none recorded). pub const HipUsage = struct { vgprs: u32, + /// False: the kernel is not launched on this chip (its image keeps a trap body there); no usage is checked. + served: bool = true, sgprs: u32 = 0, lds: u32 = 0, private: u32 = 0, diff --git a/zig/tests/hip/contract_dense.zig b/zig/tests/hip/contract_dense.zig new file mode 100644 index 0000000000..1285218083 --- /dev/null +++ b/zig/tests/hip/contract_dense.zig @@ -0,0 +1,114 @@ +//! dense_gemm_block under the contract on HIP: BF16 x times raw BF16 weights on the dot2 GEMM tile, fp32 and BF16 out. +const std = @import("std"); +const hip = @import("hip"); +const core = @import("core"); +const rig_ = @import("contract_rig.zig"); + +const contract = core.contract; +const c = contract.spec; +const g = contract.gates; +const Dense = contract.dense.Dense; + +/// A group's 64 dot2s in one fp32 chain, then one add a group of 128 into the accumulator. +fn depth(s: c.Shape) usize { + return s.k / 2 + s.k / 128 + 2; +} + +/// Qwen3-8B's projections as float64 runs them: 2,048-row passes, 12,288-wide gate/up, K to 12,288, both tile edges. +const served = [_]c.Shape{ + .{ .rows = 1, .n = 4096, .k = 4096 }, + .{ .rows = 7, .n = 1024, .k = 12288 }, + .{ .rows = 16, .n = 12288, .k = 4096 }, + .{ .rows = 129, .n = 130, .k = 4096 }, + .{ .rows = 300, .n = 256, .k = 4096 }, + .{ .rows = 2048, .n = 128, .k = 4096 }, +}; + +pub fn spec(wide: bool) c.Spec { + return .{ + .name = if (wide) "hip dense_gemm_block (bf16 weights, fp32 out)" else "hip dense_gemm_block (bf16 weights, bf16 out)", + .accum = .f32, + .depth = depth, + .out = if (wide) .f32 else .bf16, + .served = &served, + .tiles = .{ .{ .size = 128 }, .{ .size = 128 }, .{ .size = 12288, .step = 128, .min = 128 } }, + }; +} + +pub const DenseGemm = struct { + rig: rig_.Rig, + l: *const hip.Launcher, + /// fp32 out, or BF16 rounded once in the epilogue. + wide: bool, + d: ?Dense = null, + x: u64 = 0, + w: u64 = 0, + out: u64 = 0, + + pub fn init(a: std.mem.Allocator, rt: *const hip.Runtime, s: hip.Stream, l: *const hip.Launcher, wide: bool) DenseGemm { + return .{ .rig = .{ .a = a, .r = rt, .s = s }, .l = l, .wide = wide }; + } + + pub fn deinit(self: *DenseGemm) void { + if (self.d) |d| d.deinit(); + self.rig.deinit(); + } + + pub fn kernel(self: *DenseGemm) g.Kernel { + return .{ .spec = spec(self.wide), .ctx = self, .vt = &vt }; + } + + const vt = g.hooks(DenseGemm); + + /// x from the dense family; W as BF16 in [-1, 1), or the stress fill's mirrored integers; FP16 builds serve none. + pub fn prepare(self: *DenseGemm, shape: c.Shape, fill: g.Fill, seed: u64) !bool { + if (shape.rows == 0 or shape.n == 0 or shape.k == 0 or shape.k % 128 != 0) return false; + if (self.d) |d| d.deinit(); + self.d = null; + self.rig.reset(); + const a = self.rig.a; + const d = try Dense.init(a, shape, 1); + self.d = d; + d.fillX(fill, seed, 1); + const bits = try a.alloc(u16, d.w.len); + defer a.free(bits); + var prng = std.Random.DefaultPrng.init(seed ^ 0xD3); + const rng = prng.random(); + const k = shape.k; + for (bits) |*b| b.* = c.toBf16(rng.float(f32) * 2 - 1); + if (fill == .stress) for (0..shape.n) |j| for (0..k / 2) |i| { + const v: f32 = @floatFromInt(1 + rng.uintLessThan(u32, 7)); + bits[j * k + i] = c.toBf16(v); + bits[j * k + k / 2 + i] = c.toBf16(-v); + }; + if (fill == .nonfinite and shape.n > 5) { + bits[1 * k + k / 4] = c.toBf16(std.math.nan(f32)); + bits[5 * k] = c.toBf16(std.math.inf(f32)); + } + for (d.w, bits) |*w, b| w.* = c.fromBf16(b); + self.x = try self.rig.input(std.mem.sliceAsBytes(d.x)); + self.w = try self.rig.input(std.mem.sliceAsBytes(bits)); + self.out = try self.rig.output(d.outputs() * @as(usize, if (self.wide) 4 else 2)); + return true; + } + + pub fn run(self: *DenseGemm) !void { + const s = self.d.?.shape; + const out32: u64 = if (self.wide) self.out else 0; + const out16: u64 = if (self.wide) 0 else self.out; + try self.l.dense.gemm(self.l.r, self.x, self.w, out32, out16, @intCast(s.rows), @intCast(s.n), @intCast(s.k), self.rig.s.handle); + try self.rig.s.synchronize(); + } + + pub fn count(self: *DenseGemm) usize { + return self.d.?.outputs(); + } + + pub fn reference(self: *DenseGemm, want: []f64, mag: []f64) !void { + try self.d.?.reference(want, mag); + } + + pub fn simulate(self: *DenseGemm, acc: c.Accum, order: g.Order, out: []f64) !void { + try self.d.?.simulate(acc, order, out); + } +}; diff --git a/zig/tests/hip/contract_limits.zig b/zig/tests/hip/contract_limits.zig index 23e11b25ee..56ae3c123b 100644 --- a/zig/tests/hip/contract_limits.zig +++ b/zig/tests/hip/contract_limits.zig @@ -72,6 +72,11 @@ pub fn live(a: std.mem.Allocator, r: *const hip.Runtime, ctx: hip.Context, arch: std.debug.print(say, .{ word, arch, d[0], got, d[2] }); } for (all) |rec| { + // a planted control is checked on every chip: only a registered launch may be unserved somewhere + if (!rec.usage[at].served and !isControl(rec)) { + std.debug.print("N/A limits live {s}: not served on {s}\n", .{ rec.launch.name, arch }); + continue; + } const bytes = try image(rec.launch.image); const listed = try hsaco.kernels(a, bytes); defer a.free(listed); @@ -108,6 +113,11 @@ pub fn objects(a: std.mem.Allocator, io: std.Io, dir: []const u8) !usize { var line: std.ArrayList(u8) = .empty; defer line.deinit(a); for (limits.rdna, 0..) |t, i| { + if (!rec.usage[i].served and !control) { + std.debug.print("N/A limits {s} {s}: not served there\n", .{ t.name, rec.launch.name }); + try line.appendSlice(a, ".{ .vgprs = 0, .served = false }, "); + continue; + } const path = try file(a, dir, t.name, rec.launch.image); defer a.free(path); const bytes = std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 28)) catch |err| { diff --git a/zig/tests/hip/contract_main.zig b/zig/tests/hip/contract_main.zig index fce0236f1f..e2b159baf2 100644 --- a/zig/tests/hip/contract_main.zig +++ b/zig/tests/hip/contract_main.zig @@ -19,6 +19,7 @@ const kv_k = @import("contract_kv_write.zig"); const argmax_k = @import("contract_argmax.zig"); const pair_k = @import("contract_relu2_pair.zig"); const heads_rows_k = @import("contract_heads_rows.zig"); +const dense_k = @import("contract_dense.zig"); const gate_k = @import("contract_moe_gate.zig"); const router_k = @import("contract_router.zig"); const kit = @import("contract_kit.zig"); @@ -156,6 +157,13 @@ fn run(a: std.mem.Allocator, io: std.Io) !usize { var heads_rows = heads_rows_k.HeadsRows.init(a, &r, s, &l); defer heads_rows.deinit(); try kernels.append(a, heads_rows.kernel()); + var dense: [2]dense_k.DenseGemm = undefined; + for (&dense, [_]bool{ true, false }) |*d, fp32_out| d.* = .init(a, &r, s, &l, fp32_out); + defer for (&dense) |*d| d.deinit(); + // the BF16 dense GEMM runs on BF16 builds only; an FP16 build keeps its trap body and serves no shape of it + if (l.dense.bf16) { + for (&dense) |*d| try kernels.append(a, d.kernel()); + } else std.debug.print("N/A hip dense_gemm_block: not served on this FP16 build\n", .{}); var gates: [3]gate_k.Gate = undefined; for (&gates, std.enums.values(gate_k.Which)) |*m, w| m.* = .init(a, &r, s, &l, w); defer for (&gates) |*m| m.deinit(); From 266ed67aabccdf300ccc95cb227e63a355354461 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 3 Nov 2026 00:31:08 +0000 Subject: [PATCH 13/24] test(qwen3): refuse embedding references of other inputs or values The embedding gate's limit comes from the BF16 bar's distance to the fp32 target, so a bar of other inputs or a non-finite vector could loosen it or make it pass on NaN. The oracle checker and the GPU test now refuse references that differ in dtype, texts, token ids or pins, or hold a ragged, zero or non-finite vector, and fail on any non-finite distance before comparing. --- tools/zig/oracle_qwen3_embedding.py | 23 +++++++- tools/zig/test_oracle_qwen3_embedding.py | 8 +++ zig/src/families/qwen3/forward_test.zig | 75 +++++++++++++++++++++++- 3 files changed, 103 insertions(+), 3 deletions(-) diff --git a/tools/zig/oracle_qwen3_embedding.py b/tools/zig/oracle_qwen3_embedding.py index 5e4a16f2e4..fd634e20a4 100644 --- a/tools/zig/oracle_qwen3_embedding.py +++ b/tools/zig/oracle_qwen3_embedding.py @@ -206,6 +206,20 @@ def cosine(a, b) -> float: return dot / math.sqrt(sum(x * x for x in a) * sum(y * y for y in b)) +def check_vectors(name: str, made: dict, count: int) -> int: + """The reference's width, after checking it holds `count` finite, nonzero vectors of one width.""" + vectors = made.get("vectors") + if not isinstance(vectors, list) or len(vectors) != count or not vectors: + raise SystemExit(f"FAIL: {name} does not hold {count} vectors") + width = len(vectors[0]) + for i, v in enumerate(vectors): + if len(v) != width or not all(isinstance(x, (int, float)) and math.isfinite(x) for x in v): + raise SystemExit(f"FAIL: {name} vector {i} is not {width} finite values") + if sum(x * x for x in v) == 0: + raise SystemExit(f"FAIL: {name} vector {i} is zero") + return width + + def check(args) -> None: ref = json.load(open(args.ref)) bar = json.load(open(args.bar)) @@ -214,8 +228,13 @@ def check(args) -> None: check_pins("--ref", ref) check_pins("--bar", bar) texts = ref["texts"] - n, width = len(texts), len(ref["vectors"][0]) + n = len(texts) + width = check_vectors("--ref", ref, n) + if check_vectors("--bar", bar, n) != width or ref.get("ids") != bar.get("ids"): + raise SystemExit("FAIL: --ref and --bar differ in width or token ids") bar_worst = max(1 - cosine(b, r) for b, r in zip(bar["vectors"], ref["vectors"])) + if not math.isfinite(bar_worst): + raise SystemExit("FAIL: the bar's distance is not finite") limit = min(args.max_distance, args.bar_factor * bar_worst) print(f"bar: transformers bf16 worst distance {bar_worst:.2e}; limit {limit:.2e}") enc = {"encoding_format": "base64"} @@ -229,7 +248,7 @@ def check(args) -> None: norm = math.sqrt(sum(x * x for x in v)) same, by_ids = batch[i] == alone[i], batch[i] == ids[i] print(f"input {i}: distance {dist:.2e}, norm {norm:.6f}, batch == alone {same}, ids == text {by_ids}") - ok &= dist <= limit and same and by_ids and abs(norm - 1) < 1e-4 + ok &= math.isfinite(dist) and dist <= limit and same and by_ids and abs(norm - 1) < 1e-4 vs = [floats(b) for b in batch] scores = [[sum(x * y for x, y in zip(vs[q], vs[2 + d])) for d in range(2)] for q in range(2)] worst = max(abs(scores[q][d] - CARD_SCORES[q][d]) for q in range(2) for d in range(2)) diff --git a/tools/zig/test_oracle_qwen3_embedding.py b/tools/zig/test_oracle_qwen3_embedding.py index a58acb52f0..3f4b7d459a 100644 --- a/tools/zig/test_oracle_qwen3_embedding.py +++ b/tools/zig/test_oracle_qwen3_embedding.py @@ -32,6 +32,14 @@ def test_a_changed_or_missing_record_is_refused(self): with self.assertRaises(SystemExit): oracle.check_pins("--ref", {}) + def test_a_reference_with_a_nonfinite_zero_or_ragged_vector_is_refused(self): + good = {"vectors": [[0.6, 0.8], [1.0, 0.0]]} + self.assertEqual(2, oracle.check_vectors("--bar", good, 2)) + for vectors in ([[float("nan"), 1.0], [1.0, 0.0]], [[0.0, 0.0], [1.0, 0.0]], [[0.6, 0.8], [1.0]], + [[0.6, 0.8]], [[float("inf"), 0.0], [1.0, 0.0]]): + with self.assertRaises(SystemExit): + oracle.check_vectors("--bar", {"vectors": vectors}, 2) + def test_pins_are_checked_before_torch_or_transformers_is_imported(self): body = inspect.getsource(oracle.reference) self.assertLess(body.index("pinned(args.model)"), body.index("import torch")) diff --git a/zig/src/families/qwen3/forward_test.zig b/zig/src/families/qwen3/forward_test.zig index edcec15fab..41027beae4 100644 --- a/zig/src/families/qwen3/forward_test.zig +++ b/zig/src/families/qwen3/forward_test.zig @@ -100,6 +100,77 @@ fn cosine(x: []const f64, y: []const f64) f64 { return xy / @sqrt(xx * yy); } +/// The pins tools/zig/oracle_qwen3_embedding.py writes into each reference: the checkpoint and the Transformers tree. +const embed_revision = "1d8ad4ca9b3dd8059ad90a75d4983776a23d44af"; +const transformers_tree = "b08cead6d1b841ef9723106a81a92ec37d67aff3e43a198e56b7dc76538d377b"; + +fn text(v: std.json.Value, key: []const u8) ?[]const u8 { + const x = (if (v == .object) v.object.get(key) else null) orelse return null; + return if (x == .string) x.string else null; +} + +/// Refuses, before any GPU work, references that are not the fp32 target and BF16 bar of the same pinned inputs. +fn checkReferences(ref: std.json.Value, bar: std.json.Value) !void { + try t.expectEqualStrings("float32", text(ref, "dtype") orelse ""); + try t.expectEqualStrings("bfloat16", text(bar, "dtype") orelse ""); + for ([_]std.json.Value{ ref, bar }) |r| { + const pins = r.object.get("pins") orelse return error.UnpinnedReference; + try t.expectEqualStrings(embed_revision, text(pins, "revision") orelse ""); + try t.expectEqualStrings(transformers_tree, text(pins, "transformers_tree") orelse ""); + } + for ([_][]const u8{ "texts", "ids" }) |key| { + const x = ref.object.get(key) orelse return error.IncompleteReference; + const y = bar.object.get(key) orelse return error.IncompleteReference; + try t.expect(std.meta.activeTag(x) == .array and std.meta.activeTag(y) == .array and x.array.items.len == y.array.items.len); + for (x.array.items, y.array.items) |p, q| try t.expect(jsonEqual(p, q)); + } + for ([_]std.json.Value{ ref, bar }) |r| for (r.object.get("vectors").?.array.items) |v| { + var square: f64 = 0; + for (v.array.items) |x| { + const f = number(x); + try t.expect(std.math.isFinite(f)); + square += f * f; + } + try t.expect(square > 0); + }; +} + +fn jsonEqual(p: std.json.Value, q: std.json.Value) bool { + if (std.meta.activeTag(p) != std.meta.activeTag(q)) return false; + return switch (p) { + .string => |x| std.mem.eql(u8, x, q.string), + .integer => |x| x == q.integer, + .array => |x| x.items.len == q.array.items.len and for (x.items, q.array.items) |a, b| { + if (!jsonEqual(a, b)) break false; + } else true, + else => false, + }; +} + +test "Qwen3-Embedding references: a bar of other inputs, an unpinned one or a non-finite vector is refused" { + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const good = + \\{"dtype": "float32", "pins": {"revision": "1d8ad4ca9b3dd8059ad90a75d4983776a23d44af", + \\ "transformers_tree": "b08cead6d1b841ef9723106a81a92ec37d67aff3e43a198e56b7dc76538d377b"}, + \\ "texts": ["a"], "ids": [[1, 2]], "vectors": [[0.6, 0.8]]} + ; + const ref = try std.json.parseFromSliceLeaky(std.json.Value, a, good, .{}); + const bar_text = try std.mem.replaceOwned(u8, a, good, "float32", "bfloat16"); + try checkReferences(ref, try std.json.parseFromSliceLeaky(std.json.Value, a, bar_text, .{})); + const bad = [_][2][]const u8{ + .{ "[[1, 2]]", "[[1, 3]]" }, + .{ "\"texts\": [\"a\"]", "\"texts\": [\"b\"]" }, + .{ "1d8ad4ca9b3dd8059ad90a75d4983776a23d44af", "0000000000000000000000000000000000000000" }, + .{ "[[0.6, 0.8]]", "[[0.0, 0.0]]" }, + }; + for (bad) |swap| { + const other = try std.mem.replaceOwned(u8, a, bar_text, swap[0], swap[1]); + try t.expect(std.meta.isError(checkReferences(ref, try std.json.parseFromSliceLeaky(std.json.Value, a, other, .{})))); + } +} + test "Qwen3-Embedding: the card's inputs batched equal them alone, and sit within the BF16 bar of the fp32 vectors" { var arena = std.heap.ArenaAllocator.init(gpa); defer arena.deinit(); @@ -107,6 +178,7 @@ test "Qwen3-Embedding: the card's inputs batched equal them alone, and sit withi const dir = try env(a, "TF_QWEN3_EMBED_DIR") orelse return error.SkipZigTest; const ref = try readJson(a, try env(a, "TF_QWEN3_EMBED_REF") orelse return error.SkipZigTest); const bar = try readJson(a, try env(a, "TF_QWEN3_EMBED_BAR") orelse return error.SkipZigTest); + try checkReferences(ref, bar); var g: Gpu = undefined; try g.open(4 << 30); defer g.close(); @@ -143,6 +215,7 @@ test "Qwen3-Embedding: the card's inputs batched equal them alone, and sit withi } bar_worst = @max(bar_worst, 1 - cosine(bv, rv)); } + try t.expect(std.math.isFinite(bar_worst)); const limit = @min(0.001, 2 * bar_worst); var begin: u32 = 0; for (ends, 0..) |e, i| { @@ -157,7 +230,7 @@ test "Qwen3-Embedding: the card's inputs batched equal them alone, and sit withi } const distance = 1 - cosine(got, want); std.debug.print("embedding input {d}: cosine distance {e:.3} (limit {e:.3})\n", .{ i, distance, limit }); - try t.expect(distance <= limit); + try t.expect(std.math.isFinite(distance) and distance <= limit); begin = e; } } From 5db9b253e14bf304e2ac9ed3fe869f4b50fbce83 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 3 Nov 2026 00:31:08 +0000 Subject: [PATCH 14/24] test(hip): bound the dense GEMM's tests by its own rounding chain The unit test and the kernel contract bounded the BF16 dense GEMM's error with different rounding chains, and the contract's ignored the fold every 128 inputs, allowing about 28 times more error at K 12288. Both now take one chain from the launcher: a group's dot2s, then one add a group. Every contract shape still passes within it. --- zig/src/hip/kernel_tests/dense.zig | 3 +-- zig/src/hip/launches/dense.zig | 5 +++++ zig/tests/hip/contract_dense.zig | 4 ++-- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/zig/src/hip/kernel_tests/dense.zig b/zig/src/hip/kernel_tests/dense.zig index 2d86400116..6d93a2d2aa 100644 --- a/zig/src/hip/kernel_tests/dense.zig +++ b/zig/src/hip/kernel_tests/dense.zig @@ -58,8 +58,7 @@ test "dense GEMM: BF16 weights against float64, fp32 within the dot2 chain's bou defer gpa.free(g32); const g16 = try rig.download(u16, y16); defer gpa.free(g16); - // a group's 64 dot2s (two roundings each at most), then one add a group: the longest rounding chain of a sum - const depth: f64 = 128 + @as(f64, @floatFromInt(c.k / 128)) + 2; + const depth: f64 = @floatFromInt(@import("../launches/dense.zig").Kernels.chain(c.k)); var worst: f64 = 0; for (0..c.m) |r| for (0..c.n) |j| { var want: f64 = 0; diff --git a/zig/src/hip/launches/dense.zig b/zig/src/hip/launches/dense.zig index ec771a0fd0..ecbb867cc2 100644 --- a/zig/src/hip/launches/dense.zig +++ b/zig/src/hip/launches/dense.zig @@ -67,6 +67,11 @@ pub const Kernels = struct { try hl.launch(k.tile, .{ .grid = .{ .x = cdiv(n), .y = cdiv(m) }, .block = .{ .x = 256 } }, .{ .r = r, .handle = s }, &args); } + /// The longest rounding chain of an output: a group's 64 dot2s (two roundings each at most), one add a group. + pub fn chain(k: usize) usize { + return 2 * (@as(usize, @intCast(group)) / 2) + k / @as(usize, @intCast(group)) + 2; + } + fn cdiv(v: c_int) c_uint { return @intCast(@divTrunc(v + side - 1, side)); } diff --git a/zig/tests/hip/contract_dense.zig b/zig/tests/hip/contract_dense.zig index 1285218083..28a4f7e8a6 100644 --- a/zig/tests/hip/contract_dense.zig +++ b/zig/tests/hip/contract_dense.zig @@ -9,9 +9,9 @@ const c = contract.spec; const g = contract.gates; const Dense = contract.dense.Dense; -/// A group's 64 dot2s in one fp32 chain, then one add a group of 128 into the accumulator. +/// The kernel's own chain (launches/dense.zig): a group's dot2s, then one add a group into the accumulator. fn depth(s: c.Shape) usize { - return s.k / 2 + s.k / 128 + 2; + return @FieldType(hip.Launcher, "dense").chain(s.k); } /// Qwen3-8B's projections as float64 runs them: 2,048-row passes, 12,288-wide gate/up, K to 12,288, both tile edges. From c1b0cf8ae41fb9d88f32c142243c0d843c12c0a8 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 3 Nov 2026 00:37:41 +0000 Subject: [PATCH 15/24] test(qwen3): require one complete pin record in both GPU references The GPU gate checked only the revision and the Transformers tree of its references, so a record without its checkpoint, weight or library hashes still passed where the oracle's checker refuses it. Both references must now carry the complete record, the same in each, and a refused reference fails without printing a diff. --- zig/src/families/qwen3/forward_test.zig | 49 ++++++++++++++++++++----- 1 file changed, 39 insertions(+), 10 deletions(-) diff --git a/zig/src/families/qwen3/forward_test.zig b/zig/src/families/qwen3/forward_test.zig index 41027beae4..af46244c9e 100644 --- a/zig/src/families/qwen3/forward_test.zig +++ b/zig/src/families/qwen3/forward_test.zig @@ -109,29 +109,47 @@ fn text(v: std.json.Value, key: []const u8) ?[]const u8 { return if (x == .string) x.string else null; } +fn same(expected: []const u8, got: ?[]const u8) !void { + if (!std.mem.eql(u8, expected, got orelse "")) return error.ReferenceMismatch; +} + +/// A silent check: a refused reference is an error, without the expect helpers' diff on stderr. +fn holds(ok: bool) !void { + if (!ok) return error.ReferenceMismatch; +} + /// Refuses, before any GPU work, references that are not the fp32 target and BF16 bar of the same pinned inputs. fn checkReferences(ref: std.json.Value, bar: std.json.Value) !void { - try t.expectEqualStrings("float32", text(ref, "dtype") orelse ""); - try t.expectEqualStrings("bfloat16", text(bar, "dtype") orelse ""); + try same("float32", text(ref, "dtype")); + try same("bfloat16", text(bar, "dtype")); + // the oracle checks every hash of the pin record; here both references carry one complete, identical record for ([_]std.json.Value{ ref, bar }) |r| { const pins = r.object.get("pins") orelse return error.UnpinnedReference; - try t.expectEqualStrings(embed_revision, text(pins, "revision") orelse ""); - try t.expectEqualStrings(transformers_tree, text(pins, "transformers_tree") orelse ""); + try same("Qwen/Qwen3-Embedding-8B", text(pins, "repo")); + try same(embed_revision, text(pins, "revision")); + try same("5.22.1", text(pins, "transformers")); + try same(transformers_tree, text(pins, "transformers_tree")); + for ([_][]const u8{ "checkpoint", "weights", "library" }) |key| { + const files = pins.object.get(key) orelse return error.UnpinnedReference; + try holds(files == .object and files.object.count() > 0); + for (files.object.values()) |hash| try holds(hash == .string and hash.string.len == 64); + } } + try holds(jsonEqual(ref.object.get("pins").?, bar.object.get("pins").?)); for ([_][]const u8{ "texts", "ids" }) |key| { const x = ref.object.get(key) orelse return error.IncompleteReference; const y = bar.object.get(key) orelse return error.IncompleteReference; - try t.expect(std.meta.activeTag(x) == .array and std.meta.activeTag(y) == .array and x.array.items.len == y.array.items.len); - for (x.array.items, y.array.items) |p, q| try t.expect(jsonEqual(p, q)); + try holds(std.meta.activeTag(x) == .array and std.meta.activeTag(y) == .array and x.array.items.len == y.array.items.len); + for (x.array.items, y.array.items) |p, q| try holds(jsonEqual(p, q)); } for ([_]std.json.Value{ ref, bar }) |r| for (r.object.get("vectors").?.array.items) |v| { var square: f64 = 0; for (v.array.items) |x| { const f = number(x); - try t.expect(std.math.isFinite(f)); + try holds(std.math.isFinite(f)); square += f * f; } - try t.expect(square > 0); + try holds(square > 0); }; } @@ -143,6 +161,10 @@ fn jsonEqual(p: std.json.Value, q: std.json.Value) bool { .array => |x| x.items.len == q.array.items.len and for (x.items, q.array.items) |a, b| { if (!jsonEqual(a, b)) break false; } else true, + .object => |x| x.count() == q.object.count() and for (x.keys(), x.values()) |k, v| { + const w = q.object.get(k) orelse break false; + if (!jsonEqual(v, w)) break false; + } else true, else => false, }; } @@ -152,8 +174,12 @@ test "Qwen3-Embedding references: a bar of other inputs, an unpinned one or a no defer arena.deinit(); const a = arena.allocator(); const good = - \\{"dtype": "float32", "pins": {"revision": "1d8ad4ca9b3dd8059ad90a75d4983776a23d44af", - \\ "transformers_tree": "b08cead6d1b841ef9723106a81a92ec37d67aff3e43a198e56b7dc76538d377b"}, + \\{"dtype": "float32", "pins": {"repo": "Qwen/Qwen3-Embedding-8B", + \\ "revision": "1d8ad4ca9b3dd8059ad90a75d4983776a23d44af", "transformers": "5.22.1", + \\ "transformers_tree": "b08cead6d1b841ef9723106a81a92ec37d67aff3e43a198e56b7dc76538d377b", + \\ "checkpoint": {"config.json": "dcccf7c7890c8debc4af8dace8d6acd8dd50bd56d50d4d8b60529381ba3de3d7"}, + \\ "weights": {"model-00001-of-00004.safetensors": "99b343597fe840706146144699a8b9188dd3387e43eb61faf0231b70b249d451"}, + \\ "library": {"models/qwen3/modeling_qwen3.py": "63a6213c15d6fb04e58a1c64fd05a4cdfb49476443fbeef524f0e0486043824b"}}, \\ "texts": ["a"], "ids": [[1, 2]], "vectors": [[0.6, 0.8]]} ; const ref = try std.json.parseFromSliceLeaky(std.json.Value, a, good, .{}); @@ -164,6 +190,9 @@ test "Qwen3-Embedding references: a bar of other inputs, an unpinned one or a no .{ "\"texts\": [\"a\"]", "\"texts\": [\"b\"]" }, .{ "1d8ad4ca9b3dd8059ad90a75d4983776a23d44af", "0000000000000000000000000000000000000000" }, .{ "[[0.6, 0.8]]", "[[0.0, 0.0]]" }, + .{ "99b343597fe840706146144699a8b9188dd3387e43eb61faf0231b70b249d451", "00b343597fe840706146144699a8b9188dd3387e43eb61faf0231b70b249d451" }, + .{ "\"weights\": {\"model-00001-of-00004.safetensors\": \"99b343597fe840706146144699a8b9188dd3387e43eb61faf0231b70b249d451\"},", "" }, + .{ "\"transformers\": \"5.22.1\"", "\"transformers\": \"5.22.0\"" }, }; for (bad) |swap| { const other = try std.mem.replaceOwned(u8, a, bar_text, swap[0], swap[1]); From 1b6d90931dc726ad24ba7bbff4a6a2de64a52481 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 3 Nov 2026 00:42:12 +0000 Subject: [PATCH 16/24] test(qwen3): bind the GPU gate's references to the files it loads Two references could carry the same wrong pin maps and still pass the GPU gate, which never looked at the checkpoint it loads. The test now hashes every pinned checkpoint and weight file in the directory it loads and requires every shard of the index to be pinned, before any GPU work; a host test refuses a wrong hash and an unpinned shard. --- zig/src/families/qwen3/forward_test.zig | 53 +++++++++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/zig/src/families/qwen3/forward_test.zig b/zig/src/families/qwen3/forward_test.zig index af46244c9e..116cd22b61 100644 --- a/zig/src/families/qwen3/forward_test.zig +++ b/zig/src/families/qwen3/forward_test.zig @@ -153,6 +153,39 @@ fn checkReferences(ref: std.json.Value, bar: std.json.Value) !void { }; } +/// The SHA256 of a file, as lowercase hex. +fn fileSha256(io: std.Io, path: []const u8) ![64]u8 { + const f = try std.Io.Dir.cwd().openFile(io, path, .{}); + defer f.close(io); + const size = (try f.stat(io)).size; + var h = std.crypto.hash.sha2.Sha256.init(.{}); + var buf: [1 << 22]u8 = undefined; + var at: u64 = 0; + while (at < size) { + const n = try f.readPositionalAll(io, buf[0..@intCast(@min(buf.len, size - at))], at); + if (n == 0) return error.ShortRead; + h.update(buf[0..n]); + at += n; + } + return std.fmt.bytesToHex(h.finalResult(), .lower); +} + +/// The references' pinned files are the checkpoint the test loads: each one's SHA256, and every indexed shard pinned. +fn checkFiles(a: std.mem.Allocator, io: std.Io, dir: []const u8, pins: std.json.Value) !void { + for ([_][]const u8{ "checkpoint", "weights" }) |key| { + const files = pins.object.get(key).?.object; + for (files.keys(), files.values()) |name, hash| { + const got = try fileSha256(io, try std.fs.path.join(a, &.{ dir, name })); + try holds(std.mem.eql(u8, &got, hash.string)); + } + } + const index = try readJson(a, try std.fs.path.join(a, &.{ dir, "model.safetensors.index.json" })); + const shards = index.object.get("weight_map").?.object; + const weights = pins.object.get("weights").?.object; + for (shards.values()) |file| try holds(file == .string and weights.get(file.string) != null); + try holds(pins.object.get("checkpoint").?.object.get("model.safetensors.index.json") != null); +} + fn jsonEqual(p: std.json.Value, q: std.json.Value) bool { if (std.meta.activeTag(p) != std.meta.activeTag(q)) return false; return switch (p) { @@ -200,6 +233,25 @@ test "Qwen3-Embedding references: a bar of other inputs, an unpinned one or a no } } +test "Qwen3-Embedding references: pins that do not hash to the checkpoint's files are refused, whoever made both" { + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var tmp = t.tmpDir(.{}); + defer tmp.cleanup(); + try tmp.dir.writeFile(t.io, .{ .sub_path = "s.safetensors", .data = "weights" }); + try tmp.dir.writeFile(t.io, .{ .sub_path = "model.safetensors.index.json", .data = "{\"weight_map\": {\"x\": \"s.safetensors\"}}" }); + const dir = try tmp.dir.realPathFileAlloc(t.io, ".", a); + const shard = try fileSha256(t.io, try std.fs.path.join(a, &.{ dir, "s.safetensors" })); + const index = try fileSha256(t.io, try std.fs.path.join(a, &.{ dir, "model.safetensors.index.json" })); + const good = try std.fmt.allocPrint(a, "{{\"checkpoint\": {{\"model.safetensors.index.json\": \"{s}\"}}, \"weights\": {{\"s.safetensors\": \"{s}\"}}}}", .{ &index, &shard }); + try checkFiles(a, t.io, dir, try std.json.parseFromSliceLeaky(std.json.Value, a, good, .{})); + const wrong = try std.mem.replaceOwned(u8, a, good, &shard, "0000000000000000000000000000000000000000000000000000000000000000"); + try t.expect(std.meta.isError(checkFiles(a, t.io, dir, try std.json.parseFromSliceLeaky(std.json.Value, a, wrong, .{})))); + const unpinned = try std.fmt.allocPrint(a, "{{\"checkpoint\": {{\"model.safetensors.index.json\": \"{s}\"}}, \"weights\": {{}}}}", .{&index}); + try t.expect(std.meta.isError(checkFiles(a, t.io, dir, try std.json.parseFromSliceLeaky(std.json.Value, a, unpinned, .{})))); +} + test "Qwen3-Embedding: the card's inputs batched equal them alone, and sit within the BF16 bar of the fp32 vectors" { var arena = std.heap.ArenaAllocator.init(gpa); defer arena.deinit(); @@ -208,6 +260,7 @@ test "Qwen3-Embedding: the card's inputs batched equal them alone, and sit withi const ref = try readJson(a, try env(a, "TF_QWEN3_EMBED_REF") orelse return error.SkipZigTest); const bar = try readJson(a, try env(a, "TF_QWEN3_EMBED_BAR") orelse return error.SkipZigTest); try checkReferences(ref, bar); + try checkFiles(a, t.io, dir, ref.object.get("pins").?); var g: Gpu = undefined; try g.open(4 << 30); defer g.close(); From 73bcf745b5eeef73c48407523adfffbc0d917c34 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 3 Nov 2026 01:00:00 +0000 Subject: [PATCH 17/24] fix(embed): refuse blank inputs and name an input past the context Probing SGLang's Qwen3-Embedding-8B server with the same requests as this one showed two gaps: whitespace-only text was embedded, and one input longer than the context was refused only as an over-full batch. Text that is empty or only Unicode whitespace, as Python's strip sees it, is now refused with its index in a list, and an input past the context is refused with its token count before the batch's own limit applies. --- zig/src/server/embedding_protocol_test.zig | 31 ++++++++++++++++++++++ zig/src/server/embedding_request.zig | 31 +++++++++++++++++++--- 2 files changed, 58 insertions(+), 4 deletions(-) diff --git a/zig/src/server/embedding_protocol_test.zig b/zig/src/server/embedding_protocol_test.zig index c2f175dd9b..af3dba7369 100644 --- a/zig/src/server/embedding_protocol_test.zig +++ b/zig/src/server/embedding_protocol_test.zig @@ -117,3 +117,34 @@ test "embeddings protocol: float and little-endian base64 preserve MRL rows and var nan = [_]f32{ 3, 4, std.math.nan(f32) }; try t.expectError(error.Refused, response.encode(&cx, "m", parsed, &nan)); } + +test "embeddings protocol: whitespace-only text is refused with its index; an over-long input names its length" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + var cx: errors.Cx = .{ .a = arena.allocator() }; + try t.expectError(error.Refused, parse(&cx, "{\"input\":\" \\t\\n\"}", cap)); + try t.expectEqualStrings("embedding text must not be empty or whitespace only", cx.message); + try t.expectError(error.Refused, parse(&cx, "{\"input\":[\"a\",\" \"]}", cap)); + try t.expectEqualStrings("embedding input at index 1 must not be empty or whitespace only", cx.message); + // the context is 4 tokens here: a 5-token input names itself, not the batch + try t.expectError(error.Refused, parse(&cx, "{\"input\":\"abcde\"}", cap)); + try t.expectEqual(errors.Kind.context_length, cx.kind); + try t.expectEqualStrings("the input (5 tokens) is longer than the model's context length (4 tokens)", cx.message); + try t.expectError(error.Refused, parse(&cx, "{\"input\":[1,2,3,4,5]}", cap)); + try t.expectEqual(errors.Kind.context_length, cx.kind); + // past both the context (4) and the batch (8): the input's own length is what is refused, counted whole + try t.expectError(error.Refused, parse(&cx, "{\"input\":\"abcdefghi\"}", cap)); + try t.expectEqualStrings("the input (9 tokens) is longer than the model's context length (4 tokens)", cx.message); + try t.expectError(error.Refused, parse(&cx, "{\"input\":[1,2,3,4,5,6,7,8,9]}", cap)); + try t.expectEqualStrings("the input (9 tokens) is longer than the model's context length (4 tokens)", cx.message); + // Unicode spaces are spaces, as Python's strip() takes them; a letter among them is text + try t.expectError(error.Refused, parse(&cx, "{\"input\":\"\\u00a0\\u2003\\u3000\"}", cap)); + try t.expectError(error.Refused, parse(&cx, "{\"input\":[\"a\",\"\\u2003\"]}", cap)); + try t.expectEqualStrings("embedding input at index 1 must not be empty or whitespace only", cx.message); + _ = try parse(&cx, "{\"input\":\"\\u00a0a\"}", cap); + // inputs within the context that together pass the batch's 8 tokens + try t.expectError(error.Refused, parse(&cx, "{\"input\":[\"abcd\",\"efgh\",\"i\"]}", cap)); + try t.expectEqualStrings("the embedding batch exceeds the admitted 8 tokens", cx.message); + const full = try parse(&cx, "{\"input\":[\"abcd\",\"efgh\"]}", cap); + try t.expectEqual(@as(usize, 8), full.tokens.len); +} diff --git a/zig/src/server/embedding_request.zig b/zig/src/server/embedding_request.zig index 20dd00fef2..82120ba7f4 100644 --- a/zig/src/server/embedding_request.zig +++ b/zig/src/server/embedding_request.zig @@ -41,8 +41,9 @@ pub fn parse(cx: *Cx, text: Text, cap: api.EmbeddingInfo, context: u32, body: Va try sequence(cx, text, cap, context, input, &tokens, &ends); } else if (first == .string or first == .array) { if (input.array.len > cap.max_inputs) return cx.refuse("too many embedding inputs"); - for (input.array) |item| { + for (input.array, 0..) |item, i| { if (@as(std.meta.Tag(Value), item) != @as(std.meta.Tag(Value), first)) return cx.refuse("embedding input forms cannot be mixed"); + if (item == .string and blank(item.string)) return cx.fail(.request, "embedding input at index {d} must not be empty or whitespace only", .{i}); try sequence(cx, text, cap, context, item, &tokens, &ends); } } else return cx.refuse("input must contain strings or token IDs"); @@ -56,14 +57,16 @@ fn sequence(cx: *Cx, text: Text, cap: api.EmbeddingInfo, context: u32, input: Va if (ends.items.len >= cap.max_inputs) return cx.refuse("too many embedding inputs"); const at = tokens.items.len; if (input == .string) { - if (input.string.len == 0) return cx.refuse("embedding text must not be empty"); + if (blank(input.string)) return cx.refuse("embedding text must not be empty or whitespace only"); const ids = text.encode(cx.a, input.string, true) catch |err| return switch (err) { error.OutOfMemory => error.OutOfMemory, error.Template => cx.refuse("the tokenizer cannot encode this embedding input"), }; defer cx.a.free(ids); + try fits(cx, context, ids.len); for (ids) |id| try token(cx, cap, id, tokens); } else if (input == .array) { + try fits(cx, context, input.array.len); for (input.array) |item| { if (item != .int) return cx.refuse("token IDs must be integers"); const id = std.fmt.parseInt(u32, item.int, 10) catch return cx.refuse("token ID is outside the vocabulary"); @@ -72,12 +75,32 @@ fn sequence(cx: *Cx, text: Text, cap: api.EmbeddingInfo, context: u32, input: Va } else return cx.refuse("each embedding input must be text or token IDs"); const count = tokens.items.len - at; if (count == 0) return cx.refuse("embedding input must contain tokens"); - if (context > 0 and count > context) return cx.fail(.context_length, "embedding input exceeds the {d}-token context", .{context}); try ends.append(cx.a, @intCast(tokens.items.len)); } +/// One input longer than the context is refused as such, before the batch's own limit. +fn fits(cx: *Cx, context: u32, count: usize) errors.Refused!void { + if (context > 0 and count > context) return cx.fail(.context_length, "the input ({d} tokens) is longer than the model's context length ({d} tokens)", .{ count, context }); +} + +/// Empty or whitespace only, by Python's str.isspace (as SGLang strips): nothing to embed; invalid UTF-8 is not blank. +fn blank(s: []const u8) bool { + const view = std.unicode.Utf8View.init(s) catch return false; + var it = view.iterator(); + while (it.nextCodepoint()) |c| if (!space(c)) return false; + return true; +} + +fn space(c: u21) bool { + return switch (c) { + 0x09...0x0d, 0x1c...0x20, 0x85, 0xa0, 0x1680, 0x2000...0x200a, 0x2028, 0x2029, 0x202f, 0x205f, 0x3000 => true, + else => false, + }; +} + fn token(cx: *Cx, cap: api.EmbeddingInfo, id: u32, tokens: *std.ArrayList(u32)) errors.Refused!void { if (id >= cap.vocabulary) return cx.refuse("token ID is outside the vocabulary"); - if (tokens.items.len >= cap.max_batch_tokens) return cx.refuse("embedding batch exceeds the admitted token limit"); + // the input's own length was checked before: here only the whole batch can pass the admitted tokens + if (tokens.items.len >= cap.max_batch_tokens) return cx.fail(.request, "the embedding batch exceeds the admitted {d} tokens", .{cap.max_batch_tokens}); try tokens.append(cx.a, id); } From 46fd9fb1999158bf65c32a0d8fa5086c3889c912 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 3 Nov 2026 01:02:17 +0000 Subject: [PATCH 18/24] fix(qwen3): refuse a HIP build without BF16 dot2 before loading The BF16 dense GEMM runs only where the build has BF16 activations; on an FP16 build (RDNA2) its kernel keeps only a trap body, and the qwen3 engine failed late as a scratch count error. It now refuses before the checkpoint loads, naming the GPU it was built for. --- zig/src/families/qwen3/native.zig | 2 ++ 1 file changed, 2 insertions(+) diff --git a/zig/src/families/qwen3/native.zig b/zig/src/families/qwen3/native.zig index e75de11ccb..7efa2b9540 100644 --- a/zig/src/families/qwen3/native.zig +++ b/zig/src/families/qwen3/native.zig @@ -136,6 +136,8 @@ fn release(ptr: *anyopaque) void { /// The engine for `dir` on `dev`, its head from the folder and its window from the memory budget. fn openEngine(a: Allocator, gpa: Allocator, io: std.Io, dev: hip.Device, dir: []const u8, o: anytype, problem: *[]const u8) anyerror!Loaded { if (o.world > 1) return refuse(a, problem, "--tp {d}: Qwen3 dense checkpoints run on one HIP GPU", .{o.world}); + // the BF16 weights run on the BF16 dot2 GEMM tile: gfx11 and gfx12, not an FP16 build (RDNA2) + if (dev.caps.act != .bf16) return refuse(a, problem, "Qwen3 dense checkpoints on HIP need a GPU with BF16 dot2 (gfx11 or gfx12); {s} runs FP16 activations", .{hip.kernels.arch}); const head = try headOf(a, io, dir, problem); const e = Engine.load(gpa, io, dir, dev.index, dev.policy.choices()) catch |err| return refuse(a, problem, "the native HIP engine cannot load {s} ({s})", .{ dir, @errorName(err) }); From 39836cb7995f675c6967006a8002e805f53fb538 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 3 Nov 2026 01:03:35 +0000 Subject: [PATCH 19/24] docs(recipes): serve Qwen3-Embedding-8B on native HIP The new qwen3 HIP family needs a recipe like the other native models: what it serves and refuses, how to build and serve it on gfx1201 with ROCm 10.1's hipcc, the request form the model card uses, and the pinned oracle and GPU gate with their results on a Radeon AI PRO R9700. --- docs/recipes/qwen3-embedding-hip.md | 87 +++++++++++++++++++++++++++++ 1 file changed, 87 insertions(+) create mode 100644 docs/recipes/qwen3-embedding-hip.md diff --git a/docs/recipes/qwen3-embedding-hip.md b/docs/recipes/qwen3-embedding-hip.md new file mode 100644 index 0000000000..4304d00bf3 --- /dev/null +++ b/docs/recipes/qwen3-embedding-hip.md @@ -0,0 +1,87 @@ +# Qwen3-Embedding-8B on native HIP + +The native HIP engine serves `Qwen/Qwen3-Embedding-8B` as published (BF16 +safetensors, revision `1d8ad4ca9b3dd8059ad90a75d4983776a23d44af`) through +`POST /v1/embeddings`. It is qualified on one Radeon AI PRO R9700 (gfx1201). +gfx11 builds run the same code; FP16 builds (RDNA2) refuse the checkpoint. + +| Responsibility | Files | +| --- | --- | +| Config, BF16 weights, batched prompt pass | `zig/src/families/qwen3/` (`config.zig`, `model.zig`, `forward.zig`) | +| Prompt-only engine, lane backend, embedding head | `zig/src/families/qwen3/` (`engine.zig`, `prompt.zig`, `native.zig`) | +| BF16 dense GEMM (dot2 tile, BF16 weight decoder) | `zig/kernels/hip/quant/bf16_decoder.hpp`, `dense_tiles.hpp`, `zig/src/hip/launches/dense.zig` | +| Head-major attention rows | `tf_heads_rows` in `zig/kernels/hip/ops/attention.hpp` | +| Stored-dtype dense weights | `zig/src/core/quant/dense.zig` (`readStored`) | +| Device contracts | `zig/tests/hip/contract_dense.zig`, `contract_heads_rows.zig` | +| Pinned model-card oracle | `tools/zig/oracle_qwen3_embedding.py` | + +The weights stay BF16 on the GPU (14.1 GiB). Activations are BF16 and every +product sums in fp32. The head comes from the sentence-transformers folder: +`1_Pooling/config.json` must pool the last token only, and `modules.json`'s +`Normalize` makes every full-width vector unit length. A folder without +`1_Pooling` is refused. + +## Build and serve + +```sh +zig build native-hip -Dhip=true -Dhip-arch=gfx1201 -Dhipcc=/opt/rocm/core/bin/hipcc +zig-out/native-hip/bin/tensorfold-native serve MODEL_DIRECTORY \ + --name Qwen3-Embedding-8B --host 127.0.0.1 --port 8090 +``` + +With ROCm 10.1, pass the hipcc under `/opt/rocm/core`: the build takes the +ROCm root from the hipcc path, and `/opt/rocm` lacks the device math headers. + +The server sizes one pass to `--context` (32,768 tokens by default, at most +the config's 40,960) and shrinks it by halves to fit the memory budget, down +to 2,048; the startup line names the window. A request holds at most 128 +inputs and the window's tokens in all. `dimensions` crops to any width from +32 to 4,096 (the card's MRL range) and renormalizes. The chat and completion +routes are refused. + +Queries take the card's instruction prefix; documents are sent as they are: + +```sh +curl -s http://127.0.0.1:8090/v1/embeddings -H 'Content-Type: application/json' -d '{ + "input": ["Instruct: Given a web search query, retrieve relevant passages that answer the query\nQuery:Explain gravity", + "Gravity is a force that attracts two bodies towards each other."]}' +``` + +Text inputs end in `<|endoftext|>` through the tokenizer's post-processor, +which last-token pooling reads. Token-ID inputs are used as sent, so a client +must append 151643 itself. + +## Gates + +The oracle runs the model card's Transformers example in a venv with +Transformers 5.22.1. It refuses a snapshot or a Transformers install whose +files differ from its pins before any model code runs: + +```sh +python3 tools/zig/oracle_qwen3_embedding.py sources --model MODEL_DIRECTORY +python3 tools/zig/oracle_qwen3_embedding.py ref --model MODEL_DIRECTORY --dtype float32 --device cpu --out ref32.json +python3 tools/zig/oracle_qwen3_embedding.py ref --model MODEL_DIRECTORY --dtype bfloat16 --device cuda --out ref16.json +python3 tools/zig/oracle_qwen3_embedding.py check --ref ref32.json --bar ref16.json --url http://127.0.0.1:8090 +``` + +`check` passes when every vector lies within `min(0.001, 2 x transformers +BF16's worst distance)` of the fp32 vector (distance = 1 - cosine), has unit +norm, is the same fp32 bytes alone, in a batch and from token IDs, and the +card's 2 x 2 scores are within 0.01. The same gate runs on the GPU without +a server: + +```sh +TF_QWEN3_EMBED_DIR=MODEL_DIRECTORY TF_QWEN3_EMBED_REF=ref32.json TF_QWEN3_EMBED_BAR=ref16.json \ + zig build qwen3-hip-gpu-test -Dhip=true -Dhip-arch=gfx1201 -Dhipcc=/opt/rocm/core/bin/hipcc +``` + +Results on the R9700 (ROCm 10.1): + +| Check | Result | +| --- | --- | +| Card example, distance to fp32 (4 inputs) | 1.32e-4, 9.30e-5, 7.76e-5, 9.27e-5 (limit 2.26e-4; transformers BF16: 1.13e-4) | +| Card 2 x 2 scores | [[0.7479, 0.0759], [0.0886, 0.6324]] (card: [[0.7493, 0.0751], [0.0880, 0.6318]]) | +| Batch, alone and token-ID vectors | identical fp32 bytes | +| `hip-kernel-test`, `dense_gemm_block` | fp32 within the dot2 chain's bound against float64; BF16 out is the fp32 sum rounded once | +| `hip-kernel-contract` | dense GEMM (fp32 and BF16 out) and `tf_heads_rows` pass gates 1 to 4; their gate 5 records are current (gfx1030 does not serve the dense GEMM) | +| Load | 4 s from the page cache | From f5b2c0c3e2299ddb87c3a0e552bc743a7dcb6bc4 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 3 Nov 2026 07:50:48 +0000 Subject: [PATCH 20/24] docs(api): document the shared /v1/embeddings refusals Review of #1978 asked for the request refusals this branch adds to every embedding family to be written down. A new Embeddings section in docs/api.md lists the whitespace-only text refusal, the over-long input that is named with its token count and the context_length_exceeded code before the batch limit, the batch limit itself, and how dimensions crops to an MRL width, and the CHANGELOG gains the Qwen3-Embedding-8B entry that points to it. --- CHANGELOG.md | 2 ++ docs/api.md | 16 ++++++++++++++++ 2 files changed, 18 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 53e64cda2c..984be25440 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -40,6 +40,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Docs: the GLM CUDA recipe says what `TF_GLM_CHUNK_SOLO` gains at TP=3 with `TF_GLM_ROCE_MULTI` (a little under half of its gain without it) (#1996). Thanks to @Bizuayeu. - DeepSeek-V4.1 on CUDA reads `--prompt-cache-gib` as the kept states' resident limit (0: none kept) unless `TF_DSV41_SESSION_RESIDENT_GIB` is set, which wins; with TP > 1 the variable is still needed on every rank (#1625). - DeepSeek-V4.1 on CUDA reads `--prompt-cache-gib` as the kept states' resident limit on every rank (0: none resident in RAM; with `TF_DSV41_SESSION_DISK` an evicted chat is parked on disk): with TP > 1 rank 0 sends its limit to the other ranks at startup. A set `TF_DSV41_SESSION_RESIDENT_GIB` wins on rank 0; set on another rank, it must equal rank 0's limit or every rank refuses to start (#1625). +- The native HIP engine serves `Qwen/Qwen3-Embedding-8B` as Qwen publishes it (BF16) through `/v1/embeddings` on a Radeon AI PRO R9700 (gfx1201): BF16 dense weights stay BF16 on a new dot2 GEMM tile, inputs of up to 32,768 tokens batch to the same vectors as alone, and `dimensions` takes any MRL width from 32 to 4,096. Every family's `/v1/embeddings` now refuses a whitespace-only text as SGLang does, and names an input longer than the context with its token count (`context_length_exceeded`) before the batch limit; see [the API notes](docs/api.md#embeddings) (#1978, part of #1977). Thanks to @laudney. + - Eight remaining host BF16 roundings in test oracles (GLM's Metal checks and CUDA FP8 KV reference, Nemotron's training scan, LFM2's DSpark taps, Flash Next's HIP kernel tests) and the XPU FP64 GDN reference use core's shared conversion, so a negative NaN gives a NaN instead of overflowing in Debug builds; every other non-NaN input gives the same bits. No kernel or served path changes (part of #1863). - Living Weights on Qwen4 dense, step 1: the QSA backward pass on the host, checked against a binary64 reference; nothing served changes (#2021, part of #1804). diff --git a/docs/api.md b/docs/api.md index 1586675e0b..e044f3cff0 100644 --- a/docs/api.md +++ b/docs/api.md @@ -498,6 +498,22 @@ On a discrete card (an RTX 3090, an A100) the server sizes its window from the c reports it, less a reserve of a 32nd of the card and at least 1 GiB. `TENSORFOLD_MEMORY_RESERVE_GIB` sets that one from 0.5 GiB up. +## Embeddings + +`POST /v1/embeddings` takes OpenAI's body: `input` is one text, a list of texts, one list of token IDs or a list of +them, with `encoding_format` (`float` or `base64`) and `dimensions`. Every family that serves embeddings shares these +refusals, each an HTTP 400 `invalid_request_error`: + +| Input | Refusal | +| --- | --- | +| a text that is empty or whitespace only (Python's whitespace set, which SGLang strips before it refuses) | "embedding text must not be empty or whitespace only"; in a list, "embedding input at index N must not be empty or whitespace only" | +| one input longer than the model's context | code `context_length_exceeded`: "the input (N tokens) is longer than the model's context length (C tokens)" | +| inputs that each fit but together reach the server's batch limit | "the embedding batch exceeds the admitted N tokens" | + +An over-long input is named before the batch limit is checked, so a client learns that the input itself is too long, +not that the batch is too full. `dimensions` crops each vector to that many leading values and renormalizes it; a +model with Matryoshka (MRL) widths takes any width in its declared range (32 to 4,096 for Qwen3-Embedding-8B). + ## Responses `choices[0].message.content` holds the answer. Reasoning uses `reasoning_content`, or From 5855df29129eb42187fea9777127095c010d5534 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 3 Nov 2026 08:23:15 +0000 Subject: [PATCH 21/24] test(qwen3): widen the embedding fidelity set to 32K and MRL Review of #1978 asked for the pinned fidelity set to cover 2K to 32K inputs and MRL widths 32, 1024 and 4096. The oracle adds five long inputs, each the longest whole-line prefix of a pinned file within 2K to 32K tokens, pins every input's text and token IDs by SHA256, runs the long ones alone at full length, and checks a server in requests that fit its window, at full width and at each MRL width against the cropped fp32 vector; the GPU test covers the same nine inputs, and the recipe records the measured distances. --- docs/recipes/qwen3-embedding-hip.md | 29 ++-- tools/zig/oracle_qwen3_embedding.py | 169 ++++++++++++++++++++--- tools/zig/test_oracle_qwen3_embedding.py | 49 +++++++ zig/src/families/qwen3/forward_test.zig | 156 +++++++++++++++++---- 4 files changed, 347 insertions(+), 56 deletions(-) diff --git a/docs/recipes/qwen3-embedding-hip.md b/docs/recipes/qwen3-embedding-hip.md index 4304d00bf3..1ab5ff3864 100644 --- a/docs/recipes/qwen3-embedding-hip.md +++ b/docs/recipes/qwen3-embedding-hip.md @@ -55,7 +55,13 @@ must append 151643 itself. The oracle runs the model card's Transformers example in a venv with Transformers 5.22.1. It refuses a snapshot or a Transformers install whose -files differ from its pins before any model code runs: +files differ from its pins before any model code runs. Its fidelity set is +the card's 4 inputs and 5 long inputs of 2,044 to 32,759 tokens, each the +longest whole-line prefix of a pinned file within 2K, 4K, 8K, 16K or 32K +tokens (the card's `LICENSE` and `README.md`, three Transformers sources); +every input's text and token IDs are pinned by SHA256. The card's inputs run +as the card runs them, one padded batch; each long input runs alone. The +fp32 reference takes about 20 minutes on 32 CPU cores: ```sh python3 tools/zig/oracle_qwen3_embedding.py sources --model MODEL_DIRECTORY @@ -64,11 +70,15 @@ python3 tools/zig/oracle_qwen3_embedding.py ref --model MODEL_DIRECTORY --dtype python3 tools/zig/oracle_qwen3_embedding.py check --ref ref32.json --bar ref16.json --url http://127.0.0.1:8090 ``` -`check` passes when every vector lies within `min(0.001, 2 x transformers -BF16's worst distance)` of the fp32 vector (distance = 1 - cosine), has unit -norm, is the same fp32 bytes alone, in a batch and from token IDs, and the -card's 2 x 2 scores are within 0.01. The same gate runs on the GPU without -a server: +`check` sends the inputs in requests that fit the server's window (the card's +and the 2K to 16K inputs together, the 32K one alone). It passes when every +vector lies within `min(0.001, 2 x transformers BF16's worst distance)` of +the fp32 vector (distance = 1 - cosine), has unit norm, is the same fp32 +bytes alone, in a batch and from token IDs, and the card's 2 x 2 scores are +within 0.01. At MRL widths 32, 1024 and 4096 (`dimensions`) every vector +must have unit norm and lie within the same bound, taken at that width, of +the fp32 vector cropped and renormalized. The full-width gate runs on the +GPU without a server too: ```sh TF_QWEN3_EMBED_DIR=MODEL_DIRECTORY TF_QWEN3_EMBED_REF=ref32.json TF_QWEN3_EMBED_BAR=ref16.json \ @@ -79,9 +89,12 @@ Results on the R9700 (ROCm 10.1): | Check | Result | | --- | --- | -| Card example, distance to fp32 (4 inputs) | 1.32e-4, 9.30e-5, 7.76e-5, 9.27e-5 (limit 2.26e-4; transformers BF16: 1.13e-4) | +| Card example, distance to fp32 (4 inputs) | 1.32e-4, 9.30e-5, 7.76e-5, 9.27e-5 | +| Long inputs, distance to fp32 (2,044, 4,043, 8,177, 16,383, 32,759 tokens) | 1.16e-4, 5.24e-5, 7.00e-5, 9.90e-5, 9.22e-5 | +| Full-width limit (2 x transformers BF16's worst, 1.57e-4) | 3.14e-4 | +| MRL width 32, 1024, 4096: worst distance (limit) | 7.92e-5 (3.77e-4), 1.18e-4 (3.25e-4), 1.32e-4 (3.14e-4) | | Card 2 x 2 scores | [[0.7479, 0.0759], [0.0886, 0.6324]] (card: [[0.7493, 0.0751], [0.0880, 0.6318]]) | -| Batch, alone and token-ID vectors | identical fp32 bytes | +| Batch, alone and token-ID vectors (all 9 inputs) | identical fp32 bytes | | `hip-kernel-test`, `dense_gemm_block` | fp32 within the dot2 chain's bound against float64; BF16 out is the fp32 sum rounded once | | `hip-kernel-contract` | dense GEMM (fp32 and BF16 out) and `tf_heads_rows` pass gates 1 to 4; their gate 5 records are current (gfx1030 does not serve the dense GEMM) | | Load | 4 s from the page cache | diff --git a/tools/zig/oracle_qwen3_embedding.py b/tools/zig/oracle_qwen3_embedding.py index fd634e20a4..9109b347e0 100644 --- a/tools/zig/oracle_qwen3_embedding.py +++ b/tools/zig/oracle_qwen3_embedding.py @@ -14,13 +14,15 @@ USAGE = """ Pins checked before any model code runs (the snapshot's revision, every file's SHA256, Transformers 5.22.1's tree): python tools/zig/oracle_qwen3_embedding.py sources --model DIR -Reference (writes JSON with token ids, vectors and the card's 2 x 2 scores): +Reference (writes JSON with token ids, vectors and the card's 2 x 2 scores; the card's 4 inputs, then 5 long inputs +of about 2K to 32K tokens cut from pinned files): python tools/zig/oracle_qwen3_embedding.py ref --model DIR --dtype float32 --device cpu --out ref32.json -Compare a running server's /v1/embeddings (same texts, one batch and one by one) with the fp32 target and the -transformers BF16 bar: +Compare a running server's /v1/embeddings (same texts in batches that fit its window, and one by one) with the fp32 +target and the transformers BF16 bar, at full width and at MRL widths 32, 1024 and 4096: python tools/zig/oracle_qwen3_embedding.py check --ref ref32.json --bar ref16.json --url http://127.0.0.1:8080 -Gates: cosine distance (1 - cos) to fp32 at most 0.001 and at most --bar-factor times the BF16 bar's worst distance; -unit norm; the card's scores within 0.01; batch, alone and token-id vectors equal as fp32 bytes (base64). +Gates, at each width: cosine distance (1 - cos) to fp32 (cropped and renormalized for MRL) at most 0.001 and at most +--bar-factor times the BF16 bar's worst distance at that width; unit norm. At full width also: the card's scores +within 0.01; batch, alone and token-id vectors equal as fp32 bytes (base64). """ CHECKPOINT = ("Qwen/Qwen3-Embedding-8B", "1d8ad4ca9b3dd8059ad90a75d4983776a23d44af") @@ -50,12 +52,40 @@ VERSION = "5.22.1" # tree_sha256 of the whole installed Transformers package, so the shared modeling and tokenizer code is pinned too. TRANSFORMERS_TREE = "b08cead6d1b841ef9723106a81a92ec37d67aff3e43a198e56b7dc76538d377b" -# The Qwen3 model, configuration and tokenizer code the card's example runs (package-relative). +# The Qwen3 model, configuration and tokenizer code the card's example runs, and the files the long inputs are cut +# from (package-relative). LIBRARY = { "models/qwen3/modeling_qwen3.py": "63a6213c15d6fb04e58a1c64fd05a4cdfb49476443fbeef524f0e0486043824b", "models/qwen3/configuration_qwen3.py": "49d9ccd0f29ccd977b93dba2004f84b867891b3d0509b71fd8ba3ed01054e64d", "models/qwen2/tokenization_qwen2.py": "fac4e6576bfe2369731be147a4e530f262bdf32f2ac50436f96f0d8bdd2fc628", + "tokenization_utils_base.py": "a70a4eda47e15756cb4ddd10ecd3a4c2c92d00b556fc00c47e7ce8a6bdb15dbd", + "generation/utils.py": "966f70ba83d94c9ce07b21150be49ff84b24eadbf9c47de052a1fb0f8ba38679", + "modeling_utils.py": "251aaae7733affbb7cdb5aa05a84157c46cac02518d868898cd6d6282cc14571", } +# Long inputs: (most tokens with the closing <|endoftext|>, where the file is pinned, file). Each is the file's longest +# prefix of whole lines within that many tokens, so the 32K one fits a 32,768-token window. +LONG = [ + (2048, "checkpoint", "LICENSE"), + (4096, "checkpoint", "README.md"), + (8192, "library", "tokenization_utils_base.py"), + (16384, "library", "generation/utils.py"), + (32768, "library", "modeling_utils.py"), +] +# Every fidelity input, in order (the card's 4, then LONG): its tokens, the SHA256 of its UTF-8 text and the SHA256 of +# its token IDs as little-endian uint32. A reference whose inputs differ is refused, whatever its pin record says. +INPUTS = [ + (27, "b11f8d3ae64c3c7029275c5d8d82172ef31a9a354223c12845b61047a0c972e8", "3c1537f6fc0a496f2b68c1aaa38484d8a82030338ae52bbd955b45b50c9efc24"), + (23, "bd6e4e8bb1b5c2e15db6003175c53f084690aeb6c09b34f1808161ef3353d524", "a08f60368f1da49a05dd929f406c503bcd0864ee464f25807866753bffe3062c"), + (8, "b6d0a6ed3a94d5e532ae83127502407102e51342d1db011caf4c5b9808a1e131", "c1e33d8973efc767e95f15a7b31af3c68eccc5ca89e0c2648f5337fec3c8f2c6"), + (31, "ba5bd6b5de9d9f827336e76d982b5dbae8aa3cc075c115f57405ba370bba8dbc", "4bd9685251a5e3587a645c0bf96dbc76ca70998073523b5537a1513c15f4c8ea"), + (2044, "f60c57fd5911f2c10704c555aef3a6e38543afa78604385ad5d8dde87050e0ff", "e10c109c63916e2cacd0143d8295fa5fa449c9eb09b7d3126405bfec9a6c9d92"), + (4043, "93c6978d85521f7b15d47dfedd386b72f063fd25ea5ada5ccf38fefdc88e7766", "a87db7cfe361a40e15c9b209796f6260bfb74017bc1e897e6e5be3d1a557cdbf"), + (8177, "3f5c76527fb4d032b31c827e5ec4d171a9feb94e7ab2ab96341028b1dd412505", "17c5b8a1cb0f288b40b9810549881b901eea58c83edb6e6d1b89ac1beab603eb"), + (16383, "c0a043846083a03ec205a34af84d7b1053a8d1a4ea4f3d3c028bb4c048ef1366", "528fda3905d2fc0340f8d5ca34f4bd91703613d76b40e5ac0835fab641acf173"), + (32759, "009256a9df975a6f67e5984760078c1dbcb0bd0b9aefbce0bdb8373ebd386dd7", "1dfcec470654b993960be543f8e84c461bc28d743aa5a50ee63bd81a9b8e15a8"), +] +# MRL widths checked besides the full one (the card's range is 32 to 4096). +WIDTHS = (32, 1024, 4096) TASK = "Given a web search query, retrieve relevant passages that answer the query" CARD_SCORES = [[0.7493016123771667, 0.0750647559762001], [0.08795969933271408, 0.6318399906158447]] @@ -75,6 +105,62 @@ def card_texts() -> list: return queries + documents +def long_texts(model: Path, package: Path, count) -> list: + """Each LONG file's longest prefix of whole lines whose token count (as `count` gives it) is within its bound.""" + roots = {"checkpoint": model, "library": package} + out = [] + for most, where, name in LONG: + lines = (roots[where] / name).read_text(encoding="utf-8").splitlines(keepends=True) + lo, hi = 0, len(lines) # the longest k with count(lines[:k]) <= most, by bisection on whole lines + while lo < hi: + k = (lo + hi + 1) // 2 + lo, hi = (k, hi) if count("".join(lines[:k])) <= most else (lo, k - 1) + text = "".join(lines[:lo]) + if count(text) < most // 2: + raise SystemExit(f"FAIL: {name} gives only {count(text)} of {most} tokens") + out.append(text) + return out + + +def fingerprint(text: str, ids: list) -> tuple: + """An input's tokens, text SHA256 and token-ID SHA256 (little-endian uint32), as INPUTS pins them.""" + import struct + raw = struct.pack(f"<{len(ids)}I", *ids) + return len(ids), hashlib.sha256(text.encode()).hexdigest(), hashlib.sha256(raw).hexdigest() + + +def check_inputs(name: str, texts, ids) -> None: + """Refuse inputs that are not the pinned fidelity set, text and token IDs both.""" + if not isinstance(texts, list) or not isinstance(ids, list) or len(texts) != len(INPUTS) or len(ids) != len(INPUTS): + raise SystemExit(f"FAIL: {name} does not hold the {len(INPUTS)} pinned inputs") + for i, (text, one, pin) in enumerate(zip(texts, ids, INPUTS)): + if not isinstance(text, str) or not isinstance(one, list) or not all(type(x) is int and 0 <= x <= 0xFFFFFFFF for x in one): + raise SystemExit(f"FAIL: {name} input {i} is not a text and its token IDs") + if fingerprint(text, one) != pin: + raise SystemExit(f"FAIL: {name} input {i} differs from its pin") + + +def groups(ids: list, window: int, most: int = 128) -> list: + """Consecutive input indices in requests of at most `most` inputs and `window` tokens in all.""" + out, part, used = [], [], 0 + for i, one in enumerate(ids): + if len(one) > window: + raise SystemExit(f"FAIL: input {i} has {len(one)} tokens, more than the {window}-token window") + if part and (used + len(one) > window or len(part) == most): + out.append(part) + part, used = [], 0 + part.append(i) + used += len(one) + return out + [part] if part else out + + +def crop(v, width: int) -> list: + """An MRL vector: the first `width` values, unit length again.""" + head = list(v[:width]) + norm = math.sqrt(sum(x * x for x in head)) + return [x / norm for x in head] + + def digest(path: Path) -> str: h = hashlib.sha256() with path.open("rb") as stream: @@ -157,15 +243,27 @@ def last_token_pool(last_hidden_states, attention_mask): return last_hidden_states[torch.arange(last_hidden_states.shape[0], device=last_hidden_states.device), lengths] dtype = getattr(torch, args.dtype) - texts = card_texts() tokenizer = AutoTokenizer.from_pretrained(args.model, padding_side="left", local_files_only=True) + package = Path(importlib.util.find_spec("transformers").origin).parent + longs = long_texts(args.model, package, lambda text: len(tokenizer(text)["input_ids"])) + texts = card_texts() + longs + ids = [tokenizer(t)["input_ids"] for t in texts] + check_inputs("the tokenized inputs", texts, ids) # before the model loads model = AutoModel.from_pretrained(args.model, dtype=dtype, local_files_only=True).to(args.device).eval() - batch = tokenizer(texts, padding=True, truncation=True, max_length=8192, return_tensors="pt").to(args.device) + # the card's 4 inputs as the card runs them: one left-padded batch, truncated at 8,192 tokens + batch = tokenizer(card_texts(), padding=True, truncation=True, max_length=8192, return_tensors="pt").to(args.device) with torch.no_grad(): out = model(**batch) - emb = F.normalize(last_token_pool(out.last_hidden_state, batch["attention_mask"]).float(), p=2, dim=1) - scores = (emb[:2] @ emb[2:].T).tolist() - ids = [tokenizer(t)["input_ids"] for t in texts] + emb = [F.normalize(last_token_pool(out.last_hidden_state, batch["attention_mask"]).float(), p=2, dim=1)] + # the long inputs one by one, whole: no padding, no truncation + for text in longs: + one = tokenizer(text, return_tensors="pt").to(args.device) + with torch.no_grad(): + out = model(**one) + emb.append(F.normalize(last_token_pool(out.last_hidden_state, one["attention_mask"]).float(), p=2, dim=1)) + print(f"long input: {one['input_ids'].shape[1]} tokens", flush=True) + emb = torch.cat(emb) + scores = (emb[:2] @ emb[2:4].T).tolist() print("scores", scores) print("card ", CARD_SCORES) json.dump({"dtype": args.dtype, "device": args.device, "pins": pins, "texts": texts, "ids": ids, @@ -225,6 +323,10 @@ def check(args) -> None: bar = json.load(open(args.bar)) if ref["dtype"] != "float32" or bar["dtype"] != "bfloat16" or ref["texts"] != bar["texts"]: raise SystemExit("FAIL: --ref must be the float32 reference and --bar the bfloat16 one, of the same texts") + if ref["texts"][:4] != card_texts(): + raise SystemExit("FAIL: the references do not start with the card's 4 inputs") + check_inputs("--ref", ref.get("texts"), ref.get("ids")) + check_inputs("--bar", bar.get("texts"), bar.get("ids")) check_pins("--ref", ref) check_pins("--bar", bar) texts = ref["texts"] @@ -232,24 +334,49 @@ def check(args) -> None: width = check_vectors("--ref", ref, n) if check_vectors("--bar", bar, n) != width or ref.get("ids") != bar.get("ids"): raise SystemExit("FAIL: --ref and --bar differ in width or token ids") - bar_worst = max(1 - cosine(b, r) for b, r in zip(bar["vectors"], ref["vectors"])) - if not math.isfinite(bar_worst): - raise SystemExit("FAIL: the bar's distance is not finite") - limit = min(args.max_distance, args.bar_factor * bar_worst) - print(f"bar: transformers bf16 worst distance {bar_worst:.2e}; limit {limit:.2e}") + parts = groups(ref["ids"], args.window) enc = {"encoding_format": "base64"} - batch = vectors(post(args.url, {"input": texts, **enc}), n, width) + + def served(inputs: list, dims) -> list: + """Every input's vector, from requests of the inputs in `parts`, cropped by the server to `dims` if given.""" + out = [b""] * n + for part in parts: + body = {"input": [inputs[i] for i in part], **enc} | ({"dimensions": dims} if dims else {}) + for i, raw in zip(part, vectors(post(args.url, body), len(part), dims or width)): + out[i] = raw + return out + + def limit_at(dims: int) -> float: + worst = max(1 - cosine(crop(b, dims), crop(r, dims)) for b, r in zip(bar["vectors"], ref["vectors"])) + if not math.isfinite(worst): + raise SystemExit(f"FAIL: the bar's distance at width {dims} is not finite") + limit = min(args.max_distance, args.bar_factor * worst) + print(f"width {dims}: transformers bf16 worst distance {worst:.2e}; limit {limit:.2e}") + return limit + + print(f"requests: {[[i for i in part] for part in parts]} (window {args.window})") + limit = limit_at(width) + batch = served(texts, None) alone = [vectors(post(args.url, {"input": t, **enc}), 1, width)[0] for t in texts] - ids = vectors(post(args.url, {"input": ref["ids"], **enc}), n, width) + ids = served(ref["ids"], None) ok = True for i in range(n): v = floats(batch[i]) dist = 1 - cosine(v, ref["vectors"][i]) norm = math.sqrt(sum(x * x for x in v)) same, by_ids = batch[i] == alone[i], batch[i] == ids[i] - print(f"input {i}: distance {dist:.2e}, norm {norm:.6f}, batch == alone {same}, ids == text {by_ids}") + print(f"input {i} ({len(ref['ids'][i])} tokens): distance {dist:.2e}, norm {norm:.6f}, " + f"batch == alone {same}, ids == text {by_ids}") ok &= math.isfinite(dist) and dist <= limit and same and by_ids and abs(norm - 1) < 1e-4 - vs = [floats(b) for b in batch] + for dims in args.widths: + limit = limit_at(dims) + for i, raw in enumerate(served(texts, dims)): + v = floats(raw) + dist = 1 - cosine(v, crop(ref["vectors"][i], dims)) + norm = math.sqrt(sum(x * x for x in v)) + print(f"width {dims} input {i}: distance {dist:.2e}, norm {norm:.6f}") + ok &= math.isfinite(dist) and dist <= limit and abs(norm - 1) < 1e-4 + vs = [floats(b) for b in batch[:4]] scores = [[sum(x * y for x, y in zip(vs[q], vs[2 + d])) for d in range(2)] for q in range(2)] worst = max(abs(scores[q][d] - CARD_SCORES[q][d]) for q in range(2) for d in range(2)) print("scores", scores, "card", CARD_SCORES, f"max diff {worst:.4f}") @@ -275,6 +402,8 @@ def main() -> None: c.add_argument("--url", default="http://127.0.0.1:8080") c.add_argument("--max-distance", type=float, default=0.001) c.add_argument("--max-score-diff", type=float, default=0.01) + c.add_argument("--window", type=int, default=32768, help="the server's window: the most tokens in one request") + c.add_argument("--widths", type=int, nargs="*", default=list(WIDTHS), help="MRL widths to check") args = p.parse_args() {"sources": sources, "ref": reference, "check": check}[args.cmd](args) diff --git a/tools/zig/test_oracle_qwen3_embedding.py b/tools/zig/test_oracle_qwen3_embedding.py index 3f4b7d459a..7dd03f1d58 100644 --- a/tools/zig/test_oracle_qwen3_embedding.py +++ b/tools/zig/test_oracle_qwen3_embedding.py @@ -1,7 +1,9 @@ """The Qwen3-Embedding oracle refuses references not made under its pins, and checks pins before model code.""" import copy import inspect +import tempfile import unittest +from pathlib import Path import oracle_qwen3_embedding as oracle @@ -47,5 +49,52 @@ def test_pins_are_checked_before_torch_or_transformers_is_imported(self): self.assertNotIn("transformers", {m.split(".")[0] for m in oracle.__dict__ if not m.startswith("_")}) +class FidelitySet(unittest.TestCase): + def test_every_long_input_is_cut_from_a_pinned_file_and_fits_the_window(self): + pins = {"checkpoint": oracle.SOURCES, "library": oracle.LIBRARY} + for most, where, name in oracle.LONG: + self.assertIn(name, pins[where]) + self.assertLessEqual(most, 32768) + self.assertEqual([2048, 4096, 8192, 16384, 32768], [most for most, _, _ in oracle.LONG]) + self.assertEqual((32, 1024, 4096), oracle.WIDTHS) + + def test_a_long_input_is_the_longest_prefix_of_whole_lines_within_its_bound(self): + with tempfile.TemporaryDirectory() as root: + lines = [f"line {i} " + "word " * i + "\n" for i in range(400)] + for _, where, name in oracle.LONG: + path = Path(root, where, name) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("".join(lines)) + texts = oracle.long_texts(Path(root, "checkpoint"), Path(root, "library"), len) + for (most, _, _), text in zip(oracle.LONG, texts): + self.assertLessEqual(len(text), most) + k = text.count("\n") + self.assertEqual(text, "".join(lines[:k])) + self.assertGreater(len(text) + len(lines[k]), most) + + def test_requests_hold_at_most_the_window_and_128_inputs_in_order(self): + self.assertEqual([[0, 1], [2], [3]], oracle.groups([[1] * 3, [1] * 5, [1] * 9, [1]], 9)) + self.assertEqual([[0], [1, 2]], oracle.groups([[1] * 4, [1] * 2, [1] * 2], 4)) + self.assertEqual([list(range(128)), [128]], oracle.groups([[1]] * 129, 1000)) + with self.assertRaises(SystemExit): + oracle.groups([[1] * 10], 9) + + def test_inputs_other_than_the_pinned_set_are_refused(self): + texts, ids = ["a"] * 9, [[1, 2]] * 9 + for bad in ([[1, 2]] * 9, [[1, 2]] * 8 + [[1 << 32]], [[1, 2]] * 8 + [[-1]], [[1, 2]] * 8 + [[1.0]]): + with self.assertRaises(SystemExit): + oracle.check_inputs("--ref", texts, bad) + self.assertEqual((2, "ca978112ca1bbdcafac231b39a23dc4da786eff8147c4e72b9807785afee48bb", + oracle.hashlib.sha256(bytes([1, 0, 0, 0, 2, 0, 0, 0])).hexdigest()), + oracle.fingerprint("a", [1, 2])) + self.assertEqual(4 + len(oracle.LONG), len(oracle.INPUTS)) + for (most, _, _), (tokens, _, _) in zip(oracle.LONG, oracle.INPUTS[4:]): + self.assertTrue(most // 2 <= tokens <= most) + self.assertLessEqual(sum(tokens for tokens, _, _ in oracle.INPUTS[:-1]), 32768) + + def test_an_mrl_vector_is_the_cropped_head_at_unit_length(self): + self.assertEqual([0.6, 0.8], oracle.crop([3.0, 4.0, 12.0], 2)) + + if __name__ == "__main__": unittest.main() diff --git a/zig/src/families/qwen3/forward_test.zig b/zig/src/families/qwen3/forward_test.zig index 116cd22b61..369525bed5 100644 --- a/zig/src/families/qwen3/forward_test.zig +++ b/zig/src/families/qwen3/forward_test.zig @@ -104,6 +104,66 @@ fn cosine(x: []const f64, y: []const f64) f64 { const embed_revision = "1d8ad4ca9b3dd8059ad90a75d4983776a23d44af"; const transformers_tree = "b08cead6d1b841ef9723106a81a92ec37d67aff3e43a198e56b7dc76538d377b"; +/// One fidelity input as the oracle's INPUTS pins it: tokens, SHA256 of the UTF-8 text, SHA256 of the IDs as LE u32. +const InputPin = struct { tokens: usize, text: *const [64]u8, ids: *const [64]u8 }; + +/// The card's 4 inputs, then 5 long inputs of 2K to 32K tokens cut from pinned files. +const embed_inputs = [_]InputPin{ + .{ .tokens = 27, .text = "b11f8d3ae64c3c7029275c5d8d82172ef31a9a354223c12845b61047a0c972e8", .ids = "3c1537f6fc0a496f2b68c1aaa38484d8a82030338ae52bbd955b45b50c9efc24" }, + .{ .tokens = 23, .text = "bd6e4e8bb1b5c2e15db6003175c53f084690aeb6c09b34f1808161ef3353d524", .ids = "a08f60368f1da49a05dd929f406c503bcd0864ee464f25807866753bffe3062c" }, + .{ .tokens = 8, .text = "b6d0a6ed3a94d5e532ae83127502407102e51342d1db011caf4c5b9808a1e131", .ids = "c1e33d8973efc767e95f15a7b31af3c68eccc5ca89e0c2648f5337fec3c8f2c6" }, + .{ .tokens = 31, .text = "ba5bd6b5de9d9f827336e76d982b5dbae8aa3cc075c115f57405ba370bba8dbc", .ids = "4bd9685251a5e3587a645c0bf96dbc76ca70998073523b5537a1513c15f4c8ea" }, + .{ .tokens = 2044, .text = "f60c57fd5911f2c10704c555aef3a6e38543afa78604385ad5d8dde87050e0ff", .ids = "e10c109c63916e2cacd0143d8295fa5fa449c9eb09b7d3126405bfec9a6c9d92" }, + .{ .tokens = 4043, .text = "93c6978d85521f7b15d47dfedd386b72f063fd25ea5ada5ccf38fefdc88e7766", .ids = "a87db7cfe361a40e15c9b209796f6260bfb74017bc1e897e6e5be3d1a557cdbf" }, + .{ .tokens = 8177, .text = "3f5c76527fb4d032b31c827e5ec4d171a9feb94e7ab2ab96341028b1dd412505", .ids = "17c5b8a1cb0f288b40b9810549881b901eea58c83edb6e6d1b89ac1beab603eb" }, + .{ .tokens = 16383, .text = "c0a043846083a03ec205a34af84d7b1053a8d1a4ea4f3d3c028bb4c048ef1366", .ids = "528fda3905d2fc0340f8d5ca34f4bd91703613d76b40e5ac0835fab641acf173" }, + .{ .tokens = 32759, .text = "009256a9df975a6f67e5984760078c1dbcb0bd0b9aefbce0bdb8373ebd386dd7", .ids = "1dfcec470654b993960be543f8e84c461bc28d743aa5a50ee63bd81a9b8e15a8" }, +}; + +/// Refuses references whose texts and token IDs are not exactly the pinned inputs, in order. +fn checkInputs(ref: std.json.Value, pins: []const InputPin) !void { + const texts = ref.object.get("texts") orelse return error.IncompleteReference; + const ids = ref.object.get("ids") orelse return error.IncompleteReference; + try holds(texts == .array and ids == .array and texts.array.items.len == pins.len and ids.array.items.len == pins.len); + for (texts.array.items, ids.array.items, pins) |x, one, pin| { + try holds(x == .string and one == .array and one.array.items.len == pin.tokens); + var h = std.crypto.hash.sha2.Sha256.init(.{}); + for (one.array.items) |id| { + try holds(id == .integer and id.integer >= 0 and id.integer <= std.math.maxInt(u32)); + var le: [4]u8 = undefined; + std.mem.writeInt(u32, &le, @intCast(id.integer), .little); + h.update(&le); + } + try holds(std.mem.eql(u8, &std.fmt.bytesToHex(h.finalResult(), .lower), pin.ids)); + var th: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(x.string, &th, .{}); + try holds(std.mem.eql(u8, &std.fmt.bytesToHex(th, .lower), pin.text)); + } +} + +test "Qwen3-Embedding references: inputs other than the pinned ones are refused" { + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + // "a" and the IDs [1, 2] (bytes 01 00 00 00 02 00 00 00) + const pin = [_]InputPin{.{ .tokens = 2, .text = "ca978112ca1bbdcafac231b39a23dc4da786eff8147c4e72b9807785afee48bb", .ids = "34fb5c825de7ca4aea6e712f19d439c1da0c92c37b423936c5f618545ca4fa1f" }}; + const parse = struct { + fn of(al: std.mem.Allocator, body: []const u8) !std.json.Value { + return std.json.parseFromSliceLeaky(std.json.Value, al, body, .{}); + } + }.of; + try checkInputs(try parse(a, "{\"texts\": [\"a\"], \"ids\": [[1, 2]]}"), &pin); + for ([_][]const u8{ + "{\"texts\": [\"b\"], \"ids\": [[1, 2]]}", + "{\"texts\": [\"a\"], \"ids\": [[1, 3]]}", + "{\"texts\": [\"a\"], \"ids\": [[1]]}", + "{\"texts\": [\"a\", \"a\"], \"ids\": [[1, 2], [1, 2]]}", + "{\"texts\": [\"a\"], \"ids\": [[1, -2]]}", + "{\"texts\": [\"a\"]}", + }) |body| try t.expect(std.meta.isError(checkInputs(try parse(a, body), &pin))); + try t.expectEqual(@as(usize, 9), embed_inputs.len); +} + fn text(v: std.json.Value, key: []const u8) ?[]const u8 { const x = (if (v == .object) v.object.get(key) else null) orelse return null; return if (x == .string) x.string else null; @@ -252,7 +312,35 @@ test "Qwen3-Embedding references: pins that do not hash to the checkpoint's file try t.expect(std.meta.isError(checkFiles(a, t.io, dir, try std.json.parseFromSliceLeaky(std.json.Value, a, unpinned, .{})))); } -test "Qwen3-Embedding: the card's inputs batched equal them alone, and sit within the BF16 bar of the fp32 vectors" { +/// The oracle's requests: consecutive inputs, at most 128 and `window` tokens in all (`ids` lengths in order). +fn groupEnds(a: std.mem.Allocator, lengths: []const usize, window: usize) ![]usize { + var out: std.ArrayList(usize) = .empty; + var used: usize = 0; + var count: usize = 0; + for (lengths, 0..) |n, i| { + if (n > window) return error.InputPastWindow; + if (count > 0 and (used + n > window or count == 128)) { + try out.append(a, i); + used = 0; + count = 0; + } + used += n; + count += 1; + } + try out.append(a, lengths.len); + return out.items; +} + +test "Qwen3-Embedding: the GPU test batches inputs as the oracle's check does" { + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + try t.expectEqualSlices(usize, &.{ 2, 3, 4 }, try groupEnds(a, &.{ 3, 5, 9, 1 }, 9)); + try t.expectEqualSlices(usize, &.{ 1, 3 }, try groupEnds(a, &.{ 4, 2, 2 }, 4)); + try t.expectError(error.InputPastWindow, groupEnds(a, &.{10}, 9)); +} + +test "Qwen3-Embedding: the card's and the long inputs batched equal them alone, within the BF16 bar of fp32" { var arena = std.heap.ArenaAllocator.init(gpa); defer arena.deinit(); const a = arena.allocator(); @@ -260,32 +348,31 @@ test "Qwen3-Embedding: the card's inputs batched equal them alone, and sit withi const ref = try readJson(a, try env(a, "TF_QWEN3_EMBED_REF") orelse return error.SkipZigTest); const bar = try readJson(a, try env(a, "TF_QWEN3_EMBED_BAR") orelse return error.SkipZigTest); try checkReferences(ref, bar); + try checkInputs(ref, &embed_inputs); try checkFiles(a, t.io, dir, ref.object.get("pins").?); var g: Gpu = undefined; try g.open(4 << 30); defer g.close(); - // the card's 2 queries and 2 documents in both references, each vector full width, before any GPU work + // the card's 2 queries and 2 documents, then 5 long inputs of 2K to 32K tokens, each vector full width const inputs = ref.object.get("ids").?.array.items; const refs = ref.object.get("vectors").?.array.items; const bars = bar.object.get("vectors").?.array.items; - try t.expectEqual(@as(usize, 4), inputs.len); - try t.expectEqual(@as(usize, 4), refs.len); - try t.expectEqual(@as(usize, 4), bars.len); - for (inputs) |one| try t.expect(one.array.items.len > 0); + try t.expectEqual(@as(usize, 9), inputs.len); + try t.expectEqual(inputs.len, refs.len); + try t.expectEqual(inputs.len, bars.len); + const lengths = try a.alloc(usize, inputs.len); + for (inputs, lengths) |one, *n| { + try t.expect(one.array.items.len > 0); + n.* = one.array.items.len; + } + try t.expect(lengths[inputs.len - 1] > 16384); + const window = 32768; var m = try qwen3.Model.load(gpa, t.io, &g.r, dir); defer m.deinit(); for (refs, bars) |r, b| { try t.expectEqual(m.spec.hidden, r.array.items.len); try t.expectEqual(m.spec.hidden, b.array.items.len); } - var ids: std.ArrayList(i32) = .empty; - const ends = try a.alloc(u32, inputs.len); - for (inputs, ends) |one, *e| { - try ids.appendSlice(a, try idsOf(a, one)); - e.* = @intCast(ids.items.len); - } - const batch = try g.lastRows(&m, ids.items, ends); - defer gpa.free(batch); const hidden = m.spec.hidden; var bar_worst: f64 = 0; for (bars, refs) |b, r| { @@ -299,21 +386,34 @@ test "Qwen3-Embedding: the card's inputs batched equal them alone, and sit withi } try t.expect(std.math.isFinite(bar_worst)); const limit = @min(0.001, 2 * bar_worst); - var begin: u32 = 0; - for (ends, 0..) |e, i| { - const alone = try g.lastRows(&m, ids.items[begin..e], &.{e - begin}); - defer gpa.free(alone); - try t.expectEqualSlices(u16, batch[i * hidden ..][0..hidden], alone); - const got = try a.alloc(f64, hidden); - const want = try a.alloc(f64, hidden); - for (batch[i * hidden ..][0..hidden], refs[i].array.items, got, want) |bits, r, *p, *q| { - p.* = widen(bits); - q.* = number(r); + var first: usize = 0; + for (try groupEnds(a, lengths, window)) |last| { + defer first = last; + var ids: std.ArrayList(i32) = .empty; + const ends = try a.alloc(u32, last - first); + for (inputs[first..last], ends) |one, *e| { + try ids.appendSlice(a, try idsOf(a, one)); + e.* = @intCast(ids.items.len); + } + const batch = try g.lastRows(&m, ids.items, ends); + defer gpa.free(batch); + var begin: u32 = 0; + for (ends, 0..) |e, j| { + defer begin = e; + const i = first + j; + const alone = try g.lastRows(&m, ids.items[begin..e], &.{e - begin}); + defer gpa.free(alone); + try t.expectEqualSlices(u16, batch[j * hidden ..][0..hidden], alone); + const got = try a.alloc(f64, hidden); + const want = try a.alloc(f64, hidden); + for (batch[j * hidden ..][0..hidden], refs[i].array.items, got, want) |bits, r, *p, *q| { + p.* = widen(bits); + q.* = number(r); + } + const distance = 1 - cosine(got, want); + std.debug.print("embedding input {d} ({d} tokens): cosine distance {e:.3} (limit {e:.3})\n", .{ i, e - begin, distance, limit }); + try t.expect(std.math.isFinite(distance) and distance <= limit); } - const distance = 1 - cosine(got, want); - std.debug.print("embedding input {d}: cosine distance {e:.3} (limit {e:.3})\n", .{ i, distance, limit }); - try t.expect(std.math.isFinite(distance) and distance <= limit); - begin = e; } } From 45dde72ce84b68edab8b18489b506d1a141c3227 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 3 Nov 2026 08:28:58 +0000 Subject: [PATCH 22/24] refactor(contract): keep the qwen3 HIP records in their own file Merging main put hip_records.zig past lean_check's 600-line limit, so the tf_heads_rows record joins the BF16 dense GEMM record in hip_records_qwen3.zig, whose tests are now collected from the contract module's test block; the record values and chip order are unchanged. The oracle's new comments and one docstring become single lines, so lean_check reports the same problems as main. --- tools/zig/oracle_qwen3_embedding.py | 11 ++++------- zig/src/core/contract.zig | 1 + zig/src/core/contract/hip_records.zig | 15 +-------------- ...ip_records_dense.zig => hip_records_qwen3.zig} | 11 ++++++++++- 4 files changed, 16 insertions(+), 22 deletions(-) rename zig/src/core/contract/{hip_records_dense.zig => hip_records_qwen3.zig} (69%) diff --git a/tools/zig/oracle_qwen3_embedding.py b/tools/zig/oracle_qwen3_embedding.py index 9109b347e0..39ccd33b7e 100644 --- a/tools/zig/oracle_qwen3_embedding.py +++ b/tools/zig/oracle_qwen3_embedding.py @@ -52,8 +52,7 @@ VERSION = "5.22.1" # tree_sha256 of the whole installed Transformers package, so the shared modeling and tokenizer code is pinned too. TRANSFORMERS_TREE = "b08cead6d1b841ef9723106a81a92ec37d67aff3e43a198e56b7dc76538d377b" -# The Qwen3 model, configuration and tokenizer code the card's example runs, and the files the long inputs are cut -# from (package-relative). +# The Qwen3 code the card's example runs, and the files the long inputs are cut from (package-relative). LIBRARY = { "models/qwen3/modeling_qwen3.py": "63a6213c15d6fb04e58a1c64fd05a4cdfb49476443fbeef524f0e0486043824b", "models/qwen3/configuration_qwen3.py": "49d9ccd0f29ccd977b93dba2004f84b867891b3d0509b71fd8ba3ed01054e64d", @@ -62,8 +61,7 @@ "generation/utils.py": "966f70ba83d94c9ce07b21150be49ff84b24eadbf9c47de052a1fb0f8ba38679", "modeling_utils.py": "251aaae7733affbb7cdb5aa05a84157c46cac02518d868898cd6d6282cc14571", } -# Long inputs: (most tokens with the closing <|endoftext|>, where the file is pinned, file). Each is the file's longest -# prefix of whole lines within that many tokens, so the 32K one fits a 32,768-token window. +# Long inputs (most tokens with the closing <|endoftext|>, pin record, file): see long_texts. LONG = [ (2048, "checkpoint", "LICENSE"), (4096, "checkpoint", "README.md"), @@ -71,8 +69,7 @@ (16384, "library", "generation/utils.py"), (32768, "library", "modeling_utils.py"), ] -# Every fidelity input, in order (the card's 4, then LONG): its tokens, the SHA256 of its UTF-8 text and the SHA256 of -# its token IDs as little-endian uint32. A reference whose inputs differ is refused, whatever its pin record says. +# Every input in order (the card's 4, then LONG) as fingerprint() gives it; check_inputs refuses any other. INPUTS = [ (27, "b11f8d3ae64c3c7029275c5d8d82172ef31a9a354223c12845b61047a0c972e8", "3c1537f6fc0a496f2b68c1aaa38484d8a82030338ae52bbd955b45b50c9efc24"), (23, "bd6e4e8bb1b5c2e15db6003175c53f084690aeb6c09b34f1808161ef3353d524", "a08f60368f1da49a05dd929f406c503bcd0864ee464f25807866753bffe3062c"), @@ -106,7 +103,7 @@ def card_texts() -> list: def long_texts(model: Path, package: Path, count) -> list: - """Each LONG file's longest prefix of whole lines whose token count (as `count` gives it) is within its bound.""" + """Each LONG file's longest whole-line prefix whose `count` (with the closing <|endoftext|>) is within its bound.""" roots = {"checkpoint": model, "library": package} out = [] for most, where, name in LONG: diff --git a/zig/src/core/contract.zig b/zig/src/core/contract.zig index c0fc468210..3991c3c210 100644 --- a/zig/src/core/contract.zig +++ b/zig/src/core/contract.zig @@ -40,6 +40,7 @@ test { _ = cubin; _ = hsaco; _ = hip_records; + _ = @import("contract/hip_records_qwen3.zig"); _ = @import("contract/fakes.zig"); } diff --git a/zig/src/core/contract/hip_records.zig b/zig/src/core/contract/hip_records.zig index 1bf4862e11..bde9f890d1 100644 --- a/zig/src/core/contract/hip_records.zig +++ b/zig/src/core/contract/hip_records.zig @@ -25,15 +25,6 @@ fn relu2Tile(comptime rows: u32, comptime cb: u32) limits.HipLaunch { /// HIP kernels as launched; tf-hip-kernel-contract limits DIR re-derives the usage records from the code objects. const own_records = [_]Record{ - .{ - .launch = .{ .name = "attention tf_heads_rows", .symbol = "tf_heads_rows", .image = "ops", .threads = 256 }, - .usage = .{ - .{ .vgprs = 17, .sgprs = 28, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, - .{ .vgprs = 17, .sgprs = 28, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, - .{ .vgprs = 17, .sgprs = 28, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, - .{ .vgprs = 17, .sgprs = 30, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, - }, - }, .{ .launch = .{ .name = "RMS norm tf_rms_rows", .symbol = "tf_rms_rows", .image = "ops", .threads = 256 }, .usage = .{ @@ -509,7 +500,7 @@ pub const records = own_records ++ @import("hip_records_stream.zig").records ++ @import("hip_records_router.zig").records ++ @import("hip_records_glm_ocr.zig").records ++ @import("hip_records_item.zig").records ++ - @import("hip_records_dense.zig").records; + @import("hip_records_qwen3.zig").records; /// Gate 5's planted HIP controls: a launch past its __launch_bounds__, and 72 KiB of LDS where a workgroup has 64. pub const controls = [_]Record{ @@ -579,10 +570,6 @@ pub fn controlProblems(a: std.mem.Allocator, r: Record, quiet: bool) !usize { return n; } -test { - _ = @import("hip_records_dense.zig"); -} - test "every registered HIP launch fits every RDNA chip we serve, and each planted control breaks one" { const a = std.testing.allocator; try std.testing.expectEqual(@as(usize, 0), try check(a, &records, false)); diff --git a/zig/src/core/contract/hip_records_dense.zig b/zig/src/core/contract/hip_records_qwen3.zig similarity index 69% rename from zig/src/core/contract/hip_records_dense.zig rename to zig/src/core/contract/hip_records_qwen3.zig index 26968332ec..565fd22ece 100644 --- a/zig/src/core/contract/hip_records_dense.zig +++ b/zig/src/core/contract/hip_records_qwen3.zig @@ -1,4 +1,4 @@ -//! Gate 5 for the BF16 dense GEMM: gfx1030's FP16 build keeps only its trap body, so the record does not serve it. +//! Gate 5 for qwen3's tf_heads_rows and BF16 dense GEMM; the GEMM does not serve gfx1030 (only a trap body there). const std = @import("std"); const limits = @import("limits.zig"); const hip_records = @import("hip_records.zig"); @@ -7,6 +7,15 @@ const Record = hip_records.Record; const check = hip_records.check; pub const records = [_]Record{ + .{ + .launch = .{ .name = "attention tf_heads_rows", .symbol = "tf_heads_rows", .image = "ops", .threads = 256 }, + .usage = .{ + .{ .vgprs = 17, .sgprs = 28, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 17, .sgprs = 28, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 17, .sgprs = 28, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + .{ .vgprs = 17, .sgprs = 30, .lds = 0, .private = 0, .bound = 1024, .wave = 32 }, + }, + }, .{ .launch = .{ .name = "dense gemm dense_gemm_block", .symbol = "_ZN2tf4rocm16dense_gemm_blockINS0_7DotBF16ELi8EEEvNS0_6AffineE", .image = "affine_tiles", .threads = 256 }, .usage = .{ From a0e57893f588a95a4e57d574bfc3a83c7f9662f4 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Tue, 3 Nov 2026 11:22:55 +0000 Subject: [PATCH 23/24] test(qwen3): gate embeddings at transformers BF16's own distance The review asked the embedding gate to follow the rule accepted for the reranker in #1981: transformers' own BF16 distance from fp32, rounded up, not twice it. The oracle and the GPU test now take the BF16 bar's worst cosine distance at each width, round it up to two significant digits with exact power-of-ten steps (never below it), and cap it at 0.001, so the limits are 1.6e-4, 1.7e-4 and 1.9e-4 at widths 4096, 1024 and 32; --bar-factor goes, and the recipe states the new limits beside the posted distances, which already meet them. --- docs/recipes/qwen3-embedding-hip.md | 12 ++++----- tools/zig/oracle_qwen3_embedding.py | 20 +++++++++++---- tools/zig/test_oracle_qwen3_embedding.py | 16 ++++++++++++ zig/src/families/qwen3/forward_test.zig | 32 ++++++++++++++++++++++-- 4 files changed, 67 insertions(+), 13 deletions(-) diff --git a/docs/recipes/qwen3-embedding-hip.md b/docs/recipes/qwen3-embedding-hip.md index 1ab5ff3864..364864d4d8 100644 --- a/docs/recipes/qwen3-embedding-hip.md +++ b/docs/recipes/qwen3-embedding-hip.md @@ -72,10 +72,10 @@ python3 tools/zig/oracle_qwen3_embedding.py check --ref ref32.json --bar ref16.j `check` sends the inputs in requests that fit the server's window (the card's and the 2K to 16K inputs together, the 32K one alone). It passes when every -vector lies within `min(0.001, 2 x transformers BF16's worst distance)` of -the fp32 vector (distance = 1 - cosine), has unit norm, is the same fp32 -bytes alone, in a batch and from token IDs, and the card's 2 x 2 scores are -within 0.01. At MRL widths 32, 1024 and 4096 (`dimensions`) every vector +vector lies within transformers' own BF16 distance of the fp32 vector, +rounded up to two digits and at most 0.001 (distance = 1 - cosine), has unit +norm, is the same fp32 bytes alone, in a batch and from token IDs, and the +card's 2 x 2 scores are within 0.01. At MRL widths 32, 1024 and 4096 (`dimensions`) every vector must have unit norm and lie within the same bound, taken at that width, of the fp32 vector cropped and renormalized. The full-width gate runs on the GPU without a server too: @@ -91,8 +91,8 @@ Results on the R9700 (ROCm 10.1): | --- | --- | | Card example, distance to fp32 (4 inputs) | 1.32e-4, 9.30e-5, 7.76e-5, 9.27e-5 | | Long inputs, distance to fp32 (2,044, 4,043, 8,177, 16,383, 32,759 tokens) | 1.16e-4, 5.24e-5, 7.00e-5, 9.90e-5, 9.22e-5 | -| Full-width limit (2 x transformers BF16's worst, 1.57e-4) | 3.14e-4 | -| MRL width 32, 1024, 4096: worst distance (limit) | 7.92e-5 (3.77e-4), 1.18e-4 (3.25e-4), 1.32e-4 (3.14e-4) | +| Full-width limit (transformers BF16's worst, 1.572e-4, rounded up) | 1.6e-4 | +| MRL width 32, 1024, 4096: worst distance (limit; transformers BF16's worst) | 7.92e-5 (1.9e-4; 1.884e-4), 1.18e-4 (1.7e-4; 1.624e-4), 1.32e-4 (1.6e-4; 1.572e-4) | | Card 2 x 2 scores | [[0.7479, 0.0759], [0.0886, 0.6324]] (card: [[0.7493, 0.0751], [0.0880, 0.6318]]) | | Batch, alone and token-ID vectors (all 9 inputs) | identical fp32 bytes | | `hip-kernel-test`, `dense_gemm_block` | fp32 within the dot2 chain's bound against float64; BF16 out is the fp32 sum rounded once | diff --git a/tools/zig/oracle_qwen3_embedding.py b/tools/zig/oracle_qwen3_embedding.py index 39ccd33b7e..25795a79a6 100644 --- a/tools/zig/oracle_qwen3_embedding.py +++ b/tools/zig/oracle_qwen3_embedding.py @@ -21,7 +21,7 @@ target and the transformers BF16 bar, at full width and at MRL widths 32, 1024 and 4096: python tools/zig/oracle_qwen3_embedding.py check --ref ref32.json --bar ref16.json --url http://127.0.0.1:8080 Gates, at each width: cosine distance (1 - cos) to fp32 (cropped and renormalized for MRL) at most 0.001 and at most ---bar-factor times the BF16 bar's worst distance at that width; unit norm. At full width also: the card's scores +the BF16 bar's worst distance at that width, rounded up to two digits; unit norm. At full width also: the card's scores within 0.01; batch, alone and token-id vectors equal as fp32 bytes (base64). """ @@ -151,6 +151,17 @@ def groups(ids: list, window: int, most: int = 128) -> list: return out + [part] if part else out +def round_up(x: float) -> float: + """`x` in [1e-20, 10) rounded up to two significant digits: the BF16 bar's distance as a gate, never below x.""" + scale = float(10 ** max(0, 1 - math.floor(math.log10(x)))) + up = math.ceil(x * scale) + while (up - 1) / scale >= x: + up -= 1 + while up / scale < x: + up += 1 + return up / scale + + def crop(v, width: int) -> list: """An MRL vector: the first `width` values, unit length again.""" head = list(v[:width]) @@ -345,9 +356,9 @@ def served(inputs: list, dims) -> list: def limit_at(dims: int) -> float: worst = max(1 - cosine(crop(b, dims), crop(r, dims)) for b, r in zip(bar["vectors"], ref["vectors"])) - if not math.isfinite(worst): - raise SystemExit(f"FAIL: the bar's distance at width {dims} is not finite") - limit = min(args.max_distance, args.bar_factor * worst) + if not (math.isfinite(worst) and worst > 0): + raise SystemExit(f"FAIL: the bar's distance at width {dims} is not finite and positive") + limit = min(args.max_distance, round_up(worst)) print(f"width {dims}: transformers bf16 worst distance {worst:.2e}; limit {limit:.2e}") return limit @@ -395,7 +406,6 @@ def main() -> None: c = sub.add_parser("check") c.add_argument("--ref", required=True, help="float32 reference JSON (the target)") c.add_argument("--bar", required=True, help="bfloat16 transformers reference JSON (the bar)") - c.add_argument("--bar-factor", type=float, default=2.0) c.add_argument("--url", default="http://127.0.0.1:8080") c.add_argument("--max-distance", type=float, default=0.001) c.add_argument("--max-score-diff", type=float, default=0.01) diff --git a/tools/zig/test_oracle_qwen3_embedding.py b/tools/zig/test_oracle_qwen3_embedding.py index 7dd03f1d58..dea021f051 100644 --- a/tools/zig/test_oracle_qwen3_embedding.py +++ b/tools/zig/test_oracle_qwen3_embedding.py @@ -1,6 +1,7 @@ """The Qwen3-Embedding oracle refuses references not made under its pins, and checks pins before model code.""" import copy import inspect +import math import tempfile import unittest from pathlib import Path @@ -12,6 +13,21 @@ def made() -> dict: return {"pins": copy.deepcopy(oracle.expected_pins()) | {"torch": "any"}} +class Gate(unittest.TestCase): + def test_the_bar_distance_rounds_up_to_two_digits(self): + for x, want in [(1.572e-4, 1.6e-4), (1.624e-4, 1.7e-4), (1.884e-4, 1.9e-4), (1.6e-4, 1.6e-4), + (0.1099, 0.11), (9.96e-5, 1.0e-4)]: + self.assertAlmostEqual(want, oracle.round_up(x), delta=want * 1e-12) + self.assertGreaterEqual(oracle.round_up(x), x) + + def test_two_digit_values_stay_and_the_next_float_rounds_past_them(self): + for k in range(2, 10): + for d in range(10, 100): + x = d / 10**k + self.assertEqual(x, oracle.round_up(x)) + self.assertEqual((d + 1) / 10**k, oracle.round_up(math.nextafter(x, 1))) + + class PinRecord(unittest.TestCase): def test_a_reference_made_under_the_pins_passes(self): oracle.check_pins("--ref", made()) diff --git a/zig/src/families/qwen3/forward_test.zig b/zig/src/families/qwen3/forward_test.zig index 369525bed5..8d29e1c5e8 100644 --- a/zig/src/families/qwen3/forward_test.zig +++ b/zig/src/families/qwen3/forward_test.zig @@ -340,6 +340,34 @@ test "Qwen3-Embedding: the GPU test batches inputs as the oracle's check does" { try t.expectError(error.InputPastWindow, groupEnds(a, &.{10}, 9)); } +/// `x` in [1e-20, 10) rounded up to two significant digits, as the oracle's round_up: never below x. +fn roundUp(x: f64) f64 { + var scale: f64 = 1; + var k = 1 - @floor(std.math.log10(x)); + while (k > 0) : (k -= 1) scale *= 10; + var up = @ceil(x * scale); + while ((up - 1) / scale >= x) up -= 1; + while (up / scale < x) up += 1; + return up / scale; +} + +test "Qwen3-Embedding gate: the BF16 bar's distance rounds up to two digits" { + const cases = [_][2]f64{ + .{ 1.572e-4, 1.6e-4 }, .{ 1.624e-4, 1.7e-4 }, .{ 1.884e-4, 1.9e-4 }, + .{ 1.6e-4, 1.6e-4 }, .{ 0.1099, 0.11 }, .{ 9.96e-5, 1.0e-4 }, + }; + for (cases) |c| { + try t.expectApproxEqRel(c[1], roundUp(c[0]), 1e-12); + try t.expect(roundUp(c[0]) >= c[0]); + } + var tens: f64 = 100; + while (tens <= 1e9) : (tens *= 10) for (10..100) |d| { + const x = @as(f64, @floatFromInt(d)) / tens; + try t.expectEqual(x, roundUp(x)); + try t.expectEqual(@as(f64, @floatFromInt(d + 1)) / tens, roundUp(std.math.nextAfter(f64, x, 1))); + }; +} + test "Qwen3-Embedding: the card's and the long inputs batched equal them alone, within the BF16 bar of fp32" { var arena = std.heap.ArenaAllocator.init(gpa); defer arena.deinit(); @@ -384,8 +412,8 @@ test "Qwen3-Embedding: the card's and the long inputs batched equal them alone, } bar_worst = @max(bar_worst, 1 - cosine(bv, rv)); } - try t.expect(std.math.isFinite(bar_worst)); - const limit = @min(0.001, 2 * bar_worst); + try t.expect(std.math.isFinite(bar_worst) and bar_worst > 0); + const limit = @min(0.001, roundUp(bar_worst)); var first: usize = 0; for (try groupEnds(a, lengths, window)) |last| { defer first = last; From b34b7a7ec7ed34c87387fa03ce83ab463342379a Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Wed, 4 Nov 2026 01:59:42 +0000 Subject: [PATCH 24/24] contract: hip_records.zig back under 600 lines with main's item records --- zig/src/core/contract/hip_records.zig | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/zig/src/core/contract/hip_records.zig b/zig/src/core/contract/hip_records.zig index bde9f890d1..986be387fe 100644 --- a/zig/src/core/contract/hip_records.zig +++ b/zig/src/core/contract/hip_records.zig @@ -496,11 +496,8 @@ const own_records = [_]Record{ } ++ @import("hip_records_pair.zig").records ++ @import("hip_records_prompt.zig").records ++ @import("hip_records_score.zig").records; -pub const records = own_records ++ @import("hip_records_stream.zig").records ++ - @import("hip_records_router.zig").records ++ - @import("hip_records_glm_ocr.zig").records ++ - @import("hip_records_item.zig").records ++ - @import("hip_records_qwen3.zig").records; +pub const records = own_records ++ @import("hip_records_stream.zig").records ++ @import("hip_records_router.zig").records ++ + @import("hip_records_glm_ocr.zig").records ++ @import("hip_records_item.zig").records ++ @import("hip_records_qwen3.zig").records; /// Gate 5's planted HIP controls: a launch past its __launch_bounds__, and 72 KiB of LDS where a workgroup has 64. pub const controls = [_]Record{