From 3f069e3feafd06fbcff7dccc556bda40472e02ee Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Thu, 15 Oct 2026 20:41:48 -0700 Subject: [PATCH 1/2] Add lossless NVFP4 compressed scale planes and TP slices --- zig/src/core/nvfp4_csf.md | 36 ++++++++ zig/src/core/nvfp4_csf.zig | 90 ++++++++++++++++++++ zig/src/core/nvfp4_csf_test.zig | 145 ++++++++++++++++++++++++++++++++ zig/src/core/root.zig | 2 + 4 files changed, 273 insertions(+) create mode 100644 zig/src/core/nvfp4_csf.md create mode 100644 zig/src/core/nvfp4_csf.zig create mode 100644 zig/src/core/nvfp4_csf_test.zig diff --git a/zig/src/core/nvfp4_csf.md b/zig/src/core/nvfp4_csf.md new file mode 100644 index 0000000000..02e2cdf40e --- /dev/null +++ b/zig/src/core/nvfp4_csf.md @@ -0,0 +1,36 @@ +# Lossless NVFP4 scale planes + +`nvfp4_csf.Plane` reads `byte-window4-fixed-stream-u24-exceptions/1` storage. +It preserves the original scale bytes, including signed zeros and NaN encodings; +it performs no floating-point conversion or requantization. + +Each 16-row slab contains 16 byte bases, followed by row-major pairs of four-bit +offsets. The even column occupies the low nibble. A decoded byte is its row base +plus its offset. Bases cannot exceed 240. Exceptions are little-endian uint32 +words: the low 24 bits name a row-major position, and the high byte replaces its +value. Positions must be unique, sorted, and inside the logical plane. + +Rows must be a positive multiple of 16, columns a positive multiple of two, and +the plane can contain at most 2^24 bytes. Stream extents must match exactly. +`Plane.init` validates the entire plane before returning a borrowed view. The +caller keeps the streams alive and immutable while that view is used. + +```zig +const plane = try core.nvfp4_csf.Plane.init(rows, columns, fixed_bytes, exception_bytes); +try plane.decode(rank_scales, row_begin, row_count, column_begin, column_count); +``` + +`decode` writes exactly the selected rectangle in row-major order, with no heap +allocation. Both row and column partitions can end inside a storage slab or a +packed byte. Invalid output extents and overlap with either borrowed stream are +rejected before writing. Sorted exceptions are searched per selected row, so a +rank does not need to materialize another rank's scales. + +This is a host format primitive. It does not change checkpoint admission, +projection arithmetic, CUDA kernels, or the runtime capability advertisement. +Callers must validate safetensors dtypes/shapes and select this exact codec before +constructing the view. + +The public format reference is the [byte-window implementation in b12x](https://github.com/local-inference-lab/b12x/blob/2cc7f66aaef76a9c7cbb660a1efb865cf3f26c88/b12x/_lib/quant/nvfp4_csf.py). +Tests use independently generated synthetic source bytes, handwritten slab +examples, every replacement byte, both partition axes, and malformed streams. diff --git a/zig/src/core/nvfp4_csf.zig b/zig/src/core/nvfp4_csf.zig new file mode 100644 index 0000000000..49d6e24319 --- /dev/null +++ b/zig/src/core/nvfp4_csf.zig @@ -0,0 +1,90 @@ +//! Lossless byte-window4 scale planes, with sorted u24-position/u8-value exceptions. + +const std = @import("std"); + +pub const codec = "byte-window4-fixed-stream-u24-exceptions/1"; +pub const Error = error{ BadGeometry, BadFixedStream, BadExceptions, BadSlice, OverlappingOutput }; + +/// A validated, borrowed checkpoint plane; callers keep both streams immutable until decoding finishes. +pub const Plane = struct { + rows: usize, + columns: usize, + fixed: []const u8, + exceptions: []const u8, + + /// Slabs hold 16 row bases followed by those rows' low-even/high-odd packed offsets. + pub fn init(rows: usize, columns: usize, fixed: []const u8, exceptions: []const u8) Error!Plane { + if (rows == 0 or rows % 16 != 0 or columns == 0 or columns % 2 != 0) return error.BadGeometry; + if (columns > 1 << 24 or rows > (1 << 24) / columns) return error.BadGeometry; + const stride = 16 + 8 * columns; + if (fixed.len != rows / 16 * stride) return error.BadFixedStream; + for (0..rows / 16) |slab| { + for (fixed[slab * stride ..][0..16]) |base| if (base > 240) return error.BadFixedStream; + } + if (exceptions.len % 4 != 0) return error.BadExceptions; + var previous: ?u32 = null; + var offset: usize = 0; + while (offset < exceptions.len) : (offset += 4) { + const position = word(exceptions, offset) & 0xffffff; + if (position >= rows * columns) return error.BadExceptions; + if (previous) |p| if (position <= p) return error.BadExceptions; + previous = position; + } + return .{ .rows = rows, .columns = columns, .fixed = fixed, .exceptions = exceptions }; + } + + /// Decode a rectangular TP slice into exactly sized row-major bytes, without allocating or rounding. + pub fn decode(self: Plane, out: []u8, row: usize, rows: usize, column: usize, columns: usize) Error!void { + if (row > self.rows or rows > self.rows - row or column > self.columns or columns > self.columns - column) + return error.BadSlice; + if (out.len != rows * columns) return error.BadSlice; + if (overlaps(out, self.fixed) or overlaps(out, self.exceptions)) return error.OverlappingOutput; + if (rows == 0 or columns == 0) return; + const stride = 16 + 8 * self.columns; + for (0..rows) |r| { + const source_row = row + r; + const slab = source_row / 16 * stride; + const base = self.fixed[slab + source_row % 16]; + const codes = slab + 16 + source_row % 16 * (self.columns / 2); + for (0..columns) |c| { + const source_column = column + c; + const codes_byte = self.fixed[codes + source_column / 2]; + const shift: u3 = if (source_column % 2 == 0) 0 else 4; + out[r * columns + c] = base + ((codes_byte >> shift) & 15); + } + var i = self.lowerBound(source_row * self.columns + column); + const end = source_row * self.columns + column + columns; + while (i < self.exceptions.len) : (i += 4) { + const entry = word(self.exceptions, i); + const position = entry & 0xffffff; + if (position >= end) break; + out[r * columns + position - (source_row * self.columns + column)] = @intCast(entry >> 24); + } + } + } + + fn lowerBound(self: Plane, position: usize) usize { + var lo: usize = 0; + var hi = self.exceptions.len / 4; + while (lo < hi) { + const mid = lo + (hi - lo) / 2; + if (word(self.exceptions, mid * 4) & 0xffffff < position) lo = mid + 1 else hi = mid; + } + return lo * 4; + } +}; + +fn word(bytes: []const u8, offset: usize) u32 { + return std.mem.readInt(u32, bytes[offset..][0..4], .little); +} + +fn overlaps(out: []u8, source: []const u8) bool { + if (out.len == 0 or source.len == 0) return false; + const a = @intFromPtr(out.ptr); + const b = @intFromPtr(source.ptr); + return if (a <= b) b - a < out.len else a - b < source.len; +} + +test { + _ = @import("nvfp4_csf_test.zig"); +} diff --git a/zig/src/core/nvfp4_csf_test.zig b/zig/src/core/nvfp4_csf_test.zig new file mode 100644 index 0000000000..13e63c88b7 --- /dev/null +++ b/zig/src/core/nvfp4_csf_test.zig @@ -0,0 +1,145 @@ +//! Independent byte fixtures and malformed-storage boundaries for lossless scale decoding. + +const std = @import("std"); +const Plane = @import("nvfp4_csf.zig").Plane; +const expect = std.testing.expectEqualSlices; + +test "handwritten slab preserves nibble order, bases and exception bytes" { + var fixed: [32]u8 = @splat(0); + fixed[0] = 240; + fixed[1] = 1; + fixed[16] = 0xf0; + fixed[17] = 0x21; + const exceptions = [_]u8{ 2, 0, 0, 255, 31, 0, 0, 128 }; + const plane = try Plane.init(16, 2, &fixed, &exceptions); + var out: [32]u8 = undefined; + try plane.decode(&out, 0, 16, 0, 2); + try expect(u8, &.{ 240, 255, 255, 3 }, out[0..4]); + try std.testing.expectEqual(@as(u8, 128), out[31]); + var odd: [2]u8 = undefined; + try plane.decode(&odd, 0, 2, 1, 1); + try expect(u8, &.{ 255, 3 }, &odd); +} + +const Fixture = struct { + fixed: []u8, + exceptions: []u8, + source: []u8, + rows: usize, + columns: usize, + + fn init(rows: usize, columns: usize) !Fixture { + const gpa = std.testing.allocator; + const source = try gpa.alloc(u8, rows * columns); + errdefer gpa.free(source); + const fixed = try gpa.alloc(u8, rows + source.len / 2); + errdefer gpa.free(fixed); + const exceptions = try gpa.alloc(u8, ((source.len + 96) / 97) * 4); + errdefer gpa.free(exceptions); + @memset(fixed, 0); + var cursor: usize = 0; + for (0..rows) |row| { + const base: u8 = @intCast((row * 37 + 19) % 241); + const slab = row / 16 * (16 + columns * 8); + fixed[slab + row % 16] = base; + for (0..columns) |column| { + const i = row * columns + column; + const delta: u8 = @intCast((column * 11 + row * 3) % 16); + source[i] = base + delta; + const offset = slab + 16 + row % 16 * (columns / 2) + column / 2; + fixed[offset] |= delta << @as(u3, if (column % 2 == 0) 0 else 4); + if (i % 97 == 0) { + source[i] = @intCast((i / 97) % 256); + const record = @as(u32, @intCast(i)) | (@as(u32, source[i]) << 24); + std.mem.writeInt(u32, exceptions[cursor..][0..4], record, .little); + cursor += 4; + } + } + } + return .{ .fixed = fixed, .exceptions = exceptions, .source = source, .rows = rows, .columns = columns }; + } + + fn deinit(self: Fixture) void { + const gpa = std.testing.allocator; + gpa.free(self.fixed); + gpa.free(self.exceptions); + gpa.free(self.source); + } + + fn check(self: Fixture, row: usize, rows: usize, column: usize, columns: usize) !void { + const out = try std.testing.allocator.alloc(u8, rows * columns); + defer std.testing.allocator.free(out); + @memset(out, 0xd6); + const plane = try Plane.init(self.rows, self.columns, self.fixed, self.exceptions); + try plane.decode(out, row, rows, column, columns); + for (0..rows) |r| try expect(u8, self.source[(row + r) * self.columns + column ..][0..columns], out[r * columns ..][0..columns]); + } +}; + +test "full planes and row/column TP partitions reproduce independent source bytes" { + for ([_][2]usize{ .{ 16, 2 }, .{ 32, 20 }, .{ 640, 160 }, .{ 2560, 40 } }) |shape| { + const fixture = try Fixture.init(shape[0], shape[1]); + defer fixture.deinit(); + try fixture.check(0, shape[0], 0, shape[1]); + for (0..2) |rank| { + try fixture.check(rank * shape[0] / 2, shape[0] / 2, 0, shape[1]); + try fixture.check(0, shape[0], rank * shape[1] / 2, shape[1] / 2); + } + try fixture.check(15, 1, 1, shape[1] - 1); + } +} + +test "all 256 byte values survive exception decoding unchanged" { + const fixed: [144]u8 = @splat(0); + var exceptions: [1024]u8 = undefined; + for (0..256) |i| std.mem.writeInt(u32, exceptions[i * 4 ..][0..4], @as(u32, @intCast(i)) * 0x1000001, .little); + const plane = try Plane.init(16, 16, &fixed, &exceptions); + var out: [256]u8 = undefined; + try plane.decode(&out, 0, 16, 0, 16); + for (out, 0..) |value, i| try std.testing.expectEqual(@as(u8, @intCast(i)), value); +} + +test "the last u24 position is valid and larger planes are refused" { + const columns = 1 << 20; + const fixed = try std.testing.allocator.alloc(u8, 16 + 8 * columns); + defer std.testing.allocator.free(fixed); + @memset(fixed, 0); + const plane = try Plane.init(16, columns, fixed, &.{ 255, 255, 255, 173 }); + var out: [1]u8 = undefined; + try plane.decode(&out, 15, 1, columns - 1, 1); + try expect(u8, &.{173}, &out); + try std.testing.expectError(error.BadGeometry, Plane.init(32, columns, fixed, &.{})); +} + +test "bad geometry, truncated streams, overflow bases and unordered exceptions are refused" { + var fixed: [32]u8 = @splat(0); + const bad = std.testing.expectError; + for ([_][2]usize{ .{ 0, 2 }, .{ 1, 2 }, .{ 16, 0 }, .{ 16, 3 }, .{ 16, 1 << 24 }, .{ 1 << 24, 2 } }) |shape| + try bad(error.BadGeometry, Plane.init(shape[0], shape[1], &fixed, &.{})); + try bad(error.BadFixedStream, Plane.init(16, 2, fixed[0..31], &.{})); + fixed[15] = 241; + try bad(error.BadFixedStream, Plane.init(16, 2, &fixed, &.{})); + fixed[15] = 0; + try bad(error.BadExceptions, Plane.init(16, 2, &fixed, &.{0})); + try bad(error.BadExceptions, Plane.init(16, 2, &fixed, &.{ 32, 0, 0, 1 })); + try bad(error.BadExceptions, Plane.init(16, 2, &fixed, &.{ 1, 0, 0, 2, 1, 0, 0, 3 })); + try bad(error.BadExceptions, Plane.init(16, 2, &fixed, &.{ 2, 0, 0, 2, 1, 0, 0, 3 })); +} + +test "invalid slices and aliased storage fail before output changes" { + var fixed: [32]u8 = @splat(0); + var exceptions = [_]u8{ 0, 0, 0, 1 }; + const plane = try Plane.init(16, 2, &fixed, &exceptions); + var out: [32]u8 = @splat(0xd6); + const bad = std.testing.expectError; + try bad(error.BadSlice, plane.decode(&out, 17, 0, 0, 2)); + try bad(error.BadSlice, plane.decode(&out, 0, 17, 0, 2)); + try bad(error.BadSlice, plane.decode(&out, 0, 16, 3, 0)); + try bad(error.BadSlice, plane.decode(&out, 0, 16, 1, 2)); + try bad(error.BadSlice, plane.decode(out[0..31], 0, 16, 0, 2)); + const poison: [32]u8 = @splat(0xd6); + try expect(u8, &poison, &out); + try bad(error.OverlappingOutput, plane.decode(&fixed, 0, 16, 0, 2)); + try bad(error.OverlappingOutput, plane.decode(&exceptions, 0, 2, 0, 2)); + try plane.decode(out[0..0], 16, 0, 2, 0); +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 736248a93d..4ffd5c0afb 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -14,6 +14,7 @@ pub const affine4_host = @import("affine4_host.zig"); pub const affine4_fold = @import("affine4_fold.zig"); pub const shard_edit = @import("shard_edit.zig"); pub const quant = @import("quant/quant.zig"); +pub const nvfp4_csf = @import("nvfp4_csf.zig"); pub const registry = @import("registry.zig"); pub const tuning = @import("tuning.zig"); pub const norms_ref = @import("norms_ref.zig"); @@ -25,6 +26,7 @@ test { _ = affine4_fold; _ = shard_edit; _ = quant; + _ = nvfp4_csf; _ = registry; _ = tuning; _ = norms_ref; From 0a0bb8634af8bdd6fc66c8d45ab00b6b00004dc1 Mon Sep 17 00:00:00 2001 From: akol1 <15747200+akol1@users.noreply.github.com> Date: Thu, 15 Oct 2026 20:57:50 -0700 Subject: [PATCH 2/2] core: read exact rank-local NVFP4 tensors with plain or compressed scales --- zig/src/core/nvfp4_slice.md | 29 ++++++ zig/src/core/nvfp4_slice.zig | 105 ++++++++++++++++++++++ zig/src/core/nvfp4_slice_test.zig | 141 ++++++++++++++++++++++++++++++ zig/src/core/root.zig | 2 + 4 files changed, 277 insertions(+) create mode 100644 zig/src/core/nvfp4_slice.md create mode 100644 zig/src/core/nvfp4_slice.zig create mode 100644 zig/src/core/nvfp4_slice_test.zig diff --git a/zig/src/core/nvfp4_slice.md b/zig/src/core/nvfp4_slice.md new file mode 100644 index 0000000000..b3d4355835 --- /dev/null +++ b/zig/src/core/nvfp4_slice.md @@ -0,0 +1,29 @@ +# Rank-local NVFP4 checkpoint reads + +`nvfp4_slice.read(checkpoint, allocator, module, rows, inputs)` reads a logical +rectangle without expanding the whole projection. The result owns its packed +FP4 code bytes and row-major E4M3 scale bytes; call `deinit` with the same +allocator. `global` preserves the stored FP32 tensor scale. There is no numeric +conversion or requantization. + +The source weight is U8 `[N, K/2]`. Plain scales are F8_E4M3 `[N, K/16]`. +Alternatively, the source may contain the fixed U8 and U32 exception tensors +described by [the CSF codec](nvfp4_csf.md). Both representations cannot be present +at once, and a partial compressed pair is refused. Tensor scale and optional +input calibration must be FP32 scalars. Calibration is consumed as checkpoint +metadata; this reader does not implement an activation quantizer. + +Row and input spans are half-open, nonempty and bounded by the original tensor. +Input cuts must align to groups of 16. Row cuts can cross a compressed 16-row +slab. Source geometry, dtype, byte extents and complete compressed streams are +validated before allocation. Allocation failures release every intermediate +buffer. The result contains only the requested code and scale rectangle. + +`dimensions(source)` and `slice(allocator, source, rows, inputs)` also accept +borrowed tensor views for callers with their own checkpoint mapping. The caller +owns those mappings; the returned buffers are independent of them. + +Tests compare plain and compressed row/column partitions with the independent +FP64 decoder, exercise allocator failures and malformed inputs, and read a real +synthetic safetensors file through the checkpoint reader. These are host format +checks, not tensor-parallel execution or GPU qualification. diff --git a/zig/src/core/nvfp4_slice.zig b/zig/src/core/nvfp4_slice.zig new file mode 100644 index 0000000000..5f3dc0fbc7 --- /dev/null +++ b/zig/src/core/nvfp4_slice.zig @@ -0,0 +1,105 @@ +//! Plain or CSF-compressed ModelOpt scales and packed codes, materialized only for one tensor-parallel rectangle. + +const std = @import("std"); +const ckpt = @import("checkpoint.zig"); +const csf = @import("nvfp4_csf.zig"); +const Tensor = ckpt.Tensor; +const Span = @import("quant/slice.zig").Span; + +pub const Scales = union(enum) { plain: Tensor, csf: struct { fixed: Tensor, exceptions: Tensor } }; +pub const Source = struct { weight: Tensor, scales: Scales, global: Tensor, input: ?Tensor = null }; +pub const Dims = struct { n: usize, k: usize }; + +/// Owned rank-local bytes: low-even FP4 nibbles, row-major E4M3 scales and the unchanged FP32 global scale. +pub const Linear = struct { + n: usize, + k: usize, + codes: []u8, + scales: []u8, + global: f32, + + pub fn deinit(self: *Linear, gpa: std.mem.Allocator) void { + gpa.free(self.codes); + gpa.free(self.scales); + self.* = undefined; + } +}; + +/// Validate a projection's original geometry before allocating rank-local storage. +pub fn dimensions(source: Source) !Dims { + const w = source.weight; + if (w.rank != 2 or w.dtype != .u8 or w.shape[0] == 0 or w.shape[1] == 0) return error.BadNvfp4; + const k = std.math.mul(usize, w.shape[1], 2) catch return error.BadNvfp4; + const count = std.math.mul(usize, w.shape[0], w.shape[1]) catch return error.BadNvfp4; + if (w.bytes.len != count or k % 16 != 0) return error.BadNvfp4; + const d: Dims = .{ .n = w.shape[0], .k = k }; + if (!scalar(source.global)) return error.BadNvfp4; + if (source.input) |a| if (!scalar(a)) return error.BadNvfp4; + switch (source.scales) { + .plain => |s| if (!s.is(.f8_e4m3, &.{ d.n, d.k / 16 }) or s.bytes.len != count / 8) return error.BadNvfp4, + .csf => |s| { + if (s.exceptions.dtype != .u32 or s.exceptions.rank != 1) return error.BadNvfp4; + if (s.exceptions.shape[0] > std.math.maxInt(usize) / 4 or s.exceptions.bytes.len != s.exceptions.shape[0] * 4) + return error.BadNvfp4; + if (!s.fixed.is(.u8, &.{ d.n / 16, 16 + 8 * (d.k / 16) })) return error.BadNvfp4; + _ = try csf.Plane.init(d.n, d.k / 16, s.fixed.bytes, s.exceptions.bytes); + }, + } + return d; +} + +/// Rows and input channels are logical weight axes; input cuts must preserve complete groups of 16. +pub fn slice(gpa: std.mem.Allocator, source: Source, rows: Span, inputs: Span) !Linear { + const d = try dimensions(source); + if (rows.from >= rows.to or rows.to > d.n or inputs.from >= inputs.to or inputs.to > d.k) + return error.BadNvfp4Slice; + if (inputs.from % 16 != 0 or inputs.to % 16 != 0) return error.BadNvfp4Slice; + const n = rows.to - rows.from; + const k = inputs.to - inputs.from; + const codes = try gpa.alloc(u8, n * (k / 2)); + errdefer gpa.free(codes); + const scales = try gpa.alloc(u8, n * (k / 16)); + errdefer gpa.free(scales); + for (0..n) |row| { + const begin = (rows.from + row) * (d.k / 2) + inputs.from / 2; + @memcpy(codes[row * (k / 2) ..][0 .. k / 2], source.weight.bytes[begin..][0 .. k / 2]); + } + switch (source.scales) { + .plain => |s| for (0..n) |row| { + const begin = (rows.from + row) * (d.k / 16) + inputs.from / 16; + @memcpy(scales[row * (k / 16) ..][0 .. k / 16], s.bytes[begin..][0 .. k / 16]); + }, + .csf => |s| { + const plane = try csf.Plane.init(d.n, d.k / 16, s.fixed.bytes, s.exceptions.bytes); + try plane.decode(scales, rows.from, n, inputs.from / 16, k / 16); + }, + } + return .{ .n = n, .k = k, .codes = codes, .scales = scales, .global = @bitCast(std.mem.readInt(u32, source.global.bytes[0..4], .little)) }; +} + +/// Read exact stored tensors, consuming optional calibration and refusing ambiguous plain/compressed scales. +pub fn read(ck: *ckpt.Checkpoint, gpa: std.mem.Allocator, module: []const u8, rows: Span, inputs: Span) !Linear { + var buf: [512]u8 = undefined; + const w = try ck.get(try std.fmt.bufPrint(&buf, "{s}.weight", .{module})); + const g = try ck.get(try std.fmt.bufPrint(&buf, "{s}.weight_scale_2", .{module})); + const input_name = try std.fmt.bufPrint(&buf, "{s}.input_scale", .{module}); + const a: ?Tensor = if (ck.has(input_name)) try ck.get(input_name) else null; + const fixed_name = try std.fmt.bufPrint(&buf, "{s}.weight_scale.nvfp4_csf_fixed", .{module}); + const fixed: ?Tensor = if (ck.has(fixed_name)) try ck.get(fixed_name) else null; + const exception_name = try std.fmt.bufPrint(&buf, "{s}.weight_scale.nvfp4_csf_exceptions", .{module}); + const exceptions: ?Tensor = if (ck.has(exception_name)) try ck.get(exception_name) else null; + if ((fixed != null) != (exceptions != null)) return error.BadNvfp4; + const plain_name = try std.fmt.bufPrint(&buf, "{s}.weight_scale", .{module}); + const plain = ck.has(plain_name); + if (plain == (fixed != null)) return error.BadNvfp4; + const scales: Scales = if (plain) .{ .plain = try ck.get(plain_name) } else .{ .csf = .{ .fixed = fixed.?, .exceptions = exceptions.? } }; + return slice(gpa, .{ .weight = w, .scales = scales, .global = g, .input = a }, rows, inputs); +} + +fn scalar(t: Tensor) bool { + return t.dtype == .f32 and t.bytes.len == 4 and (t.rank == 0 or (t.rank == 1 and t.shape[0] == 1)); +} + +test { + _ = @import("nvfp4_slice_test.zig"); +} diff --git a/zig/src/core/nvfp4_slice_test.zig b/zig/src/core/nvfp4_slice_test.zig new file mode 100644 index 0000000000..1ca0a128df --- /dev/null +++ b/zig/src/core/nvfp4_slice_test.zig @@ -0,0 +1,141 @@ +//! Rank-local NVFP4 bytes against the independent FP64 decoder and real safetensors headers. + +const std = @import("std"); +const reader = @import("nvfp4_slice.zig"); +const st = @import("safetensors.zig"); +const ckpt = @import("checkpoint.zig"); +const quant = @import("quant/quant.zig"); + +const Fixture = struct { + codes: [512]u8 = undefined, + scales: [64]u8 = undefined, + fixed: [48]u8 = @splat(0), + exceptions: [4]u8 = .{ 3, 0, 0, 0x40 }, + global: [4]u8 = undefined, + + fn init() Fixture { + var f: Fixture = .{}; + std.mem.writeInt(u32, &f.global, @bitCast(@as(f32, 0.25)), .little); + for (&f.codes, 0..) |*b, i| b.* = @truncate(i * 29 + 3); + for (0..16) |row| { + const base: u8 = @intCast(0x30 + row % 8); + f.fixed[row] = base; + f.fixed[16 + row * 2] = 0x10; + f.fixed[17 + row * 2] = 0x32; + for (0..4) |column| f.scales[row * 4 + column] = base + @as(u8, @intCast(column)); + } + f.scales[3] = 0x40; + return f; + } + + fn source(f: *const Fixture, compressed: bool) reader.Source { + return .{ + .weight = tensor(.u8, &.{ 16, 32 }, &f.codes), + .scales = if (compressed) .{ .csf = .{ + .fixed = tensor(.u8, &.{ 1, 48 }, &f.fixed), + .exceptions = tensor(.u32, &.{1}, &f.exceptions), + } } else .{ .plain = tensor(.f8_e4m3, &.{ 16, 4 }, &f.scales) }, + .global = tensor(.f32, &.{}, &f.global), + }; + } +}; + +fn tensor(dtype: st.DType, shape: []const usize, bytes: []const u8) st.Tensor { + var out: st.Tensor = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .bytes = bytes }; + @memcpy(out.shape[0..shape.len], shape); + return out; +} + +fn reference(linear: reader.Linear, out: []f64) !void { + try quant.nvfp4.reference(.{ + .weight = tensor(.u8, &.{ linear.n, linear.k / 2 }, linear.codes), + .scale = tensor(.f8_e4m3, &.{ linear.n, linear.k / 16 }, linear.scales), + .global = &.{linear.global}, + }, out); +} + +test "plain and compressed rows/columns preserve every FP64 decoded value" { + const f = Fixture.init(); + const a = std.testing.allocator; + var whole = try reader.slice(a, f.source(false), .{ .from = 0, .to = 16 }, .{ .from = 0, .to = 64 }); + defer whole.deinit(a); + var expected: [1024]f64 = undefined; + try reference(whole, &expected); + for ([_]bool{ false, true }) |compressed| { + for (0..2) |axis| for (0..2) |rank| { + const r0 = if (axis == 0) rank * 8 else 0; + const rn: usize = if (axis == 0) 8 else 16; + const c0 = if (axis == 1) rank * 32 else 0; + const cn: usize = if (axis == 1) 32 else 64; + var part = try reader.slice(a, f.source(compressed), .{ .from = r0, .to = r0 + rn }, .{ .from = c0, .to = c0 + cn }); + defer part.deinit(a); + var out: [512]f64 = undefined; + try reference(part, &out); + for (0..rn) |row| try std.testing.expectEqualSlices(f64, expected[(r0 + row) * 64 + c0 ..][0..cn], out[row * cn ..][0..cn]); + }; + } +} + +test "bad compressed tensor metadata and cuts cannot become a rank projection" { + const f = Fixture.init(); + var source = f.source(true); + const bad = std.testing.expectError; + source.scales.csf.exceptions.dtype = .i32; + try bad(error.BadNvfp4, reader.dimensions(source)); + source = f.source(true); + source.scales.csf.fixed.shape[1] -= 1; + try bad(error.BadNvfp4, reader.dimensions(source)); + source = f.source(false); + source.global.dtype = .bf16; + try bad(error.BadNvfp4, reader.dimensions(source)); + const a = std.testing.allocator; + try bad(error.BadNvfp4Slice, reader.slice(a, f.source(true), .{ .from = 0, .to = 17 }, .{ .from = 0, .to = 64 })); + try bad(error.BadNvfp4Slice, reader.slice(a, f.source(true), .{ .from = 8, .to = 8 }, .{ .from = 0, .to = 64 })); + try bad(error.BadNvfp4Slice, reader.slice(a, f.source(true), .{ .from = 0, .to = 16 }, .{ .from = 1, .to = 32 })); +} + +test "allocation failures free every rank-local buffer" { + const f = Fixture.init(); + for (0..3) |fail_at| { + var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{ .fail_index = fail_at }); + const a = failing.allocator(); + if (reader.slice(a, f.source(true), .{ .from = 0, .to = 16 }, .{ .from = 0, .to = 32 })) |value| { + var owned = value; + owned.deinit(a); + } else |e| try std.testing.expectEqual(error.OutOfMemory, e); + try std.testing.expectEqual(failing.allocated_bytes, failing.freed_bytes); + } +} + +test "checkpoint reading consumes CSF streams and preserves optional calibration" { + const f = Fixture.init(); + const io = std.testing.io; + const a = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const header = + \\{"l.weight":{"dtype":"U8","shape":[16,32],"data_offsets":[0,512]}, + \\ "l.weight_scale.nvfp4_csf_fixed":{"dtype":"U8","shape":[1,48],"data_offsets":[512,560]}, + \\ "l.weight_scale.nvfp4_csf_exceptions":{"dtype":"U32","shape":[1],"data_offsets":[560,564]}, + \\ "l.weight_scale_2":{"dtype":"F32","shape":[],"data_offsets":[564,568]}, + \\ "l.input_scale":{"dtype":"F32","shape":[1],"data_offsets":[568,572]}} + ; + var image: [8 + header.len + 572]u8 = undefined; + std.mem.writeInt(u64, image[0..8], header.len, .little); + @memcpy(image[8..][0..header.len], header); + const data = image[8 + header.len ..]; + @memcpy(data[0..512], &f.codes); + @memcpy(data[512..560], &f.fixed); + @memcpy(data[560..564], &f.exceptions); + @memcpy(data[564..568], &f.global); + @memcpy(data[568..572], &f.global); + try tmp.dir.writeFile(io, .{ .sub_path = "model.safetensors", .data = &image }); + const dir = try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}", .{tmp.sub_path}); + defer a.free(dir); + var ck = try ckpt.Checkpoint.openModel(a, io, dir); + defer ck.close(); + var part = try reader.read(&ck, a, "l", .{ .from = 8, .to = 16 }, .{ .from = 32, .to = 64 }); + defer part.deinit(a); + try std.testing.expectEqual(@as(usize, 0), ck.unused()); + for (0..8) |r| try std.testing.expectEqualSlices(u8, f.scales[(8 + r) * 4 + 2 ..][0..2], part.scales[r * 2 ..][0..2]); +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 4ffd5c0afb..9febe4107b 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -15,6 +15,7 @@ pub const affine4_fold = @import("affine4_fold.zig"); pub const shard_edit = @import("shard_edit.zig"); pub const quant = @import("quant/quant.zig"); pub const nvfp4_csf = @import("nvfp4_csf.zig"); +pub const nvfp4_slice = @import("nvfp4_slice.zig"); pub const registry = @import("registry.zig"); pub const tuning = @import("tuning.zig"); pub const norms_ref = @import("norms_ref.zig"); @@ -27,6 +28,7 @@ test { _ = shard_edit; _ = quant; _ = nvfp4_csf; + _ = nvfp4_slice; _ = registry; _ = tuning; _ = norms_ref;