core/quant: FP8 and NVFP4 values and float64 references in one place

zig/build/cuda.zig

@@ -125,6 +125,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu
     b.installArtifact(b.addExecutable(.{ .name = "tensorfold", .root_module = cli }));
     const runner = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/main.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip });
     runner.addImport("cuda", cuda);
+    runner.addImport("core", mods.core);
     runner.addImport("lanes", mods.lanes);
     runner.addImport("nemotron", mods.nemotron);
     const exlr8 = exlr8_build.coreAlias(b, target, optimize, mods.core);

zig/src/cli/cuda_precision.zig

@@ -4,33 +4,13 @@ const std = @import("std");
 const cuda = @import("cuda");
 const core = @import("core");
 const nemotron = @import("nemotron");
+const e4m3 = core.quant.fp8.e4m3;
+const toE4m3 = core.quant.fp8.toE4m3;
+const e2m1 = core.quant.nvfp4.e2m1;
 
 /// bf16 rows as read back from a dump file.
 const Rows16 = []align(1) const u16;
 
-fn e4m3(b: u8) f64 {
-    const sign: f64 = if (b & 0x80 != 0) -1 else 1;
-    const e: i32 = (b >> 3) & 0xF;
-    const m: f64 = @floatFromInt(b & 7);
-    return sign * if (e == 0) m / 8.0 * 0.015625 else (1 + m / 8.0) * std.math.pow(f64, 2, @floatFromInt(e - 7));
-}
-
-/// x rounded to the nearest e4m3 (ties to even), saturating at 448, as an FP8 activation quantizer does.
-fn toE4m3(x: f64) f64 {
-    const a = @min(@abs(x), 448.0);
-    if (a == 0) return 0;
-    const e = @max(std.math.floor(std.math.log2(a)), -6.0); // below 2^-6 the step is the subnormals' 2^-9
-    const step = std.math.pow(f64, 2, e - 3);
-    var q = @round(a / step);
-    if (@abs(a / step - std.math.trunc(a / step) - 0.5) < 1e-12 and @mod(q, 2) == 1) q -= 1;
-    return std.math.copysign(@min(q * step, 448.0), x);
-}
-
-fn e2m1(n: u8) f64 {
-    const mag = [8]f64{ 0, 0.5, 1, 1.5, 2, 3, 4, 6 };
-    return (if (n & 8 != 0) @as(f64, -1) else 1) * mag[n & 7];
-}
-
 fn bf(v: u16) f64 {
     return @floatCast(@as(f32, @bitCast(@as(u32, v) << 16)));
 }

zig/src/core/quant/fp8.zig

