core: read rank-local NVFP4 checkpoint slices

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.

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");
+}

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);
+}

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.