@@ -0,0 +1,89 @@
+//! Block-FP8 weights (e4m3 codes, a scale_inv a 128 x 128 block) on the host: code values and fp64 references.
+const std = @import("std");
+const types = @import("types.zig");
+const convert = @import("convert.zig");
+
+const Tensor = types.Tensor;
+
+/// Rows and columns one scale_inv entry covers.
+pub const block = 128;
+
+/// An e4m3 code's value; 0x7f and 0xff are NaN, as the GPU's conversion reads them.
+pub fn e4m3(b: u8) f64 {
+    if (b & 0x7f == 0x7f) return std.math.nan(f64);
+    const sign: f64 = if (b & 0x80 != 0) -1 else 1;
+    const e: i32 = (b >> 3) & 0xF;
+    const m: f64 = @floatFromInt(b & 7);
+    return sign * if (e == 0) m / 8.0 * std.math.pow(f64, 2, -6) else (1 + m / 8.0) * std.math.pow(f64, 2, @floatFromInt(e - 7));
+}
+
+/// x rounded to the nearest e4m3 value (ties to even), saturating at 448, as an FP8 activation quantizer does.
+pub fn toE4m3(x: f64) f64 {
+    const a = @min(@abs(x), 448.0);
+    if (a == 0) return 0;
+    const e = @max(std.math.floor(std.math.log2(a)), -6.0); // below 2^-6 the step is the subnormals' 2^-9
+    const step = std.math.pow(f64, 2, e - 3);
+    var q = @round(a / step);
+    if (@abs(a / step - std.math.trunc(a / step) - 0.5) < 1e-12 and @mod(q, 2) == 1) q -= 1;
+    return std.math.copysign(@min(q * step, 448.0), x);
+}
+
+/// A projection as stored: `weight` e4m3 [..., N, K], `scale_inv` [..., ceil(N / 128), ceil(K / 128)] (fp32 or bf16).
+pub const Host = struct { weight: Tensor, scale_inv: Tensor };
+
+/// Weights of a projection: its rows (a stack's leading dimensions folded) by K.
+pub fn elements(h: Host) usize {
+    return h.weight.numel();
+}
+
+/// fp64 dequantization of `h`'s rows (a stack's leading dimensions folded), code * scale_inv, into `out` [rows, K].
+pub fn reference(h: Host, out: []f64) error{UnexpectedTensor}!void {
+    const rank = h.weight.rank;
+    if (rank < 2 or h.scale_inv.rank != rank or h.weight.dtype != .f8_e4m3 or !convert.isFloat(h.scale_inv.dtype)) return error.UnexpectedTensor;
+    const n = h.weight.shape[rank - 2];
+    const k = h.weight.shape[rank - 1];
+    const nb = h.scale_inv.shape[rank - 2];
+    const kb = h.scale_inv.shape[rank - 1];
+    if (nb != (n + block - 1) / block or kb != (k + block - 1) / block or out.len != h.weight.numel()) return error.UnexpectedTensor;
+    const stack = h.weight.numel() / (n * k);
+    if (h.scale_inv.numel() != stack * nb * kb) return error.UnexpectedTensor;
+    for (0..stack) |e| for (0..n) |r| for (0..k) |c| {
+        const at = (e * n + r) * k + c;
+        const s = convert.load(h.scale_inv.dtype, h.scale_inv.bytes, (e * nb + r / block) * kb + c / block);
+        out[at] = e4m3(h.weight.bytes[at]) * s;
+    };
+}
+
+test "e4m3 codes and their rounding" {
+    try std.testing.expectEqual(@as(f64, 1), e4m3(0x38));
+    try std.testing.expectEqual(@as(f64, -2), e4m3(0xc0));
+    try std.testing.expectEqual(@as(f64, 448), e4m3(0x7e));
+    try std.testing.expectEqual(std.math.pow(f64, 2, -9), e4m3(0x01));
+    try std.testing.expect(std.math.isNan(e4m3(0x7f)) and std.math.isNan(e4m3(0xff)));
+    for (0..256) |i| {
+        const b: u8 = @intCast(i);
+        if (b & 0x7f == 0x7f) continue;
+        try std.testing.expectEqual(e4m3(b), toE4m3(e4m3(b)));
+    }
+    try std.testing.expectEqual(@as(f64, 1), toE4m3(1.0625)); // halfway between 1 and 1.125: to even
+    try std.testing.expectEqual(@as(f64, 1.25), toE4m3(1.1875));
+    try std.testing.expectEqual(@as(f64, -448), toE4m3(-1e6));
+}
+
+test "a block-FP8 stack against hand-decoded values" {
+    // two experts of [2, 130]: the last two columns fall in each row's second scale block
+    var codes: [2 * 2 * 130]u8 = @splat(0x38);
+    codes[129] = 0xc0;
+    codes[260 + 130 + 128] = 0x40;
+    const inv = [_]f32{ 0.5, 4, 1, 8 };
+    const h: Host = .{
+        .weight = .{ .dtype = .f8_e4m3, .rank = 3, .shape = .{ 2, 2, 130, 1, 1 }, .bytes = &codes },
+        .scale_inv = .{ .dtype = .f32, .rank = 3, .shape = .{ 2, 1, 2, 1, 1 }, .bytes = std.mem.sliceAsBytes(&inv) },
+    };
+    var out: [codes.len]f64 = undefined;
+    try reference(h, &out);
+    try std.testing.expectEqual(@as(f64, 0.5), out[0]);
+    try std.testing.expectEqual(@as(f64, -8), out[129]);
+    try std.testing.expectEqual(@as(f64, 1), out[260]);
+    try std.testing.expectEqual(@as(f64, 16), out[260 + 130 + 128]);
+}

zig/src/core/quant/nvfp4.zig

@@ -0,0 +1,62 @@
+//! ModelOpt NVFP4 weights (e2m1 codes, an e4m3 scale every 16 inputs, a global scale) on the host: fp64 references.
+const std = @import("std");
+const types = @import("types.zig");
+const fp8 = @import("fp8.zig");
+
+const Tensor = types.Tensor;
+
+/// Inputs one e4m3 scale covers.
+pub const group = 16;
+
+/// An e2m1 nibble's value.
+pub fn e2m1(n: u8) f64 {
+    const mag = [8]f64{ 0, 0.5, 1, 1.5, 2, 3, 4, 6 };
+    return (if (n & 8 != 0) @as(f64, -1) else 1) * mag[n & 7];
+}
+
+/// As stored: `weight` u8 [..., N, K / 2] (even columns low nibbles), `scale` e4m3 [..., N, K / 16], `global` a matrix.
+pub const Host = struct { weight: Tensor, scale: Tensor, global: []const f32 };
+
+/// Weights of a projection: its rows (a stack's leading dimensions folded) by K.
+pub fn elements(h: Host) usize {
+    return h.weight.numel() * 2;
+}
+
+/// fp64 dequantization of every row of `h` (a stack's leading dimensions folded), code * scale * global, into `out`.
+pub fn reference(h: Host, out: []f64) error{UnexpectedTensor}!void {
+    const rank = h.weight.rank;
+    if (rank < 2 or h.scale.rank != rank or h.weight.dtype != .u8 or h.scale.dtype != .f8_e4m3) return error.UnexpectedTensor;
+    const n = h.weight.shape[rank - 2];
+    const k = h.weight.shape[rank - 1] * 2;
+    if (h.scale.shape[rank - 2] != n or h.scale.shape[rank - 1] * group != k or out.len != elements(h)) return error.UnexpectedTensor;
+    const stack = elements(h) / (n * k);
+    if (h.global.len != stack) return error.UnexpectedTensor;
+    for (0..stack) |e| for (0..n) |r| for (0..k) |c| {
+        const row = e * n + r;
+        const byte = h.weight.bytes[row * (k / 2) + c / 2];
+        const nib: u8 = if (c & 1 == 0) byte & 0xF else byte >> 4;
+        out[row * k + c] = e2m1(nib) * fp8.e4m3(h.scale.bytes[row * (k / group) + c / group]) * h.global[e];
+    };
+}
+
+test "an NVFP4 stack against hand-decoded values" {
+    // two experts of [1, 32]: two scale groups a row; byte 0 holds columns 0 (low) and 1 (high)
+    var codes: [2 * 16]u8 = @splat(0);
+    codes[0] = 0xF2; // column 0: e2m1 2 -> 1.0; column 1: e2m1 0xF -> -6
+    codes[8] = 0x07; // column 16: 6
+    codes[16 + 15] = 0x30; // expert 1, column 31: e2m1 3 -> 1.5
+    const scales = [_]u8{ 0x38, 0x40, 0x38, 0x30 }; // 1, 2 | 1, 0.5
+    const global = [_]f32{ 0.25, 2 };
+    const h: Host = .{
+        .weight = .{ .dtype = .u8, .rank = 3, .shape = .{ 2, 1, 16, 1, 1 }, .bytes = &codes },
+        .scale = .{ .dtype = .f8_e4m3, .rank = 3, .shape = .{ 2, 1, 2, 1, 1 }, .bytes = &scales },
+        .global = &global,
+    };
+    var out: [64]f64 = undefined;
+    try reference(h, &out);
+    try std.testing.expectEqual(@as(f64, 0.25), out[0]);
+    try std.testing.expectEqual(@as(f64, -1.5), out[1]);
+    try std.testing.expectEqual(@as(f64, 3), out[16]);
+    try std.testing.expectEqual(@as(f64, 1.5), out[32 + 31]);
+    try std.testing.expectEqual(@as(f64, 0), out[2]);
+}

zig/src/core/quant/quant.zig

@@ -7,6 +7,9 @@ pub const slice = @import("slice.zig");
 pub const convert = @import("convert.zig");
 pub const mlx = @import("mlx.zig");
 pub const dense = @import("dense.zig");
+/// FP8 and NVFP4 values and fp64 references; host only (no Format: their kernels' layouts are each backend's).
+pub const fp8 = @import("fp8.zig");
+pub const nvfp4 = @import("nvfp4.zig");
 
 const Allocator = std.mem.Allocator;
 
@@ -197,4 +200,6 @@ test {
     _ = convert;
     _ = mlx;
     _ = dense;
+    _ = fp8;
+    _ = nvfp4;
 }

zig/tests/cuda/modelopt_tests.zig

@@ -1,24 +1,15 @@
 //! ModelOpt formats on the GPU: repacks against host twins; FP8, BF16 rows and relu^2 NVFP4 experts against fp64.
 const std = @import("std");
 const cuda = @import("cuda");
+const core = @import("core");
 const check = @import("check.zig");
 const Gpu = check.Gpu;
 
 const modelopt = cuda.modelopt;
 const qmmf = cuda.qmmf;
 const fx4 = cuda.experts.Nvfp4Experts;
-
-fn e4m3(b: u8) f64 {
-    const sign: f64 = if (b & 0x80 != 0) -1 else 1;
-    const e: i32 = (b >> 3) & 0xF;
-    const m: f64 = @floatFromInt(b & 7);
-    return sign * if (e == 0) m / 8.0 * std.math.pow(f64, 2, -6) else (1 + m / 8.0) * std.math.pow(f64, 2, @floatFromInt(e - 7));
-}
-
-fn e2m1(n: u8) f64 {
-    const mag = [8]f64{ 0, 0.5, 1, 1.5, 2, 3, 4, 6 };
-    return (if (n & 8 != 0) @as(f64, -1) else 1) * mag[n & 7];
-}
+const e4m3 = core.quant.fp8.e4m3;
+const e2m1 = core.quant.nvfp4.e2m1;
 
 fn bf(v: u16) f64 {
     return @floatCast(@as(f32, @bitCast(@as(u32, v) << 16)));

zig/tests/cuda/quant_ref_tests.zig

@@ -1,6 +1,7 @@
 //! The quantized products against float64 references on synthetic weights, and every row's bits at every row count.
 const std = @import("std");
 const cuda = @import("cuda");
+const core = @import("core");
 const check = @import("check.zig");
 const Gpu = check.Gpu;
 
@@ -10,6 +11,8 @@ const qmmf = cuda.qmmf;
 const qlinear = cuda.qlinear;
 const grouped = cuda.grouped;
 const experts = cuda.experts;
+const e4m3 = core.quant.fp8.e4m3;
+const e2m1 = core.quant.nvfp4.e2m1;
 
 const Rig = struct {
     gpu: Gpu,
@@ -75,21 +78,6 @@ fn bf(x: u16) f64 {
     return @as(f32, @bitCast(@as(u32, x) << 16));
 }
 
-/// An e4m3 byte (no NaN codes in use).
-fn e4m3(b: u8) f64 {
-    const sign: f64 = if (b & 0x80 != 0) -1 else 1;
-    const e: i32 = (b >> 3) & 0xF;
-    const m: f64 = @floatFromInt(b & 7);
-    if (e == 0) return sign * m / 8.0 * std.math.pow(f64, 2, -6);
-    return sign * (1 + m / 8.0) * std.math.pow(f64, 2, @floatFromInt(e - 7));
-}
-
-/// An e2m1 nibble.
-fn e2m1(n: u8) f64 {
-    const vals = [8]f64{ 0, 0.5, 1, 1.5, 2, 3, 4, 6 };
-    return (if (n & 8 != 0) @as(f64, -1) else 1) * vals[n & 7];
-}
-
 fn fp8Code(r: *Rig) u8 {
     while (true) {
         const b = r.rng.int(u8);