From 0a123fd2f81a95887e5738ef077ff0a27d20265a Mon Sep 17 00:00:00 2001 From: "L.B.R." <106273+laudney@users.noreply.github.com> Date: Fri, 6 Nov 2026 10:47:16 +0000 Subject: [PATCH 01/36] fix(hip): zero masked V per row in the wide prompt tile On RDNA 3.5, v_dot2_f32_bf16 and v_dot2_f32_f16 round a sum differently when one lane's product is -0 instead of +0. The 64-row prompt tile multiplies P by V two keys at a time and zeroed V only past the span, so when a row saw an odd number of keys in its last tile, the partner key fed +0 in a run whose chunk ended there but p = 0 times its real V, a -0 for a negative value, in a run whose chunk ended later, and a row's bytes hung on where its chunk began (#2317). Each row now zeroes every value past its own causal edge, so all runs feed the same dot2 lanes on every chip. This also drops the 0 x inf NaN past the causal edge from the prefill contract's known finding, adds cut 1 to the chunk-start test, moves the gfx1151 pin to the new bits (gfx1152 reproduced the old pin exactly), and refreshes the prompt and image tiles' VGPR records on every chip. --- docs/kernel-contract.md | 3 +-- zig/kernels/hip/attention/prefill.hip | 14 +++++++--- zig/src/core/contract/hip_records_prompt.zig | 28 ++++++++++---------- zig/src/core/contract/hip_records_vision.zig | 8 +++--- zig/src/hip/kernel_tests/attention.zig | 3 ++- zig/src/hip/kernel_tests/attention_pin.zig | 2 +- zig/tests/hip/contract_prefill.zig | 7 +++-- 7 files changed, 36 insertions(+), 29 deletions(-) diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index b19db41e0a..ea54bee455 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -133,8 +133,7 @@ Python family's Triton glue it was ported from. Nemotron's Metal `tf_xup_w3` giv the NaN: its relu2 is `metal::max(bfloat(acc), 0.0f)`; the other up pipelines give its bytes on the same inputs. HIP's prompt attention tile (`tf_fa_paged_*`, `tf_fa_wide_*`) writes 0 for a row whose every score is NaN where float64 gives NaN, as its running max is `fmaxf` and a row at max -inf takes p = 0; and an infinite value at a key -whose fp32 weight is 0 (past the row's causal edge in its 64-key tile, or under fp32's least) gives NaN where float64 -gives inf or a finite value, as its P x V takes every key of the tile. +the row sees under fp32's least weight gives NaN where float64 gives inf, as 0 x inf enters its P x V. Qwen3.8 Flash Next's `tf_sattn_nax` and `tf_attn256_nax` score masked keys -FLT_MAX, the value each row's running max starts from, so a row whose every score is -Inf weighs its masked keys 1 and gives a finite average where float64 gives NaN. `tf_sattn_nax`'s padding lanes also read key 0 at p = 0, so a non-finite value in key 0's row reaches every diff --git a/zig/kernels/hip/attention/prefill.hip b/zig/kernels/hip/attention/prefill.hip index 4dd231474a..5d6d4c8405 100644 --- a/zig/kernels/hip/attention/prefill.hip +++ b/zig/kernels/hip/attention/prefill.hip @@ -221,11 +221,13 @@ __device__ void fa_dot2(const float* __restrict__ q, const Cache* __restrict__ k // the online softmax in registers: a row's 64 keys sit on its 16 lanes, key pair pp's p on lane pp / 2 pair_t ph[4][2]; float alpha_r[4]; + int seen[4]; // the keys of this tile each row sees: none to 0, all from 64 #pragma unroll for (int i = 0; i < 4; ++i) { const int row = sr + 16 * i; int vis = CAUSAL ? q_pos0 + q0 + row + 1 : span; if (vis > span) vis = span; + seen[i] = vis - k0; float tile_max = -INFINITY; int nan_row = 0; // !CAUSAL: a NaN score makes its row NaN, as in Torch #pragma unroll @@ -257,8 +259,9 @@ __device__ void fa_dot2(const float* __restrict__ q, const Cache* __restrict__ k // the chunk's key pairs as (key, key + 1) pairs at each of its dims const int pair = tid / kKChunks, g = tid % kKChunks; using elem8 = elem __attribute__((ext_vector_type(8))); - const elem8 ha = __builtin_bit_cast(elem8, kept(va, FULL || 2 * pair < rows)); - const elem8 hb = __builtin_bit_cast(elem8, kept(vb, FULL || 2 * pair + 1 < rows)); + // a value past the span reads the span's last (load16); the products zero each key past a row's edge + const elem8 ha = __builtin_bit_cast(elem8, va); + const elem8 hb = __builtin_bit_cast(elem8, vb); pair_t lo[4], hi[4]; #pragma unroll for (int m = 0; m < 4; ++m) lo[m] = pair_t{ha[m], hb[m]}, hi[m] = pair_t{ha[4 + m], hb[4 + m]}; @@ -281,8 +284,13 @@ __device__ void fa_dot2(const float* __restrict__ q, const Cache* __restrict__ k #pragma unroll for (int i = 0; i < 4; ++i) { const pair_t pp4 = row_lane(ph[i][pp % 2], lane); + // a key past the row's edge takes value 0: RDNA 3.5's dot2 rounds a -0 lane (0 x v < 0) off +0's + const int left = FULL ? 2 : seen[i] - 2 * pp; + const unsigned keep = left >= 2 ? ~0u : left == 1 ? 0xffffu : 0u; // key 2pp is the low half #pragma unroll - for (int m = 0; m < 4; ++m) pv[i][m] = T::dot(pp4, vv[m], pv[i][m]); + for (int m = 0; m < 4; ++m) + pv[i][m] = T::dot(pp4, __builtin_bit_cast(pair_t, __builtin_bit_cast(unsigned, vv[m]) & keep), + pv[i][m]); } } #pragma unroll diff --git a/zig/src/core/contract/hip_records_prompt.zig b/zig/src/core/contract/hip_records_prompt.zig index e737e37650..fb938dfd28 100644 --- a/zig/src/core/contract/hip_records_prompt.zig +++ b/zig/src/core/contract/hip_records_prompt.zig @@ -12,27 +12,27 @@ fn promptTile(comptime symbol: []const u8, usage: [limits.rdna.len]U) Record { pub const records = [_]Record{ promptTile("tf_fa_wide_f16", .{ - .{ .vgprs = 230, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 230, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 230, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 230, .sgprs = 49, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 252, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 236, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 236, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 236, .sgprs = 49, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, }), promptTile("tf_fa_wide_bf16", .{ .{ .vgprs = 246, .sgprs = 46, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 230, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 230, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 230, .sgprs = 49, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 236, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 236, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 236, .sgprs = 49, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, }), promptTile("tf_fa_paged_f16", .{ - .{ .vgprs = 231, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 231, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 231, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 231, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 255, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 239, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 239, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 239, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, }), promptTile("tf_fa_paged_bf16", .{ .{ .vgprs = 249, .sgprs = 52, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 231, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 231, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 231, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 239, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 239, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 239, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, }), }; diff --git a/zig/src/core/contract/hip_records_vision.zig b/zig/src/core/contract/hip_records_vision.zig index 89fd81b464..aea8f078f5 100644 --- a/zig/src/core/contract/hip_records_vision.zig +++ b/zig/src/core/contract/hip_records_vision.zig @@ -41,10 +41,10 @@ pub const records = [_]Record{ .{ .launch = .{ .name = "BF16 full image attention", .symbol = "tf_fa_full_bf16", .image = "prefill", .threads = 256 }, .usage = .{ - .{ .vgprs = 245, .sgprs = 45, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 229, .sgprs = 46, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 229, .sgprs = 46, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 229, .sgprs = 48, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 246, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 230, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 230, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 230, .sgprs = 49, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, }, }, }; diff --git a/zig/src/hip/kernel_tests/attention.zig b/zig/src/hip/kernel_tests/attention.zig index 52283569fb..ba9e9f69fe 100644 --- a/zig/src/hip/kernel_tests/attention.zig +++ b/zig/src/hip/kernel_tests/attention.zig @@ -337,10 +337,11 @@ test "a query row's prompt attention keeps its bits wherever its chunk starts, o const t = try Rig.open(1 << 20); defer t.close(); var rng: rig.Rng = .{ .state = 0x7A1C5 }; - // the 64-row tile (wide on, d 128), then the 16-row tile (attention=f32 turns wide off) + // the 64-row tile (wide on, d 128), then the 16-row tile (attention=f32 turns wide off); odd cuts split a key pair for ([_]bool{ true, false }) |wide| { t.on.wide = wide; defer t.on.wide = true; + try shifted(t, &rng, 32, 8, 128, 2400, 1); try shifted(t, &rng, 32, 8, 128, 2400, 87); try shifted(t, &rng, 32, 8, 128, 2400, 2100); } diff --git a/zig/src/hip/kernel_tests/attention_pin.zig b/zig/src/hip/kernel_tests/attention_pin.zig index 51ba5126ce..8ee920fa5a 100644 --- a/zig/src/hip/kernel_tests/attention_pin.zig +++ b/zig/src/hip/kernel_tests/attention_pin.zig @@ -13,7 +13,7 @@ const d = 256; /// The digest of every case's output on an architecture whose bits were recorded. const pins = [_]struct { arch: []const u8, digest: u64 }{ - .{ .arch = "gfx1151", .digest = 0x30a7911a6c016c2b }, + .{ .arch = "gfx1151", .digest = 0xe9e35ddbb2a6ad10 }, }; /// A prompt chunk: `qlen` queries from position `q_pos0`. diff --git a/zig/tests/hip/contract_prefill.zig b/zig/tests/hip/contract_prefill.zig index f1758173de..cbf65ce444 100644 --- a/zig/tests/hip/contract_prefill.zig +++ b/zig/tests/hip/contract_prefill.zig @@ -39,11 +39,10 @@ const names = [2][2][]const u8{ .{ "hip prompt tile tf_fa_paged_f16 (fp16 cache pages)", "hip prompt tile tf_fa_paged_bf16 (bf16 cache pages)" }, }; -/// What gate 1 finds on the tile: NaN scores leave its max, and its P x V takes every key of a tile. +/// What gate 1 finds on the tile: NaN scores leave its max, and its P x V takes a seen key whose weight underflows. const finding = "a row whose every score is NaN writes 0 where float64 gives NaN (its running max is fmaxf, " ++ - "which drops NaN, and a row at max -inf takes p = 0), and an infinite value at a key whose fp32 weight is 0 " ++ - "(past the row's causal edge in its tile, or under fp32's least) gives NaN where float64 gives inf or a finite " ++ - "value, as 0 x inf enters P x V"; + "which drops NaN, and a row at max -inf takes p = 0), and an infinite value at a key the row sees under " ++ + "fp32's least weight gives NaN where float64 gives inf, as 0 x inf enters P x V"; /// A row's P x V chain: a dot2 a key pair and a fold a tile, under its keys plus a tile. fn keyChain(s: c.Shape) usize { From b6a612c470984bbd0568d734728f314a2b4aee89 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Fri, 6 Nov 2026 15:06:09 +0000 Subject: [PATCH 02/36] feat(hip): name the chips whose dot2 rounds a -0 lane apart RDNA 3.5's dot2 rounds a sum with a -0 lane apart from the same sum with +0, which made #2317's prompt rows depend on their chunk start, while gfx1201's does not. A caps field, dot2_signed_zero, records that fact per generation and reaches the kernels as TF_DOT2_SIGNED_ZERO, so the prompt tile can keep main's code on the chips that do not need the per-row mask. --- zig/build/hip.zig | 1 + zig/src/hip/caps.zig | 12 ++++++++---- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/zig/build/hip.zig b/zig/build/hip.zig index 079a3cfcb4..41d140fe47 100644 --- a/zig/build/hip.zig +++ b/zig/build/hip.zig @@ -478,6 +478,7 @@ fn compileSource(b: *std.Build, cc: Hipcc, arch: []const u8, source: LazyPath, o run.addArg(b.fmt("-DTF_SDOT4={d}", .{@intFromBool(c.sdot4)})); run.addArg(b.fmt("-DTF_SDOT8={d}", .{@intFromBool(c.sdot8)})); run.addArg(b.fmt("-DTF_MATRIX={d}", .{@intFromBool(c.matrix != .none)})); + run.addArg(b.fmt("-DTF_DOT2_SIGNED_ZERO={d}", .{@intFromBool(c.dot2_signed_zero)})); run.addArg(b.fmt("--rocm-path={s}", .{root})); const lib = cc.device_lib orelse b.fmt("{s}/lib/llvm/amdgcn/bitcode", .{root}); run.addArg(b.fmt("--rocm-device-lib-path={s}", .{lib})); diff --git a/zig/src/hip/caps.zig b/zig/src/hip/caps.zig index 8ee61f99f2..db9b6af44d 100644 --- a/zig/src/hip/caps.zig +++ b/zig/src/hip/caps.zig @@ -13,6 +13,8 @@ pub const Caps = struct { sdot8: bool, matrix: Matrix, act: Act, + /// dot2 rounds a sum with a -0 lane apart from the same sum with +0, so masked lanes must enter as +0. + dot2_signed_zero: bool, generation: Generation, /// The caps of `name` ("gfx1030", "gfx1151:xnack-" ...), or null for a GPU outside the table. @@ -35,10 +37,10 @@ const Row = struct { names: []const []const u8, caps: Caps }; /// One row a generation; every part of it shares the instructions the kernels use. const table = [_]Row{ - .{ .names = &.{ "gfx1030", "gfx1031", "gfx1032", "gfx1033", "gfx1034", "gfx1035", "gfx1036" }, .caps = .{ .wave = 32, .dot2_f16 = true, .dot2_bf16 = false, .sdot4 = true, .sdot8 = true, .matrix = .none, .act = .f16, .generation = .rdna2 } }, - .{ .names = &.{ "gfx1100", "gfx1101", "gfx1102", "gfx1103" }, .caps = .{ .wave = 32, .dot2_f16 = true, .dot2_bf16 = true, .sdot4 = true, .sdot8 = true, .matrix = .wmma11, .act = .bf16, .generation = .rdna3 } }, - .{ .names = &.{ "gfx1150", "gfx1151", "gfx1152", "gfx1153" }, .caps = .{ .wave = 32, .dot2_f16 = true, .dot2_bf16 = true, .sdot4 = true, .sdot8 = true, .matrix = .wmma11, .act = .bf16, .generation = .rdna3_5 } }, - .{ .names = &.{ "gfx1200", "gfx1201" }, .caps = .{ .wave = 32, .dot2_f16 = true, .dot2_bf16 = true, .sdot4 = true, .sdot8 = true, .matrix = .wmma12, .act = .bf16, .generation = .rdna4 } }, + .{ .names = &.{ "gfx1030", "gfx1031", "gfx1032", "gfx1033", "gfx1034", "gfx1035", "gfx1036" }, .caps = .{ .wave = 32, .dot2_f16 = true, .dot2_bf16 = false, .sdot4 = true, .sdot8 = true, .matrix = .none, .act = .f16, .dot2_signed_zero = false, .generation = .rdna2 } }, + .{ .names = &.{ "gfx1100", "gfx1101", "gfx1102", "gfx1103" }, .caps = .{ .wave = 32, .dot2_f16 = true, .dot2_bf16 = true, .sdot4 = true, .sdot8 = true, .matrix = .wmma11, .act = .bf16, .dot2_signed_zero = false, .generation = .rdna3 } }, + .{ .names = &.{ "gfx1150", "gfx1151", "gfx1152", "gfx1153" }, .caps = .{ .wave = 32, .dot2_f16 = true, .dot2_bf16 = true, .sdot4 = true, .sdot8 = true, .matrix = .wmma11, .act = .bf16, .dot2_signed_zero = true, .generation = .rdna3_5 } }, + .{ .names = &.{ "gfx1200", "gfx1201" }, .caps = .{ .wave = 32, .dot2_f16 = true, .dot2_bf16 = true, .sdot4 = true, .sdot8 = true, .matrix = .wmma12, .act = .bf16, .dot2_signed_zero = false, .generation = .rdna4 } }, }; test "each generation gets its instructions" { @@ -48,6 +50,8 @@ test "each generation gets its instructions" { try std.testing.expect(rdna3.dot2_bf16 and rdna3.matrix == .wmma11 and rdna3.act == .bf16); try std.testing.expectEqual(Generation.rdna3_5, Caps.of("gfx1151:xnack-").?.generation); try std.testing.expectEqual(Matrix.wmma12, Caps.of("gfx1201").?.matrix); + try std.testing.expect(Caps.of("gfx1152").?.dot2_signed_zero and !Caps.of("gfx1201").?.dot2_signed_zero); + try std.testing.expect(!rdna2.dot2_signed_zero and !rdna3.dot2_signed_zero); for (table) |row| try std.testing.expectEqual(@as(u8, 32), row.caps.wave); } From e5d49e281649126a1089d2a25d0e7da33093062d Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Fri, 6 Nov 2026 15:44:30 +0000 Subject: [PATCH 03/36] fix(hip): keep main's prompt tile off RDNA 3.5, with an edge pass The per-row V mask changes no bits on gfx1201 but cost 0.8 to 0.9% of its prompt speed (#2357 review). The tile's body moves to prefill.hpp with MASKED and EDGE instances: the served tile is masked only where dot2_signed_zero holds and is main's code elsewhere, the same instructions on gfx1201. There a new edge launch follows each tile: a block scans the V rows its rows are masked from, page ids read once so it does not spill, and reruns masked only when one is non-finite. A kernel test checks float64's classes, the tile's own bytes in every unplanted dim, and chunk starts, paged and flat. --- zig/build/hip.zig | 19 +- zig/kernels/hip/attention/prefill.hip | 373 +++--------------- zig/kernels/hip/attention/prefill.hpp | 369 +++++++++++++++++ .../core/contract/hip_records_families.zig | 4 +- zig/src/core/contract/hip_records_prompt.zig | 44 ++- zig/src/core/contract/hip_records_vision.zig | 8 +- zig/src/hip/kernel_tests.zig | 1 + zig/src/hip/kernel_tests/attention.zig | 8 +- zig/src/hip/kernel_tests/attention_edge.zig | 227 +++++++++++ zig/src/hip/launches.zig | 10 + zig/src/hip/launches/attention.zig | 3 + 11 files changed, 710 insertions(+), 356 deletions(-) create mode 100644 zig/kernels/hip/attention/prefill.hpp create mode 100644 zig/src/hip/kernel_tests/attention_edge.zig diff --git a/zig/build/hip.zig b/zig/build/hip.zig index 41d140fe47..fcbf5188e4 100644 --- a/zig/build/hip.zig +++ b/zig/build/hip.zig @@ -417,16 +417,15 @@ const groups = [_]Source{ /// 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", "quant/bf16_decoder.hpp", "quant/dense_tiles.hpp", "tiles/wmma12_gemm.hpp", - "ops/columns.hpp", "columns_controls.hpp", "ops/vision.hpp", "ops/dense_text.hpp", + "attention/attention_fa.hip", "attention/prefill.hpp", "decode/pages.hpp", "decode/plan.hpp", "ops/attention.hpp", + "ops/common.hpp", "ops/draw.hpp", "ops/tp.hpp", "ops/elementwise.hpp", "ops/linear.hpp", + "ops/moe.hpp", "ops/norms.hpp", "ops/rope.hpp", "common/arch.hpp", "common/dot2.hpp", + "common/vec.hpp", "common/wmma.hpp", "quant/act.hpp", "quant/mlx.hpp", "quant/mlx_decoder.hpp", + "quant/mlx_pieces.hpp", "quant/mlx_tiles.hpp", "tiles/dot2.hpp", "tiles/epilogue.hpp", "tiles/gemm.hpp", + "tiles/gemm_kp.hpp", "tiles/matrix_gemm.hpp", "tiles/plan.hpp", "tiles/stream.hpp", "flashnext/streams.hpp", + "flashnext/ngram.hpp", "flashnext/indexer.hpp", "flashnext/sparse.hpp", "ops/nemotron_moe.hpp", "quant/bf16_decoder.hpp", + "quant/dense_tiles.hpp", "tiles/wmma12_gemm.hpp", "ops/columns.hpp", "columns_controls.hpp", "ops/vision.hpp", + "ops/dense_text.hpp", }; /// The flags of the kernels' first build: no contraction, wave32 on RDNA, C++20. diff --git a/zig/kernels/hip/attention/prefill.hip b/zig/kernels/hip/attention/prefill.hip index 5d6d4c8405..10eb7c92ff 100644 --- a/zig/kernels/hip/attention/prefill.hip +++ b/zig/kernels/hip/attention/prefill.hip @@ -1,340 +1,22 @@ // Prefill attention for long prompts: 64 query rows share a 64-key tile; dot2 products and online softmax in fp32. -#include -#include -#include +#include "attention/prefill.hpp" -#include "decode/pages.hpp" - -namespace { - -constexpr int kRows = 64; // query rows a block -constexpr int kKeys = 64; // keys a tile -constexpr int kChunk = 64; // head dimensions a staged key or value chunk -constexpr int kThreads = 256; -using word4 = unsigned __attribute__((ext_vector_type(4))); // a 16-byte load held as one value -static_assert(kKeys == kPageTokens, "a key tile is one page"); - -// A cache element type's pair for dot2, the rounding into it and the product. -template -struct Dot; - -template <> -struct Dot<__half> { - using elem = _Float16; - using pair = _Float16 __attribute__((ext_vector_type(2))); - static __device__ __forceinline__ float dot(pair a, pair b, float acc) { - return __builtin_amdgcn_fdot2(a, b, acc, false); - } -}; - -template <> -struct Dot { - using elem = __bf16; - using pair = __bf16 __attribute__((ext_vector_type(2))); - static __device__ __forceinline__ float dot(pair a, pair b, float acc) { -#if TF_DOT2_BF16 - return __builtin_amdgcn_fdot2_f32_bf16(a, b, acc, false); -#else - // RDNA2 has no bf16 dot2: its caches are fp16, this only keeps the kernel building - return fmaf(static_cast(a.y), static_cast(b.y), fmaf(static_cast(a.x), static_cast(b.x), acc)); -#endif - } -}; - -// A 16-byte LDS read as four pairs. -template -__device__ __forceinline__ void pairs4(const void* at, Pair out[4]) { - const uint4 w = *reinterpret_cast(at); - out[0] = __builtin_bit_cast(Pair, w.x); - out[1] = __builtin_bit_cast(Pair, w.y); - out[2] = __builtin_bit_cast(Pair, w.z); - out[3] = __builtin_bit_cast(Pair, w.w); -} - -// A 16-byte load of row `row` at element `col`; unless FULL, a row from `rows` on reads row `rows` - 1 instead. -template -__device__ __forceinline__ word4 load16(const Cache* base, long long stride, int row, int rows, int col) { - return *reinterpret_cast(base + (FULL || row < rows ? row : rows - 1) * stride + col); -} - -// A staged piece, zero for a row past the span. -__device__ __forceinline__ uint4 kept(word4 w, bool in) { - return __builtin_bit_cast(uint4, in ? w : word4{0, 0, 0, 0}); -} - -// x from lane ^ M of its 16-lane row (M < 16): a DPP row_xmask move, all lanes active. -template -__device__ __forceinline__ float row_xor(float x) { - return __builtin_bit_cast(float, __builtin_amdgcn_mov_dpp(__builtin_bit_cast(int, x), 0x160 + M, 0xf, 0xf, false)); -} - -// A 16-lane row's max and sum by xor butterflies over lane distances 8, 4, 2, 1, each lane taking own then partner. -__device__ __forceinline__ float row_max16(float x) { - x = fmaxf(x, row_xor<8>(x)); - x = fmaxf(x, row_xor<4>(x)); - x = fmaxf(x, row_xor<2>(x)); - return fmaxf(x, row_xor<1>(x)); -} - -__device__ __forceinline__ float row_sum16(float x) { - x += row_xor<8>(x); - x += row_xor<4>(x); - x += row_xor<2>(x); - return x + row_xor<1>(x); -} - -// One lane's x on every lane of its 16-lane row, `sel` naming that lane in each nibble: v_permlane16. -template -__device__ __forceinline__ P row_lane(P x, unsigned sel) { - const int w = __builtin_bit_cast(int, x); - return __builtin_bit_cast(P, __builtin_amdgcn_permlane16(w, w, sel, sel, false, false)); -} - -// q rounds to the cache type once into LDS, p stays in registers; PAGED: pooled k, v; !CAUSAL: rows see every key. -template -__device__ void fa_dot2(const float* __restrict__ q, const Cache* __restrict__ k, const Cache* __restrict__ v, - float* __restrict__ out, int qlen, int span, int heads, int kv_heads, int d, float scale, - int q_pos0, long long k_sb, long long k_sh, long long k_ss, long long v_sb, long long v_sh, - long long v_ss, const unsigned* table, int count) { - constexpr int kQChunks = 256 / 8; // 16-byte chunks a q row - constexpr int kKChunks = kChunk / 8; // a staged key row's - __shared__ uint4 sQ[kRows * kQChunks]; - __shared__ uint4 sKV[kKeys * kKChunks]; // a key chunk (key, 8 dims) or a value chunk (key pair, 4 dims) - using T = Dot; - using elem = typename T::elem; - using pair_t = typename T::pair; - - const int q0 = blockIdx.x * kRows; - const int head = blockIdx.y; - const int batch = blockIdx.z; - if (q0 >= qlen) return; - const int tid = threadIdx.x; - int q_hi = qlen - q0; - if (q_hi > kRows) q_hi = kRows; - int max_vis = CAUSAL ? q_pos0 + q0 + q_hi : span; - if (max_vis > span) max_vis = span; - const int kv = head / (heads / kv_heads); - const float* qrows = q + ((static_cast(batch) * heads + head) * qlen + q0) * d; - float* orows = out + ((static_cast(batch) * heads + head) * qlen + q0) * d; - const Cache* kbase = PAGED ? k : k + static_cast(batch) * k_sb + static_cast(kv) * k_sh; - const Cache* vbase = PAGED ? v : v + static_cast(batch) * v_sb + static_cast(kv) * v_sh; - const long long kstride = PAGED ? d : k_ss; - const long long vstride = PAGED ? d : v_ss; - const int chunks = d / 8; - - for (int e = tid; e < kRows * chunks; e += kThreads) { - const int row = e / chunks, c = e % chunks; - elem h[8]; -#pragma unroll - for (int m = 0; m < 8; ++m) - h[m] = static_cast(row < q_hi ? qrows[static_cast(row) * d + 8 * c + m] : 0.f); - sQ[row * kQChunks + (c ^ (row & 7))] = __builtin_bit_cast(uint4, h); - } - - // rows sr + 16i: scores against keys 4sk .. 4sk + 3, PV at dims 4sk .. 4sk + 3 of each chunk - const int sr = tid / 16, sk = tid % 16; - float m_run[4], l_run[4]; -#pragma unroll - for (int i = 0; i < 4; ++i) m_run[i] = -INFINITY, l_run[i] = 0.f; - float o[4][16]; -#pragma unroll - for (int i = 0; i < 4; ++i) -#pragma unroll - for (int m = 0; m < 16; ++m) o[i][m] = 0.f; - - // a tile ending by full_end is in range and visible to every stored row, so FULL drops its guards and masks - int full_end = CAUSAL ? q_pos0 + q0 + 1 : span; - if (full_end > span) full_end = span; - // the first key row of the tile at k0 (on page `pg` when paged) in the keys' and the values' storage - const auto ktile_at = [&](int k0, unsigned pg) __attribute__((always_inline)) { - return PAGED ? kbase + page_row(pg, count, kv, 0, d) : kbase + static_cast(k0) * k_ss; - }; - const auto vtile_at = [&](int k0, unsigned pg) __attribute__((always_inline)) { - return PAGED ? vbase + page_row(pg, count, kv, 0, d) : vbase + static_cast(k0) * v_ss; - }; - // each staged chunk's loads issue before the previous chunk's products, which hide their latency - word4 kst[2], va, vb; - const auto load_k = [&](const Cache* ktile, int rows, int c0) __attribute__((always_inline)) { -#pragma unroll - for (int t = 0; t < 2; ++t) { - const int e = tid + t * kThreads, key = e / kKChunks, c = e % kKChunks; - kst[t] = load16(ktile, kstride, key, rows, 8 * (c0 + c)); - } - }; - // a value chunk: key pair `pair`, dims 8g .. 8g + 7 of both keys - const auto load_v = [&](const Cache* vtile, int rows, int c0) __attribute__((always_inline)) { - const int pair = tid / kKChunks, g = tid % kKChunks; - va = load16(vtile, vstride, 2 * pair, rows, 8 * (c0 + g)); - vb = load16(vtile, vstride, 2 * pair + 1, rows, 8 * (c0 + g)); - }; - unsigned page = PAGED && max_vis > 0 ? table[0] : 0; - if (max_vis > 0) load_k.template operator()(ktile_at(0, page), span, 0); - const auto tile = [&](int k0) __attribute__((always_inline)) { - const Cache* ktile = ktile_at(k0, page); - const Cache* vtile = vtile_at(k0, page); - const int rows = span - k0; - const bool more = k0 + kKeys < max_vis; - unsigned next = page; - float acc[4][4]; -#pragma unroll - for (int i = 0; i < 4; ++i) -#pragma unroll - for (int j = 0; j < 4; ++j) acc[i][j] = 0.f; - for (int c0 = 0; c0 < chunks; c0 += kKChunks) { - __syncthreads(); -#pragma unroll - for (int t = 0; t < 2; ++t) { - const int e = tid + t * kThreads, key = e / kKChunks, c = e % kKChunks; - sKV[key * kKChunks + (c ^ ((key >> 2) & 7))] = kept(kst[t], FULL || key < rows); - } - __syncthreads(); - if (PAGED && c0 == 0 && more) { - // a vector index keeps the page id out of scalar registers, so its load is awaited only where used - int at_page = (k0 + kKeys) / kPageTokens; - asm volatile("" : "+v"(at_page)); - next = table[at_page]; - } - if (c0 + kKChunks < chunks) load_k.template operator()(ktile, rows, c0 + kKChunks); - else load_v.template operator()(vtile, rows, 0); -#pragma unroll 2 - for (int c = 0; c < kKChunks; ++c) { - pair_t qv[4][4], kf[4][4]; -#pragma unroll - for (int i = 0; i < 4; ++i) { - const int row = sr + 16 * i; - pairs4(&sQ[row * kQChunks + ((c0 + c) ^ (row & 7))], qv[i]); - } -#pragma unroll - for (int j = 0; j < 4; ++j) { - const int key = 4 * sk + j; - pairs4(&sKV[key * kKChunks + (c ^ ((key >> 2) & 7))], kf[j]); - } -#pragma unroll - for (int i = 0; i < 4; ++i) -#pragma unroll - for (int j = 0; j < 4; ++j) -#pragma unroll - for (int m = 0; m < 4; ++m) acc[i][j] = T::dot(qv[i][m], kf[j][m], acc[i][j]); - } - } - // the online softmax in registers: a row's 64 keys sit on its 16 lanes, key pair pp's p on lane pp / 2 - pair_t ph[4][2]; - float alpha_r[4]; - int seen[4]; // the keys of this tile each row sees: none to 0, all from 64 -#pragma unroll - for (int i = 0; i < 4; ++i) { - const int row = sr + 16 * i; - int vis = CAUSAL ? q_pos0 + q0 + row + 1 : span; - if (vis > span) vis = span; - seen[i] = vis - k0; - float tile_max = -INFINITY; - int nan_row = 0; // !CAUSAL: a NaN score makes its row NaN, as in Torch -#pragma unroll - for (int j = 0; j < 4; ++j) { - acc[i][j] = (FULL || (row < q_hi && k0 + 4 * sk + j < vis)) ? acc[i][j] * scale : -INFINITY; - if constexpr (!CAUSAL) nan_row |= isnan(acc[i][j]); - tile_max = fmaxf(tile_max, acc[i][j]); - } - tile_max = row_max16(tile_max); - if constexpr (!CAUSAL) nan_row = row_max16(nan_row ? 1.f : 0.f) > 0.f; - const float next_m = nan_row ? NAN : fmaxf(m_run[i], tile_max); - const float alpha = next_m == -INFINITY ? 1.f : expf(m_run[i] - next_m); - float p[4], sum = 0.f; -#pragma unroll - for (int j = 0; j < 4; ++j) { - p[j] = next_m == -INFINITY ? 0.f : expf(acc[i][j] - next_m); - sum += p[j]; - } - sum = row_sum16(sum); - l_run[i] = alpha * l_run[i] + sum; - m_run[i] = next_m; - ph[i][0] = pair_t{static_cast(p[0]), static_cast(p[1])}; - ph[i][1] = pair_t{static_cast(p[2]), static_cast(p[3])}; - alpha_r[i] = alpha; - } -#pragma unroll - for (int at = 0; at < 256 / kChunk; ++at) { - if (at * kKChunks >= chunks) break; - // the chunk's key pairs as (key, key + 1) pairs at each of its dims - const int pair = tid / kKChunks, g = tid % kKChunks; - using elem8 = elem __attribute__((ext_vector_type(8))); - // a value past the span reads the span's last (load16); the products zero each key past a row's edge - const elem8 ha = __builtin_bit_cast(elem8, va); - const elem8 hb = __builtin_bit_cast(elem8, vb); - pair_t lo[4], hi[4]; -#pragma unroll - for (int m = 0; m < 4; ++m) lo[m] = pair_t{ha[m], hb[m]}, hi[m] = pair_t{ha[4 + m], hb[4 + m]}; - __syncthreads(); - sKV[pair * 16 + 2 * g] = __builtin_bit_cast(uint4, lo); - sKV[pair * 16 + 2 * g + 1] = __builtin_bit_cast(uint4, hi); - __syncthreads(); - if ((at + 1) * kKChunks < chunks) load_v.template operator()(vtile, rows, (at + 1) * kKChunks); - else if (more) load_k.template operator()(ktile_at(k0 + kKeys, next), rows - kKeys, 0); - float pv[4][4]; -#pragma unroll - for (int i = 0; i < 4; ++i) -#pragma unroll - for (int m = 0; m < 4; ++m) pv[i][m] = 0.f; -#pragma unroll 4 - for (int pp = 0; pp < kKeys / 2; ++pp) { - pair_t vv[4]; - pairs4(&sKV[pp * 16 + sk], vv); - const unsigned lane = (pp / 2) * 0x11111111u; -#pragma unroll - for (int i = 0; i < 4; ++i) { - const pair_t pp4 = row_lane(ph[i][pp % 2], lane); - // a key past the row's edge takes value 0: RDNA 3.5's dot2 rounds a -0 lane (0 x v < 0) off +0's - const int left = FULL ? 2 : seen[i] - 2 * pp; - const unsigned keep = left >= 2 ? ~0u : left == 1 ? 0xffffu : 0u; // key 2pp is the low half -#pragma unroll - for (int m = 0; m < 4; ++m) - pv[i][m] = T::dot(pp4, __builtin_bit_cast(pair_t, __builtin_bit_cast(unsigned, vv[m]) & keep), - pv[i][m]); - } - } -#pragma unroll - for (int i = 0; i < 4; ++i) -#pragma unroll - for (int m = 0; m < 4; ++m) { - float& acc_o = o[i][4 * at + m]; - acc_o = fmaf(alpha_r[i], acc_o, pv[i][m]); - } - } - page = next; - }; - int k0 = 0; - for (; k0 + kKeys <= full_end; k0 += kKeys) tile.template operator()(k0); - for (; k0 < max_vis; k0 += kKeys) tile.template operator()(k0); -#pragma unroll - for (int i = 0; i < 4; ++i) { - const int row = sr + 16 * i; - if (row >= q_hi) continue; - const float inv = l_run[i] == 0.f ? 0.f : 1.f / l_run[i]; -#pragma unroll - for (int at = 0; at < 256 / kChunk; ++at) { - if (at * kKChunks >= chunks) break; -#pragma unroll - for (int m = 0; m < 4; ++m) - orows[static_cast(row) * d + 64 * at + 4 * sk + m] = o[i][4 * at + m] * inv; - } - } -} - -} // namespace +// The served tile masks each key past a row's edge only where dot2 rounds a -0 lane apart (RDNA 3.5). +constexpr bool kMasked = TF_DOT2_SIGNED_ZERO != 0; extern "C" __global__ void __launch_bounds__(kThreads, 1) tf_fa_wide_f16(const float* q, const __half* k, const __half* v, float* out, int qlen, int span, int heads, int kv_heads, int d, float scale, int q_pos0, long long k_sb, long long k_sh, long long k_ss, long long v_sb, long long v_sh, long long v_ss) { - fa_dot2<__half, false>(q, k, v, out, qlen, span, heads, kv_heads, d, scale, q_pos0, k_sb, k_sh, k_ss, v_sb, v_sh, v_ss, nullptr, 0); + fa_dot2<__half, false, true, kMasked, false>(q, k, v, out, qlen, span, heads, kv_heads, d, scale, q_pos0, k_sb, k_sh, k_ss, v_sb, v_sh, v_ss, nullptr, 0); } extern "C" __global__ void __launch_bounds__(kThreads, 1) tf_fa_wide_bf16(const float* q, const hip_bfloat16* k, const hip_bfloat16* v, float* out, int qlen, int span, int heads, int kv_heads, int d, float scale, int q_pos0, long long k_sb, long long k_sh, long long k_ss, long long v_sb, long long v_sh, long long v_ss) { - fa_dot2(q, k, v, out, qlen, span, heads, kv_heads, d, scale, q_pos0, k_sb, k_sh, k_ss, v_sb, + fa_dot2(q, k, v, out, qlen, span, heads, kv_heads, d, scale, q_pos0, k_sb, k_sh, k_ss, v_sb, v_sh, v_ss, nullptr, 0); } @@ -343,7 +25,7 @@ extern "C" __global__ void __launch_bounds__(kThreads, 1) tf_fa_full_bf16(const float* q, const hip_bfloat16* k, const hip_bfloat16* v, float* out, int qlen, int span, int heads, int kv_heads, int d, float scale, int q_pos0, long long k_sb, long long k_sh, long long k_ss, long long v_sb, long long v_sh, long long v_ss) { - fa_dot2(q, k, v, out, qlen, span, heads, kv_heads, d, scale, q_pos0, k_sb, + fa_dot2(q, k, v, out, qlen, span, heads, kv_heads, d, scale, q_pos0, k_sb, k_sh, k_ss, v_sb, v_sh, v_ss, nullptr, 0); } @@ -351,11 +33,50 @@ tf_fa_full_bf16(const float* q, const hip_bfloat16* k, const hip_bfloat16* v, fl extern "C" __global__ void __launch_bounds__(kThreads, 1) tf_fa_paged_f16(const float* q, const __half* k, const __half* v, float* out, int qlen, int span, int heads, int kv_heads, int d, float scale, int q_pos0, const unsigned* table, int count) { - fa_dot2<__half, true>(q, k, v, out, qlen, span, heads, kv_heads, d, scale, q_pos0, 0, 0, 0, 0, 0, 0, table, count); + fa_dot2<__half, true, true, kMasked, false>(q, k, v, out, qlen, span, heads, kv_heads, d, scale, q_pos0, 0, 0, 0, 0, 0, 0, table, count); } extern "C" __global__ void __launch_bounds__(kThreads, 1) tf_fa_paged_bf16(const float* q, const hip_bfloat16* k, const hip_bfloat16* v, float* out, int qlen, int span, int heads, int kv_heads, int d, float scale, int q_pos0, const unsigned* table, int count) { - fa_dot2(q, k, v, out, qlen, span, heads, kv_heads, d, scale, q_pos0, 0, 0, 0, 0, 0, 0, table, count); + fa_dot2(q, k, v, out, qlen, span, heads, kv_heads, d, scale, q_pos0, 0, 0, 0, 0, 0, 0, table, count); +} + +// Edge pass: a block reruns masked when a key it masks holds a non-finite value; a trap where the tile is masked. +template +__device__ __forceinline__ void fa_edge(A... a) { +#if TF_DOT2_SIGNED_ZERO + __builtin_trap(); +#else + fa_dot2(a...); +#endif +} + +extern "C" __global__ void __launch_bounds__(kThreads, 1) +tf_fa_wide_edge_f16(const float* q, const __half* k, const __half* v, float* out, int qlen, int span, int heads, + int kv_heads, int d, float scale, int q_pos0, long long k_sb, long long k_sh, long long k_ss, + long long v_sb, long long v_sh, long long v_ss) { + fa_edge<__half, false>(q, k, v, out, qlen, span, heads, kv_heads, d, scale, q_pos0, k_sb, k_sh, k_ss, v_sb, v_sh, + v_ss, nullptr, 0); +} + +extern "C" __global__ void __launch_bounds__(kThreads, 1) +tf_fa_wide_edge_bf16(const float* q, const hip_bfloat16* k, const hip_bfloat16* v, float* out, int qlen, int span, + int heads, int kv_heads, int d, float scale, int q_pos0, long long k_sb, long long k_sh, + long long k_ss, long long v_sb, long long v_sh, long long v_ss) { + fa_edge(q, k, v, out, qlen, span, heads, kv_heads, d, scale, q_pos0, k_sb, k_sh, k_ss, v_sb, + v_sh, v_ss, nullptr, 0); +} + +extern "C" __global__ void __launch_bounds__(kThreads, 1) +tf_fa_paged_edge_f16(const float* q, const __half* k, const __half* v, float* out, int qlen, int span, int heads, + int kv_heads, int d, float scale, int q_pos0, const unsigned* table, int count) { + fa_edge<__half, true>(q, k, v, out, qlen, span, heads, kv_heads, d, scale, q_pos0, 0, 0, 0, 0, 0, 0, table, count); +} + +extern "C" __global__ void __launch_bounds__(kThreads, 1) +tf_fa_paged_edge_bf16(const float* q, const hip_bfloat16* k, const hip_bfloat16* v, float* out, int qlen, int span, + int heads, int kv_heads, int d, float scale, int q_pos0, const unsigned* table, int count) { + fa_edge(q, k, v, out, qlen, span, heads, kv_heads, d, scale, q_pos0, 0, 0, 0, 0, 0, 0, table, + count); } diff --git a/zig/kernels/hip/attention/prefill.hpp b/zig/kernels/hip/attention/prefill.hpp new file mode 100644 index 0000000000..f926ca173a --- /dev/null +++ b/zig/kernels/hip/attention/prefill.hpp @@ -0,0 +1,369 @@ +// The 64-row prefill tile's body, shared by its entry points (prefill.hip) and the contract's controls. +#pragma once + +#include +#include +#include + +#include "decode/pages.hpp" + +#ifndef TF_DOT2_SIGNED_ZERO +#error "the caps macros come from the build" +#endif + +namespace { + +constexpr int kRows = 64; // query rows a block +constexpr int kKeys = 64; // keys a tile +constexpr int kChunk = 64; // head dimensions a staged key or value chunk +constexpr int kThreads = 256; +using word4 = unsigned __attribute__((ext_vector_type(4))); // a 16-byte load held as one value +static_assert(kKeys == kPageTokens, "a key tile is one page"); + +// A cache element type's pair for dot2, the rounding into it and the product. +template +struct Dot; + +template <> +struct Dot<__half> { + using elem = _Float16; + static constexpr unsigned top = 0x7c00u; // the exponent bits, all set in an infinity or NaN + using pair = _Float16 __attribute__((ext_vector_type(2))); + static __device__ __forceinline__ float dot(pair a, pair b, float acc) { + return __builtin_amdgcn_fdot2(a, b, acc, false); + } +}; + +template <> +struct Dot { + using elem = __bf16; + static constexpr unsigned top = 0x7f80u; + using pair = __bf16 __attribute__((ext_vector_type(2))); + static __device__ __forceinline__ float dot(pair a, pair b, float acc) { +#if TF_DOT2_BF16 + return __builtin_amdgcn_fdot2_f32_bf16(a, b, acc, false); +#else + // RDNA2 has no bf16 dot2: its caches are fp16, this only keeps the kernel building + return fmaf(static_cast(a.y), static_cast(b.y), fmaf(static_cast(a.x), static_cast(b.x), acc)); +#endif + } +}; + +// A 16-byte LDS read as four pairs. +template +__device__ __forceinline__ void pairs4(const void* at, Pair out[4]) { + const uint4 w = *reinterpret_cast(at); + out[0] = __builtin_bit_cast(Pair, w.x); + out[1] = __builtin_bit_cast(Pair, w.y); + out[2] = __builtin_bit_cast(Pair, w.z); + out[3] = __builtin_bit_cast(Pair, w.w); +} + +// A 16-byte load of row `row` at element `col`; unless FULL, a row from `rows` on reads row `rows` - 1 instead. +template +__device__ __forceinline__ word4 load16(const Cache* base, long long stride, int row, int rows, int col) { + return *reinterpret_cast(base + (FULL || row < rows ? row : rows - 1) * stride + col); +} + +// A staged piece, zero for a row past the span. +__device__ __forceinline__ uint4 kept(word4 w, bool in) { + return __builtin_bit_cast(uint4, in ? w : word4{0, 0, 0, 0}); +} + +// x from lane ^ M of its 16-lane row (M < 16): a DPP row_xmask move, all lanes active. +template +__device__ __forceinline__ float row_xor(float x) { + return __builtin_bit_cast(float, __builtin_amdgcn_mov_dpp(__builtin_bit_cast(int, x), 0x160 + M, 0xf, 0xf, false)); +} + +// A 16-lane row's max and sum by xor butterflies over lane distances 8, 4, 2, 1, each lane taking own then partner. +__device__ __forceinline__ float row_max16(float x) { + x = fmaxf(x, row_xor<8>(x)); + x = fmaxf(x, row_xor<4>(x)); + x = fmaxf(x, row_xor<2>(x)); + return fmaxf(x, row_xor<1>(x)); +} + +__device__ __forceinline__ float row_sum16(float x) { + x += row_xor<8>(x); + x += row_xor<4>(x); + x += row_xor<2>(x); + return x + row_xor<1>(x); +} + +// One lane's x on every lane of its 16-lane row, `sel` naming that lane in each nibble: v_permlane16. +template +__device__ __forceinline__ P row_lane(P x, unsigned sel) { + const int w = __builtin_bit_cast(int, x); + return __builtin_bit_cast(P, __builtin_amdgcn_permlane16(w, w, sel, sel, false, false)); +} + +// Either 16-bit half of `w` has every exponent bit of `top` set: an infinity or NaN. +__device__ __forceinline__ bool nonfinite2(unsigned w, unsigned top) { + return (w & top) == top || (w & (top << 16)) == (top << 16); +} + +// q rounds to the cache type once into LDS, p stays in registers; PAGED: pooled k, v; !CAUSAL: rows see every key. +template +__device__ void fa_dot2(const float* __restrict__ q, const Cache* __restrict__ k, const Cache* __restrict__ v, + float* __restrict__ out, int qlen, int span, int heads, int kv_heads, int d, float scale, + int q_pos0, long long k_sb, long long k_sh, long long k_ss, long long v_sb, long long v_sh, + long long v_ss, const unsigned* table, int count) { + constexpr int kQChunks = 256 / 8; // 16-byte chunks a q row + constexpr int kKChunks = kChunk / 8; // a staged key row's + __shared__ uint4 sQ[kRows * kQChunks]; + __shared__ uint4 sKV[kKeys * kKChunks]; // a key chunk (key, 8 dims) or a value chunk (key pair, 4 dims) + using T = Dot; + using elem = typename T::elem; + using pair_t = typename T::pair; + + const int q0 = blockIdx.x * kRows; + const int head = blockIdx.y; + const int batch = blockIdx.z; + if (q0 >= qlen) return; + const int tid = threadIdx.x; + int q_hi = qlen - q0; + if (q_hi > kRows) q_hi = kRows; + int max_vis = CAUSAL ? q_pos0 + q0 + q_hi : span; + if (max_vis > span) max_vis = span; + const int kv = head / (heads / kv_heads); + const float* qrows = q + ((static_cast(batch) * heads + head) * qlen + q0) * d; + float* orows = out + ((static_cast(batch) * heads + head) * qlen + q0) * d; + const Cache* kbase = PAGED ? k : k + static_cast(batch) * k_sb + static_cast(kv) * k_sh; + const Cache* vbase = PAGED ? v : v + static_cast(batch) * v_sb + static_cast(kv) * v_sh; + const long long kstride = PAGED ? d : k_ss; + const long long vstride = PAGED ? d : v_ss; + const int chunks = d / 8; + + if constexpr (EDGE) { + // the keys a live row is masked from: past row 0's edge, up to the end of the last tile read + const int lo = q_pos0 + q0 + 1; + int hi = (max_vis + kKeys - 1) / kKeys * kKeys; + if (hi > span) hi = span; + // at most 126 keys span three pages: their ids load once, uniform, before the per-thread reads + const int p0 = lo / kPageTokens; + unsigned pg[3] = {0, 0, 0}; + if constexpr (PAGED) { +#pragma unroll + for (int i = 0; i < 3; ++i) + if (lo < hi && p0 + i <= (hi - 1) / kPageTokens) pg[i] = table[p0 + i]; + } + int bad = 0; + for (int e = tid; e < (hi - lo) * chunks; e += kThreads) { + const int key = lo + e / chunks, c = e % chunks; + const int at = key / kPageTokens - p0; + const unsigned id = at == 0 ? pg[0] : at == 1 ? pg[1] : pg[2]; + const Cache* row = PAGED ? vbase + page_row(id, count, kv, key % kPageTokens, d) : vbase + key * vstride; + const word4 w = *reinterpret_cast(row + 8 * c); +#pragma unroll + for (int j = 0; j < 4; ++j) bad |= nonfinite2(w[j], T::top); + } + if (!__syncthreads_or(bad)) return; + } + + for (int e = tid; e < kRows * chunks; e += kThreads) { + const int row = e / chunks, c = e % chunks; + elem h[8]; +#pragma unroll + for (int m = 0; m < 8; ++m) + h[m] = static_cast(row < q_hi ? qrows[static_cast(row) * d + 8 * c + m] : 0.f); + sQ[row * kQChunks + (c ^ (row & 7))] = __builtin_bit_cast(uint4, h); + } + + // rows sr + 16i: scores against keys 4sk .. 4sk + 3, PV at dims 4sk .. 4sk + 3 of each chunk + const int sr = tid / 16, sk = tid % 16; + float m_run[4], l_run[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) m_run[i] = -INFINITY, l_run[i] = 0.f; + float o[4][16]; +#pragma unroll + for (int i = 0; i < 4; ++i) +#pragma unroll + for (int m = 0; m < 16; ++m) o[i][m] = 0.f; + + // a tile ending by full_end is in range and visible to every stored row, so FULL drops its guards and masks + int full_end = CAUSAL ? q_pos0 + q0 + 1 : span; + if (full_end > span) full_end = span; + // the first key row of the tile at k0 (on page `pg` when paged) in the keys' and the values' storage + const auto ktile_at = [&](int k0, unsigned pg) __attribute__((always_inline)) { + return PAGED ? kbase + page_row(pg, count, kv, 0, d) : kbase + static_cast(k0) * k_ss; + }; + const auto vtile_at = [&](int k0, unsigned pg) __attribute__((always_inline)) { + return PAGED ? vbase + page_row(pg, count, kv, 0, d) : vbase + static_cast(k0) * v_ss; + }; + // each staged chunk's loads issue before the previous chunk's products, which hide their latency + word4 kst[2], va, vb; + const auto load_k = [&](const Cache* ktile, int rows, int c0) __attribute__((always_inline)) { +#pragma unroll + for (int t = 0; t < 2; ++t) { + const int e = tid + t * kThreads, key = e / kKChunks, c = e % kKChunks; + kst[t] = load16(ktile, kstride, key, rows, 8 * (c0 + c)); + } + }; + // a value chunk: key pair `pair`, dims 8g .. 8g + 7 of both keys + const auto load_v = [&](const Cache* vtile, int rows, int c0) __attribute__((always_inline)) { + const int pair = tid / kKChunks, g = tid % kKChunks; + va = load16(vtile, vstride, 2 * pair, rows, 8 * (c0 + g)); + vb = load16(vtile, vstride, 2 * pair + 1, rows, 8 * (c0 + g)); + }; + unsigned page = PAGED && max_vis > 0 ? table[0] : 0; + if (max_vis > 0) load_k.template operator()(ktile_at(0, page), span, 0); + const auto tile = [&](int k0) __attribute__((always_inline)) { + const Cache* ktile = ktile_at(k0, page); + const Cache* vtile = vtile_at(k0, page); + const int rows = span - k0; + const bool more = k0 + kKeys < max_vis; + unsigned next = page; + float acc[4][4]; +#pragma unroll + for (int i = 0; i < 4; ++i) +#pragma unroll + for (int j = 0; j < 4; ++j) acc[i][j] = 0.f; + for (int c0 = 0; c0 < chunks; c0 += kKChunks) { + __syncthreads(); +#pragma unroll + for (int t = 0; t < 2; ++t) { + const int e = tid + t * kThreads, key = e / kKChunks, c = e % kKChunks; + sKV[key * kKChunks + (c ^ ((key >> 2) & 7))] = kept(kst[t], FULL || key < rows); + } + __syncthreads(); + if (PAGED && c0 == 0 && more) { + // a vector index keeps the page id out of scalar registers, so its load is awaited only where used + int at_page = (k0 + kKeys) / kPageTokens; + asm volatile("" : "+v"(at_page)); + next = table[at_page]; + } + if (c0 + kKChunks < chunks) load_k.template operator()(ktile, rows, c0 + kKChunks); + else load_v.template operator()(vtile, rows, 0); +#pragma unroll 2 + for (int c = 0; c < kKChunks; ++c) { + pair_t qv[4][4], kf[4][4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int row = sr + 16 * i; + pairs4(&sQ[row * kQChunks + ((c0 + c) ^ (row & 7))], qv[i]); + } +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int key = 4 * sk + j; + pairs4(&sKV[key * kKChunks + (c ^ ((key >> 2) & 7))], kf[j]); + } +#pragma unroll + for (int i = 0; i < 4; ++i) +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int m = 0; m < 4; ++m) acc[i][j] = T::dot(qv[i][m], kf[j][m], acc[i][j]); + } + } + // the online softmax in registers: a row's 64 keys sit on its 16 lanes, key pair pp's p on lane pp / 2 + pair_t ph[4][2]; + float alpha_r[4]; + int seen[4]; // the keys of this tile each row sees: none to 0, all from 64 +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int row = sr + 16 * i; + int vis = CAUSAL ? q_pos0 + q0 + row + 1 : span; + if (vis > span) vis = span; + if constexpr (MASKED) seen[i] = vis - k0; + float tile_max = -INFINITY; + int nan_row = 0; // !CAUSAL: a NaN score makes its row NaN, as in Torch +#pragma unroll + for (int j = 0; j < 4; ++j) { + acc[i][j] = (FULL || (row < q_hi && k0 + 4 * sk + j < vis)) ? acc[i][j] * scale : -INFINITY; + if constexpr (!CAUSAL) nan_row |= isnan(acc[i][j]); + tile_max = fmaxf(tile_max, acc[i][j]); + } + tile_max = row_max16(tile_max); + if constexpr (!CAUSAL) nan_row = row_max16(nan_row ? 1.f : 0.f) > 0.f; + const float next_m = nan_row ? NAN : fmaxf(m_run[i], tile_max); + const float alpha = next_m == -INFINITY ? 1.f : expf(m_run[i] - next_m); + float p[4], sum = 0.f; +#pragma unroll + for (int j = 0; j < 4; ++j) { + p[j] = next_m == -INFINITY ? 0.f : expf(acc[i][j] - next_m); + sum += p[j]; + } + sum = row_sum16(sum); + l_run[i] = alpha * l_run[i] + sum; + m_run[i] = next_m; + ph[i][0] = pair_t{static_cast(p[0]), static_cast(p[1])}; + ph[i][1] = pair_t{static_cast(p[2]), static_cast(p[3])}; + alpha_r[i] = alpha; + } +#pragma unroll + for (int at = 0; at < 256 / kChunk; ++at) { + if (at * kKChunks >= chunks) break; + // the chunk's key pairs as (key, key + 1) pairs at each of its dims + const int pair = tid / kKChunks, g = tid % kKChunks; + using elem8 = elem __attribute__((ext_vector_type(8))); + // MASKED: a value past the span reads the span's last (load16); the products zero each key past a row's edge + const uint4 ua = MASKED ? __builtin_bit_cast(uint4, va) : kept(va, FULL || 2 * pair < rows); + const uint4 ub = MASKED ? __builtin_bit_cast(uint4, vb) : kept(vb, FULL || 2 * pair + 1 < rows); + const elem8 ha = __builtin_bit_cast(elem8, ua); + const elem8 hb = __builtin_bit_cast(elem8, ub); + pair_t lo[4], hi[4]; +#pragma unroll + for (int m = 0; m < 4; ++m) lo[m] = pair_t{ha[m], hb[m]}, hi[m] = pair_t{ha[4 + m], hb[4 + m]}; + __syncthreads(); + sKV[pair * 16 + 2 * g] = __builtin_bit_cast(uint4, lo); + sKV[pair * 16 + 2 * g + 1] = __builtin_bit_cast(uint4, hi); + __syncthreads(); + if ((at + 1) * kKChunks < chunks) load_v.template operator()(vtile, rows, (at + 1) * kKChunks); + else if (more) load_k.template operator()(ktile_at(k0 + kKeys, next), rows - kKeys, 0); + float pv[4][4]; +#pragma unroll + for (int i = 0; i < 4; ++i) +#pragma unroll + for (int m = 0; m < 4; ++m) pv[i][m] = 0.f; +#pragma unroll 4 + for (int pp = 0; pp < kKeys / 2; ++pp) { + pair_t vv[4]; + pairs4(&sKV[pp * 16 + sk], vv); + const unsigned lane = (pp / 2) * 0x11111111u; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const pair_t pp4 = row_lane(ph[i][pp % 2], lane); + if constexpr (MASKED) { + // a key past the row's edge takes value 0: RDNA 3.5's dot2 rounds a -0 lane off +0's + const int left = FULL ? 2 : seen[i] - 2 * pp; + const unsigned keep = left >= 2 ? ~0u : left == 1 ? 0xffffu : 0u; // key 2pp is the low half +#pragma unroll + for (int m = 0; m < 4; ++m) + pv[i][m] = T::dot(pp4, __builtin_bit_cast(pair_t, __builtin_bit_cast(unsigned, vv[m]) & keep), + pv[i][m]); + } else { +#pragma unroll + for (int m = 0; m < 4; ++m) pv[i][m] = T::dot(pp4, vv[m], pv[i][m]); + } + } + } +#pragma unroll + for (int i = 0; i < 4; ++i) +#pragma unroll + for (int m = 0; m < 4; ++m) { + float& acc_o = o[i][4 * at + m]; + acc_o = fmaf(alpha_r[i], acc_o, pv[i][m]); + } + } + page = next; + }; + int k0 = 0; + for (; k0 + kKeys <= full_end; k0 += kKeys) tile.template operator()(k0); + for (; k0 < max_vis; k0 += kKeys) tile.template operator()(k0); +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int row = sr + 16 * i; + if (row >= q_hi) continue; + const float inv = l_run[i] == 0.f ? 0.f : 1.f / l_run[i]; +#pragma unroll + for (int at = 0; at < 256 / kChunk; ++at) { + if (at * kKChunks >= chunks) break; +#pragma unroll + for (int m = 0; m < 4; ++m) + orows[static_cast(row) * d + 64 * at + 4 * sk + m] = o[i][4 * at + m] * inv; + } + } +} + +} // namespace diff --git a/zig/src/core/contract/hip_records_families.zig b/zig/src/core/contract/hip_records_families.zig index a7bdc2c992..d12af0a236 100644 --- a/zig/src/core/contract/hip_records_families.zig +++ b/zig/src/core/contract/hip_records_families.zig @@ -16,6 +16,6 @@ pub const all = join: { test "the HIP records stay main's; the family lists only regroup GLM-OCR's and the joint schema's" { const every = @import("hip_records.zig").records; - try std.testing.expectEqual(@as(usize, 193), every.len); // main's records, GLM-OCR's 12 and the joint schema's 10 among them - try std.testing.expectEqual(@as(usize, 171), every.len - all.len); + try std.testing.expectEqual(@as(usize, 197), every.len); // main's records, GLM-OCR's 12 and the joint schema's 10 among them + try std.testing.expectEqual(@as(usize, 175), every.len - all.len); } diff --git a/zig/src/core/contract/hip_records_prompt.zig b/zig/src/core/contract/hip_records_prompt.zig index fb938dfd28..7b28373931 100644 --- a/zig/src/core/contract/hip_records_prompt.zig +++ b/zig/src/core/contract/hip_records_prompt.zig @@ -12,27 +12,51 @@ fn promptTile(comptime symbol: []const u8, usage: [limits.rdna.len]U) Record { pub const records = [_]Record{ promptTile("tf_fa_wide_f16", .{ - .{ .vgprs = 252, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 230, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 230, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, .{ .vgprs = 236, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 236, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 236, .sgprs = 49, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 230, .sgprs = 49, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, }), promptTile("tf_fa_wide_bf16", .{ .{ .vgprs = 246, .sgprs = 46, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 230, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, .{ .vgprs = 236, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 236, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 236, .sgprs = 49, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 230, .sgprs = 49, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, }), promptTile("tf_fa_paged_f16", .{ - .{ .vgprs = 255, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 239, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 239, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 231, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 231, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, .{ .vgprs = 239, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 231, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, }), promptTile("tf_fa_paged_bf16", .{ .{ .vgprs = 249, .sgprs = 52, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 231, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, .{ .vgprs = 239, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 239, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 239, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 231, .sgprs = 53, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + }), + promptTile("tf_fa_wide_edge_f16", .{ + .{ .vgprs = 249, .sgprs = 50, .lds = 41216, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 233, .sgprs = 50, .lds = 41216, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 0, .served = false }, + .{ .vgprs = 233, .sgprs = 51, .lds = 41216, .private = 0, .bound = 256, .wave = 32 }, + }), + promptTile("tf_fa_wide_edge_bf16", .{ + .{ .vgprs = 254, .sgprs = 50, .lds = 41216, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 233, .sgprs = 50, .lds = 41216, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 0, .served = false }, + .{ .vgprs = 233, .sgprs = 51, .lds = 41216, .private = 0, .bound = 256, .wave = 32 }, + }), + promptTile("tf_fa_paged_edge_f16", .{ + .{ .vgprs = 252, .sgprs = 53, .lds = 41216, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 236, .sgprs = 53, .lds = 41216, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 0, .served = false }, + .{ .vgprs = 236, .sgprs = 53, .lds = 41216, .private = 0, .bound = 256, .wave = 32 }, + }), + promptTile("tf_fa_paged_edge_bf16", .{ + .{ .vgprs = 249, .sgprs = 52, .lds = 41216, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 236, .sgprs = 53, .lds = 41216, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 0, .served = false }, + .{ .vgprs = 236, .sgprs = 53, .lds = 41216, .private = 0, .bound = 256, .wave = 32 }, }), }; diff --git a/zig/src/core/contract/hip_records_vision.zig b/zig/src/core/contract/hip_records_vision.zig index aea8f078f5..89fd81b464 100644 --- a/zig/src/core/contract/hip_records_vision.zig +++ b/zig/src/core/contract/hip_records_vision.zig @@ -41,10 +41,10 @@ pub const records = [_]Record{ .{ .launch = .{ .name = "BF16 full image attention", .symbol = "tf_fa_full_bf16", .image = "prefill", .threads = 256 }, .usage = .{ - .{ .vgprs = 246, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 230, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 230, .sgprs = 47, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, - .{ .vgprs = 230, .sgprs = 49, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 245, .sgprs = 45, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 229, .sgprs = 46, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 229, .sgprs = 46, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, + .{ .vgprs = 229, .sgprs = 48, .lds = 40960, .private = 0, .bound = 256, .wave = 32 }, }, }, }; diff --git a/zig/src/hip/kernel_tests.zig b/zig/src/hip/kernel_tests.zig index a9bf5694cb..2e10db66ef 100644 --- a/zig/src/hip/kernel_tests.zig +++ b/zig/src/hip/kernel_tests.zig @@ -25,6 +25,7 @@ test { _ = @import("kernel_tests/chain.zig"); _ = @import("kernel_tests/attention.zig"); _ = @import("kernel_tests/attention_pin.zig"); + _ = @import("kernel_tests/attention_edge.zig"); _ = @import("kernel_tests/stats.zig"); _ = @import("kernel_tests/ops.zig"); _ = @import("kernel_tests/conv_split.zig"); diff --git a/zig/src/hip/kernel_tests/attention.zig b/zig/src/hip/kernel_tests/attention.zig index ba9e9f69fe..f767f03e30 100644 --- a/zig/src/hip/kernel_tests/attention.zig +++ b/zig/src/hip/kernel_tests/attention.zig @@ -10,13 +10,13 @@ const DeviceBuffer = @import("../memory.zig").DeviceBuffer; const page = 64; /// A cache of `kv_heads` heads and `span` positions in the activation type, flat (kv_heads, span, d) and paged. -const Cache = struct { +pub const Cache = struct { k: []u16, v: []u16, /// Position block i lives in page table[i]. table: []u32, - fn make(t: *Rig, rng: *rig.Rng, kv_heads: usize, span: usize, d: usize) !Cache { + pub fn make(t: *Rig, rng: *rig.Rng, kv_heads: usize, span: usize, d: usize) !Cache { const c: Cache = .{ .k = try gpa.alloc(u16, kv_heads * span * d), .v = try gpa.alloc(u16, kv_heads * span * d), .table = try gpa.alloc(u32, span / page) }; for (c.k) |*x| x.* = t.bits(rng.unit()); for (c.v) |*x| x.* = t.bits(rng.unit()); @@ -25,14 +25,14 @@ const Cache = struct { return c; } - fn free(c: Cache) void { + pub fn free(c: Cache) void { gpa.free(c.k); gpa.free(c.v); gpa.free(c.table); } /// The pool layout of decode/pages.hpp: ((h * count + page) * 64 + slot) * d. - fn pool(c: Cache, src: []const u16, kv_heads: usize, span: usize, d: usize) ![]u16 { + pub fn pool(c: Cache, src: []const u16, kv_heads: usize, span: usize, d: usize) ![]u16 { const out = try gpa.alloc(u16, src.len); const count = span / page; for (0..kv_heads) |h| for (0..span) |pos| { diff --git a/zig/src/hip/kernel_tests/attention_edge.zig b/zig/src/hip/kernel_tests/attention_edge.zig new file mode 100644 index 0000000000..7c868c8cbe --- /dev/null +++ b/zig/src/hip/kernel_tests/attention_edge.zig @@ -0,0 +1,227 @@ +//! The prompt tile's edge pass: a non-finite value at a key a row must not see adds nothing to that row. +const std = @import("std"); +const rig = @import("rig.zig"); +const Rig = rig.Rig; +const gpa = rig.gpa; +const at = rig.at; +const attention = @import("attention.zig"); +const Cache = attention.Cache; +const DeviceBuffer = @import("../memory.zig").DeviceBuffer; + +const heads = 32; +const kv_heads = 8; +const group = heads / kv_heads; +const d = 128; + +/// A non-finite value planted in V: its KV head, key position and dim. +const Plant = struct { kv: usize, key: usize, dim: usize, value: f32 }; + +/// The output class float64 gives a row that sees keys up to `last`: the planted values it sees decide it. +const Class = enum { finite, pos_inf, neg_inf, nan }; + +fn expected(plants: []const Plant, kv: usize, dim: usize, last: usize) Class { + var pos = false; + var neg = false; + for (plants) |p| { + if (p.kv != kv or p.dim != dim or p.key > last) continue; + if (std.math.isNan(p.value)) return .nan; + if (p.value > 0) pos = true else neg = true; + } + if (pos and neg) return .nan; + return if (pos) .pos_inf else if (neg) .neg_inf else .finite; +} + +fn classOf(x: f32) Class { + if (std.math.isNan(x)) return .nan; + if (std.math.isInf(x)) return if (x > 0) .pos_inf else .neg_inf; + return .finite; +} + +fn planted(plants: []const Plant, dim: usize) bool { + for (plants) |p| if (p.dim == dim) return true; + return false; +} + +fn plant(t: *const Rig, c: Cache, plants: []const Plant, span: usize) void { + for (plants) |p| if (p.key < span) { + c.v[(p.kv * span + p.key) * d + p.dim] = t.bits(p.value); + }; +} + +/// Served and tile-only outputs of one prompt chunk, paged and flat; out (heads, qlen, d). +const Runs = struct { paged: [2][]f32, flat: [2][]f32 }; + +fn run(t: *Rig, c: Cache, hq: []const f32, qlen: usize, span: usize, q_pos0: usize) !Runs { + const scale: f32 = 1 / @sqrt(@as(f32, @floatFromInt(d))); + const kind: c_int = if (t.bf16) 1 else 0; + var q = try t.upload(hq); + defer q.free(); + var k = try t.upload(c.k); + defer k.free(); + var v = try t.upload(c.v); + defer v.free(); + const pk = try c.pool(c.k, kv_heads, span, d); + defer gpa.free(pk); + const pv = try c.pool(c.v, kv_heads, span, d); + defer gpa.free(pv); + var pool_k = try t.upload(pk); + defer pool_k.free(); + var pool_v = try t.upload(pv); + defer pool_v.free(); + var table = try t.upload(c.table); + defer table.free(); + var out = try t.alloc(heads * qlen * d * 4); + defer out.free(); + const s = t.stream.handle; + var r: Runs = undefined; + // served (the edge pass where the launcher takes it), then the tile alone + const edge = t.on.edge; + defer t.on.edge = edge; + for ([_]bool{ edge, false }, 0..) |e, i| { + t.on.edge = e; + try t.on.pagedCausal(at(q), at(pool_k), at(pool_v), at(table), at(out), qlen, span, heads, kv_heads, d, scale, q_pos0, c.table.len, kind, s); + try t.stream.synchronize(); + r.paged[i] = try rig.download(f32, out); + try t.on.tf_causal(@ptrFromInt(at(q)), @ptrFromInt(at(k)), @ptrFromInt(at(v)), @ptrFromInt(at(out)), 1, @intCast(qlen), @intCast(span), heads, kv_heads, d, scale, @intCast(q_pos0), 0, @intCast(span * d), d, 0, @intCast(span * d), d, kind, null, null, null, s, null); + try t.stream.synchronize(); + r.flat[i] = try rig.download(f32, out); + } + return r; +} + +/// Plants past rows' edges, odd start, partial row block, span past the last row: float64's class, tile's own bits. +fn edges(t: *Rig, rng: *rig.Rng) !void { + const q_pos0 = 87; + const qlen = 150; + const span = 256; // the last tile, keys 192 to 255, holds keys past the last row's 236 + const inf = std.math.inf(f32); + const plants = [_]Plant{ + .{ .kv = 0, .key = q_pos0 + 1, .dim = 7, .value = inf }, // past row 0's edge + .{ .kv = 1, .key = q_pos0 + 71, .dim = 9, .value = -inf }, + .{ .kv = 2, .key = q_pos0 + 141, .dim = 11, .value = std.math.nan(f32) }, + .{ .kv = 3, .key = q_pos0 + 30, .dim = 13, .value = inf }, // row 30 sees +inf alone, later rows both + .{ .kv = 3, .key = q_pos0 + 31, .dim = 13, .value = -inf }, + .{ .kv = 4, .key = 250, .dim = 15, .value = inf }, // past every row, inside the last tile + }; + const c = try Cache.make(t, rng, kv_heads, span, d); + defer c.free(); + plant(t, c, &plants, span); + const hq = try gpa.alloc(f32, heads * qlen * d); + defer gpa.free(hq); + for (hq) |*x| x.* = rng.unit(); + const r = try run(t, c, hq, qlen, span, q_pos0); + defer for ([_][2][]f32{ r.paged, r.flat }) |pair| for (pair) |x| gpa.free(x); + try std.testing.expectEqualSlices(u32, @ptrCast(r.flat[0]), @ptrCast(r.paged[0])); + for (0..heads) |h| for (0..qlen) |i| for (0..d) |e| { + const o = (h * qlen + i) * d + e; + const want = expected(&plants, h / group, e, q_pos0 + i); + if (classOf(r.paged[0][o]) != want) { + std.debug.print("head {d} row {d} dim {d}: {d} where float64 gives {t}\n", .{ h, i, e, r.paged[0][o], want }); + return error.TestUnexpectedResult; + } + if (planted(&plants, e)) continue; + // the same finite sums with and without the pass: a past-edge lane enters as +0 in one, as 0 x v in the other + for ([_][2][]f32{ r.paged, r.flat }) |pair| { + if (@as(u32, @bitCast(pair[0][o])) != @as(u32, @bitCast(pair[1][o]))) { + std.debug.print("head {d} row {d} dim {d}: {d} served, {d} tile alone\n", .{ h, i, e, pair[0][o], pair[1][o] }); + return error.TestUnexpectedResult; + } + } + }; +} + +/// A cache of `cap` positions on the device, flat (kv_heads, cap, d) and as a pool over its page table. +const Device = struct { + k: DeviceBuffer, + v: DeviceBuffer, + pool_k: DeviceBuffer, + pool_v: DeviceBuffer, + table: DeviceBuffer, + count: usize, + + fn upload(t: *Rig, c: Cache, cap: usize) !Device { + const pk = try c.pool(c.k, kv_heads, cap, d); + defer gpa.free(pk); + const pv = try c.pool(c.v, kv_heads, cap, d); + defer gpa.free(pv); + return .{ .k = try t.upload(c.k), .v = try t.upload(c.v), .pool_k = try t.upload(pk), .pool_v = try t.upload(pv), .table = try t.upload(c.table), .count = c.table.len }; + } + + fn free(x: *Device) void { + for ([_]*DeviceBuffer{ &x.k, &x.v, &x.pool_k, &x.pool_v, &x.table }) |b| b.free(); + } +}; + +/// Rows [0, cut) from 0, then [cut, len) from cut, one launch each, paged or flat; out (heads, len, d). +fn cutRuns(t: *Rig, hq: []const f32, dev: Device, cap: usize, len: usize, cut: usize, paged: bool) ![]f32 { + const out = try gpa.alloc(f32, heads * len * d); + errdefer gpa.free(out); + const scale: f32 = 1 / @sqrt(@as(f32, @floatFromInt(d))); + const kind: c_int = if (t.bf16) 1 else 0; + const s = t.stream.handle; + for ([_][2]usize{ .{ 0, cut }, .{ cut, len } }) |seg| { + const p = seg[0]; + const rows = seg[1] - p; + if (rows == 0) continue; + const q = try gpa.alloc(f32, heads * rows * d); + defer gpa.free(q); + for (0..heads) |h| @memcpy(q[h * rows * d ..][0 .. rows * d], hq[(h * len + p) * d ..][0 .. rows * d]); + var dq = try t.upload(q); + defer dq.free(); + var dout = try t.alloc(q.len * 4); + defer dout.free(); + if (paged) { + try t.on.pagedCausal(at(dq), at(dev.pool_k), at(dev.pool_v), at(dev.table), at(dout), rows, p + rows, heads, kv_heads, d, scale, p, dev.count, kind, s); + } else { + try t.on.tf_causal(@ptrFromInt(at(dq)), @ptrFromInt(at(dev.k)), @ptrFromInt(at(dev.v)), @ptrFromInt(at(dout)), 1, @intCast(rows), @intCast(p + rows), heads, kv_heads, d, scale, @intCast(p), 0, @intCast(cap * d), d, 0, @intCast(cap * d), d, kind, null, null, null, s, null); + } + try t.stream.synchronize(); + const got = try rig.download(f32, dout); + defer gpa.free(got); + for (0..heads) |h| @memcpy(out[(h * len + p) * d ..][0 .. rows * d], got[h * rows * d ..][0 .. rows * d]); + } + return out; +} + +/// Planted values past some rows' edges: each row has the same bits cut at 0 and at `cut`, paged and flat. +fn shiftedPlanted(t: *Rig, rng: *rig.Rng, len: usize, cut: usize) !void { + const cap = 256; // whole pages past len + const inf = std.math.inf(f32); + const plants = [_]Plant{ + .{ .kv = 0, .key = cut + 1, .dim = 7, .value = inf }, + .{ .kv = 1, .key = cut + 2, .dim = 9, .value = -inf }, + .{ .kv = 2, .key = len - 3, .dim = 11, .value = std.math.nan(f32) }, + .{ .kv = 3, .key = 40, .dim = 13, .value = inf }, + .{ .kv = 3, .key = 41, .dim = 13, .value = -inf }, + }; + const c = try Cache.make(t, rng, kv_heads, cap, d); + defer c.free(); + plant(t, c, &plants, cap); + const hq = try gpa.alloc(f32, heads * len * d); + defer gpa.free(hq); + for (hq) |*x| x.* = rng.unit(); + var dev = try Device.upload(t, c, cap); + defer dev.free(); + for ([_]bool{ true, false }) |paged| { + const alone = try cutRuns(t, hq, dev, cap, len, 0, paged); + defer gpa.free(alone); + const after = try cutRuns(t, hq, dev, cap, len, cut, paged); + defer gpa.free(after); + try std.testing.expectEqualSlices(u32, @ptrCast(alone), @ptrCast(after)); + } +} + +test "a non-finite value past a row's causal edge adds nothing to that row, in both cache types, paged and flat" { + const t = try Rig.open(1 << 20); + defer t.close(); + var rng: rig.Rng = .{ .state = 0xED6E5 }; + const act = t.bf16; + defer t.bf16 = act; + for ([_]bool{ false, true }) |bf16| { + if (bf16 and !act) continue; // RDNA2 has no bf16 dot2 and serves fp16 caches only + t.bf16 = bf16; + try edges(t, &rng); + try shiftedPlanted(t, &rng, 237, 1); + try shiftedPlanted(t, &rng, 237, 87); + } +} diff --git a/zig/src/hip/launches.zig b/zig/src/hip/launches.zig index 4b620820e1..e6b5c647cd 100644 --- a/zig/src/hip/launches.zig +++ b/zig/src/hip/launches.zig @@ -80,6 +80,8 @@ pub const Launcher = struct { fa_prefill: Triple, fa_wide: [2]Function, // the 64-row prefill tile, fp16 and bf16 caches fa_paged: [2]Function, // the same tile over a stream's pages + fa_wide_edge: [2]Function, // after each, where dot2 ignores a zero's sign: a block whose masked keys are non-finite + fa_paged_edge: [2]Function, softmax_stats: Function, sum_partials: Function, op: Ops, @@ -91,6 +93,8 @@ pub const Launcher = struct { fuse: bool, /// The 64-row prefill attention tile is on. wide: bool, + /// The edge pass runs after the prompt tile (`edge_pass`; tests turn it off to see the tile alone). + edge: 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). @@ -164,6 +168,9 @@ pub const Launcher = struct { token_prob: Function, }; + /// The prompt tile keeps main's sums and needs the edge pass after it: dot2 here ignores a zero's sign. + pub const edge_pass = if (caps.Caps.of(kernels.arch)) |c| !c.dot2_signed_zero else true; + /// Which kernels a run takes where two give the same bits; every one defaults on. pub const Choices = struct { fuse: bool = true, wide: bool = true, products: affine.Choices = .{} }; @@ -186,6 +193,7 @@ pub const Launcher = struct { const pl = l.mods[@backingInt(kernels.Group.plan)]; l.fuse = choices.fuse; l.wide = choices.wide; + l.edge = edge_pass; l.dec = .{ .router = .{ try dec.function("tf_router_decode_f16"), try dec.function("tf_router_decode_bf16") }, .router_bf16 = .{ try dec.function("tf_router_bf16_f16"), try dec.function("tf_router_bf16") }, .tail = try dec.function("tf_tail"), .conv_split = try dec.function("tf_conv_split"), .rms2 = try dec.function("tf_rms2"), .gnorm_out = try dec.function("tf_gnorm_out"), .select = try dec.function("tf_select_decode") }; l.plan = .{ .kv_write = try pl.function("tf_plan_kv_write"), @@ -206,6 +214,8 @@ pub const Launcher = struct { }; l.fa_wide = .{ try pre.function("tf_fa_wide_f16"), try pre.function("tf_fa_wide_bf16") }; l.fa_paged = .{ try pre.function("tf_fa_paged_f16"), try pre.function("tf_fa_paged_bf16") }; + l.fa_wide_edge = .{ try pre.function("tf_fa_wide_edge_f16"), try pre.function("tf_fa_wide_edge_bf16") }; + l.fa_paged_edge = .{ try pre.function("tf_fa_paged_edge_f16"), try pre.function("tf_fa_paged_edge_bf16") }; l.columns = try columns.Functions.load(ops); l.vision = try vision.Functions.load(ops, pre); l.dense_text = try dense_text.Functions.load(ops); diff --git a/zig/src/hip/launches/attention.zig b/zig/src/hip/launches/attention.zig index 00e6ea71ba..9dcaaa4207 100644 --- a/zig/src/hip/launches/attention.zig +++ b/zig/src/hip/launches/attention.zig @@ -94,6 +94,7 @@ pub fn tf_causal(l: *const Launcher, q: CF, k: C, v: C, out: F, batch: c_int, ql // a 16-bit cache of whole 64-wide heads takes the 64-row tile (attention=f32: 16-row); odd heads a wave a query if (kind < 2 and @rem(d, 64) == 0 and d <= 256 and @rem(k_ss, 8) == 0 and @rem(v_ss, 8) == 0 and l.wide) { try l.go(l.fa_wide[kind], dim(cdiv(qlen, 64), heads, batch), dim(256, 1, 1), 0, s, &a); + if (l.edge) try l.go(l.fa_wide_edge[kind], dim(cdiv(qlen, 64), heads, batch), dim(256, 1, 1), 0, s, &a); } else if (d >= 2 and @rem(d, 2) == 0) { try l.go(l.fa_prefill[kind], dim(cdiv(qlen, 16), heads, batch), dim(256, 1, 1), 0, s, &a); } else { @@ -119,6 +120,8 @@ pub fn pagedCausal(l: *const Launcher, q: u64, k: u64, v: u64, table: u64, out: try a.add(table); try a.add(@as(c_int, @intCast(count))); try l.go(l.fa_paged[@intCast(cache_kind)], dim(cdiv(qlen, 64), heads, 1), dim(256, 1, 1), 0, s, &a); + // where dot2 ignores a zero's sign the tile is main's; the edge pass reruns a block a non-finite masked key reaches + if (l.edge) try l.go(l.fa_paged_edge[@intCast(cache_kind)], dim(cdiv(qlen, 64), heads, 1), dim(256, 1, 1), 0, s, &a); } pub fn tf_attn_gate(l: *const Launcher, att: CF, qg: C, out: P, kind: c_int, len: c_int, heads: c_int, d: c_int, rows_major: c_int, s: S) Error!void { From 6add094b85680b2aa8cb952ebf47c035f7610605 Mon Sep 17 00:00:00 2001 From: laudney <106273+laudney@users.noreply.github.com> Date: Fri, 6 Nov 2026 16:10:59 +0000 Subject: [PATCH 04/36] test(hip): contract entries for non-finite values past a row's edge The prompt tile's gate-1 plant mixed an infinity past a row's causal edge with one under fp32's least weight, which still fails, so the contract stayed green with no mask or edge pass at all (#2357 review). Four entries, flat and paged in both cache types, plant +inf, -inf and NaN only past rows' edges over a span a page past the last row and expect a clean gate 1. Each has a planted control, main's tile without the edge pass, which must go red at that gate. --- zig/kernels/hip/contract_controls.hip | 33 +++++++++++ zig/tests/hip/contract_main.zig | 18 ++++++ zig/tests/hip/contract_prefill.zig | 79 ++++++++++++++++++++++++--- 3 files changed, 123 insertions(+), 7 deletions(-) diff --git a/zig/kernels/hip/contract_controls.hip b/zig/kernels/hip/contract_controls.hip index f9521ce5b1..1ef9b29107 100644 --- a/zig/kernels/hip/contract_controls.hip +++ b/zig/kernels/hip/contract_controls.hip @@ -5,6 +5,7 @@ #include #include +#include "attention/prefill.hpp" #include "quant/dense_tiles.hpp" __device__ __forceinline__ float tf_ctl_x(const uint16_t* x, int64_t i) { @@ -161,3 +162,35 @@ extern "C" __global__ void tf_ctl_js_combine_unrounded(const float* prior, const if (j >= n) return; logits[j] = tf_ctl_bf16(prior[j] + gate * (joint_scale * cosine[j] + tf_ctl_x(res, j))); } + +// Gate 1: the prompt tile with neither mask nor edge pass, so 0 x inf at a key past a row's edge gives NaN in that row. +extern "C" __global__ void __launch_bounds__(256, 1) +tf_ctl_fa_wide_unmasked_f16(const float* q, const __half* k, const __half* v, float* out, int qlen, int span, int heads, + int kv_heads, int d, float scale, int q_pos0, long long k_sb, long long k_sh, long long k_ss, + long long v_sb, long long v_sh, long long v_ss) { + fa_dot2<__half, false, true, false, false>(q, k, v, out, qlen, span, heads, kv_heads, d, scale, q_pos0, k_sb, k_sh, + k_ss, v_sb, v_sh, v_ss, nullptr, 0); +} + +extern "C" __global__ void __launch_bounds__(256, 1) +tf_ctl_fa_wide_unmasked_bf16(const float* q, const hip_bfloat16* k, const hip_bfloat16* v, float* out, int qlen, + int span, int heads, int kv_heads, int d, float scale, int q_pos0, long long k_sb, + long long k_sh, long long k_ss, long long v_sb, long long v_sh, long long v_ss) { + fa_dot2(q, k, v, out, qlen, span, heads, kv_heads, d, scale, q_pos0, k_sb, + k_sh, k_ss, v_sb, v_sh, v_ss, nullptr, 0); +} + +extern "C" __global__ void __launch_bounds__(256, 1) +tf_ctl_fa_paged_unmasked_f16(const float* q, const __half* k, const __half* v, float* out, int qlen, int span, + int heads, int kv_heads, int d, float scale, int q_pos0, const unsigned* table, int count) { + fa_dot2<__half, true, true, false, false>(q, k, v, out, qlen, span, heads, kv_heads, d, scale, q_pos0, 0, 0, 0, 0, + 0, 0, table, count); +} + +extern "C" __global__ void __launch_bounds__(256, 1) +tf_ctl_fa_paged_unmasked_bf16(const float* q, const hip_bfloat16* k, const hip_bfloat16* v, float* out, int qlen, + int span, int heads, int kv_heads, int d, float scale, int q_pos0, const unsigned* table, + int count) { + fa_dot2(q, k, v, out, qlen, span, heads, kv_heads, d, scale, q_pos0, 0, 0, + 0, 0, 0, 0, table, count); +} diff --git a/zig/tests/hip/contract_main.zig b/zig/tests/hip/contract_main.zig index dd39d3135d..8697f33c4d 100644 --- a/zig/tests/hip/contract_main.zig +++ b/zig/tests/hip/contract_main.zig @@ -188,6 +188,24 @@ fn run(a: std.mem.Allocator, io: std.Io) !usize { for (&prefills, [_]bool{ true, false }) |*p, paged| p.* = .init(a, &r, s, &l, paged); defer for (&prefills) |*p| p.deinit(); for (&prefills) |*p| try kernels.append(a, p.kernel()); + // non-finite values only past rows' edges, flat and paged, both cache types; then each without its edge pass + var edges: [8]prefill_k.Prefill = undefined; + var edge_count: usize = 0; + defer for (edges[0..edge_count]) |*p| p.deinit(); + for (0..8) |i| { + const paged = i % 2 == 1; + const kind: kit.Kind = if (i / 2 % 2 == 1) .bf16 else .f16; + if (kind == .bf16 and kit.act() != .bf16) continue; // RDNA2 serves fp16 caches only + const p = &edges[edge_count]; + p.* = .init(a, &r, s, &l, paged); + edge_count += 1; + p.k = kind; + p.mode = .edge; + if (i >= 4) p.planted = try ctl.function(([_][:0]const u8{ + "tf_ctl_fa_wide_unmasked_f16", "tf_ctl_fa_paged_unmasked_f16", "tf_ctl_fa_wide_unmasked_bf16", "tf_ctl_fa_paged_unmasked_bf16", + })[i - 4]); + try kernels.append(a, p.kernel()); + } var add_rms: [3]add_rms_k.AddRms = undefined; for (&add_rms, std.enums.values(add_rms_k.Mode)) |*m, w| m.* = .init(a, &r, s, &l, w); defer for (&add_rms) |*m| m.deinit(); diff --git a/zig/tests/hip/contract_prefill.zig b/zig/tests/hip/contract_prefill.zig index cbf65ce444..db92bcf393 100644 --- a/zig/tests/hip/contract_prefill.zig +++ b/zig/tests/hip/contract_prefill.zig @@ -39,6 +39,27 @@ const names = [2][2][]const u8{ .{ "hip prompt tile tf_fa_paged_f16 (fp16 cache pages)", "hip prompt tile tf_fa_paged_bf16 (bf16 cache pages)" }, }; +/// The edge entries: non-finite values only at keys past some rows' edges; those rows must ignore them, later rows not. +const edge_names = [2][2][]const u8{ + .{ + "hip prompt tile tf_fa_wide_f16 (flat fp16 cache), non-finite values only past a row's edge", + "hip prompt tile tf_fa_wide_bf16 (flat bf16 cache), non-finite values only past a row's edge", + }, + .{ + "hip prompt tile tf_fa_paged_f16 (fp16 cache pages), non-finite values only past a row's edge", + "hip prompt tile tf_fa_paged_bf16 (bf16 cache pages), non-finite values only past a row's edge", + }, +}; + +/// Their planted controls: the tile with no edge pass and no mask, so 0 x inf past a row's edge gives NaN there. +const control_names = [2][2][]const u8{ + .{ "control: prompt tile tf_fa_wide_f16 without its edge pass", "control: prompt tile tf_fa_wide_bf16 without its edge pass" }, + .{ "control: prompt tile tf_fa_paged_f16 without its edge pass", "control: prompt tile tf_fa_paged_bf16 without its edge pass" }, +}; + +/// The served entry, or an edge entry: its own plant and a span a page past its last row. +pub const Mode = enum { served, edge }; + /// What gate 1 finds on the tile: NaN scores leave its max, and its P x V takes a seen key whose weight underflows. const finding = "a row whose every score is NaN writes 0 where float64 gives NaN (its running max is fmaxf, " ++ "which drops NaN, and a row at max -inf takes p = 0), and an infinite value at a key the row sees under " ++ @@ -50,19 +71,21 @@ fn keyChain(s: c.Shape) usize { } /// Rows of a prompt chunk (rows), query heads (n) and the chunk's first position (k); the flat tile equals the paged. -pub fn spec(k: kit.Kind, paged: bool) c.Spec { +pub fn spec(k: kit.Kind, paged: bool, mode: Mode, planted: bool) c.Spec { const bf16 = k == .bf16; + const p = @intFromBool(paged); + const b = @intFromBool(bf16); return .{ - .name = names[@intFromBool(paged)][@intFromBool(bf16)], + .name = if (planted) control_names[p][b] else if (mode == .edge) edge_names[p][b] else names[p][b], .accum = .f32, .depth = keyChain, .ops = 1, .out = .f32, - .replaces = if (paged) null else if (bf16) "tf_fa_paged_bf16" else "tf_fa_paged_f16", + .replaces = if (paged or planted) null else if (bf16) "tf_fa_paged_bf16" else "tf_fa_paged_f16", .tiles = .{ .{ .size = 64 }, .{}, .{ .size = page, .min = 0 } }, .served = &served, .sums = .{ .why = attn.why }, - .known = .{ .gate = .nonfinite, .finding = finding }, + .known = if (mode == .served) .{ .gate = .nonfinite, .finding = finding } else null, }; } @@ -73,6 +96,9 @@ pub const Prefill = struct { k: kit.Kind, /// tf_fa_paged_* over a permuted page table; else tf_fa_wide_* over flat keys, the paged tile its baseline. paged: bool, + mode: Mode = .served, + /// A planted control (tf_ctl_fa_*_unmasked_*) launched in place of the served dispatch. + planted: ?hip.Function = null, arena: std.heap.ArenaAllocator, r: ?attn.Ref = null, shape: c.Shape = undefined, @@ -105,13 +131,15 @@ pub const Prefill = struct { } pub fn kernel(self: *Prefill) g.Kernel { - return .{ .spec = spec(self.k, self.paged), .ctx = self, .vt = &vt }; + const planted = self.planted != null; + const control: ?g.Gate = if (planted) .nonfinite else null; + return .{ .spec = spec(self.k, self.paged, self.mode, planted), .ctx = self, .vt = &vt, .control = control }; } const vt = g.hooks(Prefill); fn span(self: *const Prefill) usize { - return self.shape.k + self.shape.rows; + return self.shape.k + self.shape.rows + @as(usize, if (self.mode == .edge) page else 0); } /// Element e of kv head kh's key `key` in kc and vc. @@ -151,7 +179,7 @@ pub const Prefill = struct { self.r = try attn.Ref.init(self.rig.a, d, keys); switch (fill) { .random => {}, - .nonfinite => self.plant(), + .nonfinite => if (self.mode == .edge) self.edgePlant() else self.plant(), .stress => try self.stress(rng), } self.q32 = try self.rig.input(try kit.f32Bytes(a, self.q)); @@ -204,6 +232,20 @@ pub const Prefill = struct { } } + /// Non-finite values only past rows' edges: +inf, -inf, NaN, one column seeing both infinities, one past every row. + fn edgePlant(self: *Prefill) void { + const n = self.shape.rows; + if (n < 8) return; + const hk = self.geo.kv; + const inf = std.math.inf(f64); + self.vc[self.at(0, self.pos(0) + 1, 7)] = inf; + self.vc[self.at(1 % hk, self.pos(2) + 2, 9)] = -inf; + self.vc[self.at(0, self.pos(5) + 1, 11)] = std.math.nan(f64); + self.vc[self.at(1 % hk, self.pos(0) + 1, 13)] = inf; + self.vc[self.at(1 % hk, self.pos(0) + 2, 13)] = -inf; + self.vc[self.at(0, self.pos(n - 1) + 1, 15)] = inf; + } + /// Query row r of kv head kh's first head. fn query(self: *Prefill, r: usize, kh: usize) []f64 { const head = kh * (self.geo.heads / self.geo.kv); @@ -281,6 +323,7 @@ pub const Prefill = struct { } pub fn run(self: *Prefill) !void { + if (self.planted) |f| return self.runPlanted(f); if (self.paged) { try self.launchPaged(self.out); return self.rig.s.synchronize(); @@ -319,6 +362,28 @@ pub const Prefill = struct { try self.rig.s.synchronize(); } + /// The planted control on the served symbol's arguments and grid, with no edge pass after it. + fn runPlanted(self: *Prefill, f: hip.Function) !void { + const s = self.shape; + const geo = self.geo; + var a: hip.launch.Args = .{}; + for ([_]u64{ self.q32, if (self.paged) self.pool[0] else self.flat[0], if (self.paged) self.pool[1] else self.flat[1], self.out }) |p| try a.add(p); + for ([_]usize{ s.rows, self.span(), geo.heads, geo.kv, geo.d }) |x| try a.add(@as(c_int, @intCast(x))); + try a.add(self.scale()); + try a.add(@as(c_int, @intCast(s.k))); + if (self.paged) { + try a.add(self.table); + try a.add(@as(c_int, @intCast(self.pages))); + } else { + const ss: c_longlong = @intCast(geo.d); + const sh: c_longlong = @intCast(self.span() * geo.d); + for ([_]c_longlong{ 0, sh, ss, 0, sh, ss }) |x| try a.add(x); + } + const cfg: hip.launch.Config = .{ .grid = .{ .x = @intCast((s.rows + 63) / 64), .y = @intCast(geo.heads) }, .block = .{ .x = 256 } }; + try hip.launch.launch(f, cfg, self.rig.s, &a); + try self.rig.s.synchronize(); + } + /// The paged tile on the same keys, laid out in pages. pub fn replaced(self: *Prefill) ![]const u8 { try self.launchPaged(self.base); From 4bea033dd5df1ab59841e6fdbe09093274de5eca Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Sat, 7 Nov 2026 04:36:49 +0000 Subject: [PATCH 05/36] changelog: #2357 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index a18c171ad9..609a70b1fa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- HIP: on RDNA 3.5 (gfx1150 to gfx1153) a prompt row gets the same bytes wherever its chunk starts. The 64-row prompt tile there now takes each key past a row's causal edge as +0, as RDNA 3.5's dot2 rounds a -0 lane apart from +0, so its prompt bits change, and gfx1151's prefill is up to 0.7% slower at 2k and 8k. Other chips keep main's prompt tile and bits, and gfx1201's prompt speed is main's within noise. On every chip, a non-finite value at a key past a row's edge now adds nothing to that row (#2357). Thanks to @laudney. + - GLM CUDA at TP=2 and up: each rank names its startup phases (NCCL init, peer warm-up, weights, calibration, ready), prints NCCL's init device, version and log path, and warns once a minute while NCCL init or the warm-up is unfinished, with the transport seen so far. It only reports; transfers are unchanged (#2354, part of #2347). - Docs: the serve flags match the CLI: `--metrics-open` opens `/v1/metrics` too, `--learn-conversations` on CUDA serves the Qwen3.5 hybrids and Flash Next as well as Nemotron, Flash Next shares sessions under an explicit `--parallel` above one and refuses `--recall` only on a Mac, and the Metal Qwen3.5-2B keeps prompt states (#2367). From ed7542f0ef73f0d0d54969391015896479d4c72b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Fri, 6 Nov 2026 15:36:29 +0100 Subject: [PATCH 06/36] qwen27 metal: a keep that fails part way ends its encoder, so the error returns instead of aborting the process Runner.keep opens its command encoder before the per-layer commits and keep uploads; any error after that left the encoder open, and the objc pool's release then asserted "Command encoder released without endEncoding", so the process aborted and the error was lost. A guarded errdefer ends it, as qwen27_keep's other keep sites do, so the errors after the encoder's own end (a GPU failure, a refused path) don't end it twice. tf-qwen27-keep-error fails each of the keep's allocations in turn on a synthetic model: on main it aborts, here each returns OutOfMemory and the same round keeps on retry; then a refused path after the end returns BadPath. Found on the pipeline stage path (#2005). --- CHANGELOG.md | 1 + zig/build/qwen27.zig | 1 + zig/src/families/qwen27/decode_round.zig | 3 ++ zig/tests/qwen27/keep_error.zig | 44 ++++++++++++++++++++++++ 4 files changed, 49 insertions(+) create mode 100644 zig/tests/qwen27/keep_error.zig diff --git a/CHANGELOG.md b/CHANGELOG.md index 609a70b1fa..6dd320fd02 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -81,6 +81,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Internal: 42 more lean findings fixed in the tests' comments and the tools' docstrings, with no change in behaviour (#2315). - HIP: the decode round's residual-add and norm tail issues each row's loads before its branches, and a MoE tail loads its slots' rows together, with output bytes unchanged. On a gfx1151, one-stream decode is 2.7% faster on Qwen3.6-35B-A3B and 1.1% on the 27B. `tf_tail` is now under the kernel contract (#2247). Thanks to @AbrahamCete. +- Qwen3.8-27B on Metal: a round's keep that fails part way (an allocation, a refused keep or path) returns its error instead of aborting the process on an unended command encoder, and after a failed allocation the round can keep again; `tf-qwen27-keep-error` fails each of the keep's allocations in turn and refuses a path after the encoder has ended (#2376). Thanks to @jschmied. - Docs: `docs/optimization-ledger.md` lists which speed and memory optimizations apply to which native family, and where each stands, with links to the PRs that measured them (#2232). Thanks to @jschmied. diff --git a/zig/build/qwen27.zig b/zig/build/qwen27.zig index 322c96d366..524cf517e6 100644 --- a/zig/build/qwen27.zig +++ b/zig/build/qwen27.zig @@ -113,6 +113,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu _ = program(b, "tf-qwen27-dflash-run", "Build the explicit draft prototype and its plain comparison on real weights", "zig/tests/qwen27/dflash_run.zig", target, true, family); const engine_only: []const std.Build.Module.Import = &.{ .{ .name = "tensorfold", .module = engine }, .{ .name = "metal", .module = metal } }; _ = program(b, "tf-qwen27-empty-keep", "Build shared empty/full keep state checks", "zig/tests/qwen27/empty_keep.zig", target, false, engine_only); + _ = program(b, "tf-qwen27-keep-error", "Build the check that a failing keep returns its error, never aborting", "zig/tests/qwen27/keep_error.zig", target, false, engine_only); _ = program(b, "tf-qwen27-synthetic", "Build bounded synthetic native hybrid chunk checks, no models", "zig/tests/qwen27/synthetic.zig", target, false, engine_only); _ = program(b, "tf-qwen27-dflash-synthetic", "Build model-free native draft context/forward/tree check", "zig/tests/qwen27/dflash_runtime.zig", target, false, engine_only); _ = program(b, "tf-qwen27-widths", "Build the projections at every oQ width before tensor units: float64, row windows, mixed stacks, embedding", "zig/tests/qwen27/widths.zig", target, false, engine_only); diff --git a/zig/src/families/qwen27/decode_round.zig b/zig/src/families/qwen27/decode_round.zig index 155a3b1ee3..283b2f40f7 100644 --- a/zig/src/families/qwen27/decode_round.zig +++ b/zig/src/families/qwen27/decode_round.zig @@ -297,6 +297,8 @@ pub const Runner = struct { defer pool.pop(); const cb = r.model.queue.commandBuffer(); const e = cb.compute(.serial); + var ended = false; + errdefer if (!ended) e.end(); // an encoder the pool releases unended aborts the process, and the error is lost var gi: usize = 0; var ai: usize = 0; for (r.model.weights.layers) |layer| switch (layer.mixer) { @@ -313,6 +315,7 @@ pub const Runner = struct { }, }; e.end(); + ended = true; cb.commit(); cb.wait(); if (gpuFailed(cb, "a round's keep")) { diff --git a/zig/tests/qwen27/keep_error.zig b/zig/tests/qwen27/keep_error.zig new file mode 100644 index 0000000000..5e18657781 --- /dev/null +++ b/zig/tests/qwen27/keep_error.zig @@ -0,0 +1,44 @@ +//! A keep that fails at any of its allocations, or after its encoder ends, returns the error instead of aborting. +const std = @import("std"); +const mtl = @import("metal"); +const q = @import("tensorfold").qwen27; +const graph = @import("synthetic_graph.zig"); + +pub fn main(init: std.process.Init) !void { + const a = init.gpa; + const autorelease = mtl.objc.Pool.push(); + defer autorelease.pop(); + const model = try graph.createLayers(a, 5); + defer model.deinit(); + var failing = std.testing.FailingAllocator.init(a, .{}); + var r = try q.decode_round.Runner.init(failing.allocator(), model, 1, 256); + defer r.deinit(); + const s = q.session.Session{ .runner = &r }; + try r.reset(0); + try s.prefill(&.{ 1, 2, 3, 4, 5 }, 16); + var failed: usize = 0; + var at: usize = 0; + // each of the keep's allocations fails in turn: the encoder it opened must end before the error returns + while (true) : (at += 1) { + var round = try q.round_plan.Round.init(a, &.{.{ .slot = 0, .start = r.offsets[0], .capacity = r.capacity, .ids = &.{ 6, 7 } }}, @intCast(model.config.conv_kernel), 1, @intCast(model.config.vocab)); + defer round.deinit(); + try r.verifyHead(&round, .all); + failing.fail_index = failing.alloc_index + at; + const result = r.keep(&round, &.{&.{0}}); + failing.fail_index = std.math.maxInt(usize); + result catch |err| { + if (err != error.OutOfMemory) return err; + failed += 1; + try r.keep(&round, &.{&.{0}}); // the same round keeps once the allocation succeeds + continue; + }; + break; + } + if (failed == 0) return error.NoKeepAllocation; + // a path the window refuses fails after the encoder has ended: the error returns, with no second end + var round = try q.round_plan.Round.init(a, &.{.{ .slot = 0, .start = r.offsets[0], .capacity = r.capacity, .ids = &.{ 8, 9 } }}, @intCast(model.config.conv_kernel), 1, @intCast(model.config.vocab)); + defer round.deinit(); + try r.verifyHead(&round, .all); + if (r.keep(&round, &.{&.{1}})) |_| return error.BadPathKept else |err| if (err != error.BadPath) return err; + std.debug.print("PASS a keep failing at each of its {d} allocations returns OutOfMemory and keeps on retry; a refused path after the encoder ends returns BadPath\n", .{failed}); +} From 47bcd1504eae1139f5a5e94df9da09138a45a3af Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Sat, 7 Nov 2026 06:50:55 +0900 Subject: [PATCH 07/36] glm cuda: test-glm-snapshots-host's copies module imports core, which cuda_snaps.zig reads for TF_GLM_PROMPT_TAIL --- zig/build/glm_family_host.zig | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zig/build/glm_family_host.zig b/zig/build/glm_family_host.zig index 9a51362712..95a5a690f3 100644 --- a/zig/build/glm_family_host.zig +++ b/zig/build/glm_family_host.zig @@ -48,7 +48,7 @@ pub fn steps(b: *std.Build, cuda: *std.Build.Module, glm: *std.Build.Module, mod const admission_run = b.addRunArtifact(b.addTest(.{ .root_module = admission, .filters = &.{"GLM CUDA admission:"} })); b.step("test-glm-admission-host", "Real GLM admission capabilities through the lane worker without GPU work").dependOn(&admission_run.step); step.dependOn(&admission_run.step); - const copies = b.createModule(.{ .root_source_file = b.path("zig/src/native/cuda_snaps.zig"), .target = host, .optimize = .debug, .imports = &.{.{ .name = "engine_api", .module = api }} }); + const copies = b.createModule(.{ .root_source_file = b.path("zig/src/native/cuda_snaps.zig"), .target = host, .optimize = .debug, .imports = &.{ .{ .name = "engine_api", .module = api }, .{ .name = "core", .module = mods.core } } }); const snapshots = b.createModule(.{ .root_source_file = b.path("zig/glm_snapshot_host.zig"), .target = host, .optimize = .debug, .link_libc = true, .imports = &.{ .{ .name = "glm", .module = glm }, .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core }, .{ .name = "engine_api", .module = api }, .{ .name = "copies", .module = copies } } }); b.step("test-glm-snapshots-host", "Prompt-sized GLM state copies, page budgets and rank ordering without devices").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = snapshots })).step); b.step("test-glm-prompt-settings-host", "Actual opt-in settings and rank agreement on the CPU").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = glm, .filters = &.{"prompt K experiment"} })).step); From b0799371dddff3f83a2d7fdeac87397a59d95627 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Sat, 7 Nov 2026 04:36:52 +0000 Subject: [PATCH 08/36] changelog: #2403 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6dd320fd02..68295dee03 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Tests: `zig build test-glm-snapshots-host` builds again: its copies module imports `core`, which `cuda_snaps.zig` reads for `TF_GLM_PROMPT_TAIL` since #2226 (#2403). Thanks to @Bizuayeu. + - HIP: on RDNA 3.5 (gfx1150 to gfx1153) a prompt row gets the same bytes wherever its chunk starts. The 64-row prompt tile there now takes each key past a row's causal edge as +0, as RDNA 3.5's dot2 rounds a -0 lane apart from +0, so its prompt bits change, and gfx1151's prefill is up to 0.7% slower at 2k and 8k. Other chips keep main's prompt tile and bits, and gfx1201's prompt speed is main's within noise. On every chip, a non-finite value at a key past a row's edge now adds nothing to that row (#2357). Thanks to @laudney. - GLM CUDA at TP=2 and up: each rank names its startup phases (NCCL init, peer warm-up, weights, calibration, ready), prints NCCL's init device, version and log path, and warns once a minute while NCCL init or the warm-up is unfinished, with the transport seen so far. It only reports; transfers are unchanged (#2354, part of #2347). From c4dc04b999fa786e9697ac934646e32047e3cb31 Mon Sep 17 00:00:00 2001 From: Armin <2369138+Arminova@users.noreply.github.com> Date: Mon, 2 Nov 2026 22:44:14 +0000 Subject: [PATCH 09/36] flashnext cuda: the grouped draft head's lanes carry their page table On TF_FLASHNEXT_KV_PAGES, a shared round with two or more drafting streams ran the MTP head over one grouped pass whose lanes had no page table, so the head wrote and read its rows as if its cache were contiguous from the first page: into other streams' pages. With three equal 16K prompts at once, one reply of three differed from its solo run. The lone head's pass already carries it. The grouped lane is built in headLane now, and a host test checks that it carries its draft's page table, length and rows (without `.pg` it fails: expected 28672, found 0). --- zig/src/families/flashnext/cuda_mtp_ops.zig | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/zig/src/families/flashnext/cuda_mtp_ops.zig b/zig/src/families/flashnext/cuda_mtp_ops.zig index 5cdf3943e0..716f49ea45 100644 --- a/zig/src/families/flashnext/cuda_mtp_ops.zig +++ b/zig/src/families/flashnext/cuda_mtp_ops.zig @@ -55,6 +55,11 @@ fn cint(v: usize) c_int { return @intCast(v); } +/// A lane of the head's grouped pass: its draft's cache and pages at its length, as the lone head's pass reads them. +fn headLane(d: *Draft, rows: usize, score_rows: u64) batch.Lane { + return .{ .seq = .{ .media = null, .conv = &.{}, .state = &.{}, .kc = (&d.kc)[0..1], .vc = (&d.vc)[0..1], .ikc = (&d.ikc)[0..1], .pooled = (&d.pooled)[0..1], .pg = d.pg, .tail = d.kc, .history = &d.unused_history, .cap = d.cap, .scores = score_rows }, .pos = d.len, .rows = rows }; +} + fn copyAsync(stream: cuda.Stream, dst: u64, src: u64, n: usize) !void { if (n == 0) return; const d = stream.d; @@ -248,7 +253,7 @@ pub fn Ops(comptime Tri: type) type { var at: usize = 0; for (ds, nexts, 0..) |d, nx, i| { total += nx.len; - lane_buf[i] = .{ .seq = .{ .media = null, .conv = &.{}, .state = &.{}, .kc = (&d.kc)[0..1], .vc = (&d.vc)[0..1], .ikc = (&d.ikc)[0..1], .pooled = (&d.pooled)[0..1], .tail = d.kc, .history = &d.unused_history, .cap = d.cap, .scores = score_rows }, .pos = d.len, .rows = nx.len }; + lane_buf[i] = headLane(d, nx.len, score_rows); @memcpy(ids_buf[at..][0..nx.len], nx); at += nx.len; } @@ -476,3 +481,14 @@ pub fn Ops(comptime Tri: type) type { } }; } + +test "a grouped head's lane carries its draft's page table" { + var d: Draft = undefined; + d.pg = .{ .pt = 0x7000, .ps = 4096, .ptok = 256 }; + d.len = 300; + d.cap = 1024; + const l = headLane(&d, 3, 0); + try std.testing.expectEqual(d.pg, l.seq.pg); + try std.testing.expectEqual(@as(usize, 300), l.pos); + try std.testing.expectEqual(@as(usize, 3), l.rows); +} From 3147b9c64c5b096a3ef9deba1189646fd336c1a5 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Sat, 7 Nov 2026 04:36:54 +0000 Subject: [PATCH 10/36] changelog: #2244 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 68295dee03..356b45fdf5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Flash Next on CUDA with `TF_FLASHNEXT_KV_PAGES`: a shared round's grouped draft head reads and writes each stream's own pages, so paged streams' replies equal their runs alone again, and paged decode runs as fast as unpaged; default settings are unchanged (#2244). Thanks to @Arminova. + - Tests: `zig build test-glm-snapshots-host` builds again: its copies module imports `core`, which `cuda_snaps.zig` reads for `TF_GLM_PROMPT_TAIL` since #2226 (#2403). Thanks to @Bizuayeu. - HIP: on RDNA 3.5 (gfx1150 to gfx1153) a prompt row gets the same bytes wherever its chunk starts. The 64-row prompt tile there now takes each key past a row's causal edge as +0, as RDNA 3.5's dot2 rounds a -0 lane apart from +0, so its prompt bits change, and gfx1151's prefill is up to 0.7% slower at 2k and 8k. Other chips keep main's prompt tile and bits, and gfx1201's prompt speed is main's within noise. On every chip, a non-finite value at a key past a row's edge now adds nothing to that row (#2357). Thanks to @laudney. From d9b34bbc041e266d64addc71eac6779e3c78aeee Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Tue, 27 Oct 2026 01:24:06 +0900 Subject: [PATCH 11/36] glm cuda: a prompt step takes two chunks while no other stream decodes, each one's exchanges and glue under the other's sites --- zig/src/families/glm/cuda.zig | 2 + zig/src/families/glm/cuda_engine.zig | 14 ++- zig/src/families/glm/cuda_follow.zig | 8 +- zig/src/families/glm/cuda_forward.zig | 27 +++-- zig/src/families/glm/cuda_forward_pair.zig | 85 ++++++++++++++ zig/src/families/glm/cuda_forward_pieces.zig | 105 +++++++++-------- zig/src/families/glm/cuda_lanes_test.zig | 3 +- zig/src/families/glm/cuda_prompt_pair.zig | 114 +++++++++++++++++++ zig/src/families/glm/cuda_ranks.zig | 14 ++- zig/src/families/glm/cuda_ranks_test.zig | 24 +++- zig/src/families/glm/cuda_state.zig | 27 ++++- zig/tests/cuda/glm_pair_tests.zig | 103 +++++++++++++++++ zig/tests/cuda/main.zig | 3 + 13 files changed, 461 insertions(+), 68 deletions(-) create mode 100644 zig/src/families/glm/cuda_forward_pair.zig create mode 100644 zig/src/families/glm/cuda_prompt_pair.zig create mode 100644 zig/tests/cuda/glm_pair_tests.zig diff --git a/zig/src/families/glm/cuda.zig b/zig/src/families/glm/cuda.zig index 340c9fa8db..72041dca43 100644 --- a/zig/src/families/glm/cuda.zig +++ b/zig/src/families/glm/cuda.zig @@ -32,6 +32,7 @@ pub const prompt_sink = @import("cuda_prompt_sink.zig"); pub const prompt_solo = @import("cuda_prompt_solo.zig"); pub const prompt_logprobs = @import("cuda_prompt_logprobs.zig"); pub const engine = @import("cuda_engine.zig"); +pub const prompt_pair = @import("cuda_prompt_pair.zig"); pub const pages = @import("cuda_pages.zig"); pub const share = @import("cuda_share.zig"); pub const open_owner = @import("cuda_open_owner.zig"); @@ -58,6 +59,7 @@ test { _ = @import("dflash_load.zig"); _ = prompt_solo; _ = @import("cuda_l2pf.zig"); + _ = prompt_pair; _ = experts; _ = exl3; _ = vision; diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index d96da8b910..829c5a39ff 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/zig/src/families/glm/cuda_engine.zig @@ -19,6 +19,7 @@ const dflash = @import("dflash.zig"); const prof_mod = @import("cuda_prof.zig"); const pages_mod = @import("cuda_pages.zig"); const l2pf_mod = @import("cuda_l2pf.zig"); +const prompt_pair = @import("cuda_prompt_pair.zig"); pub const Options = @import("cuda_engine_types.zig").Options; /// Rank 0's own word on its request (cancelled or not); the other ranks have none and hear rank 0's. @@ -72,6 +73,7 @@ pub const Engine = struct { pieces: fw.Pieces, w: weights.Weights, bp: fw.Buffers, // prompt chunks + pair: ?prompt_pair.Pair = null, // TF_GLM_PROMPT_PAIR: a paired prompt step's second chunk bd: fw.Buffers, // verify windows bm: fw.Buffers, // the MTP head's rows bs: ?fw.Buffers = null, // shared verify windows (with the MTP head, as the lane core shares only then) @@ -146,7 +148,7 @@ pub const Engine = struct { const block32 = !exl3.isExl3(gpa, text) and try @import("cuda_experts.zig").wantsBlock32(gpa, text); var m = try fw.Modules.load(d, ctx, block32); errdefer m.deinit(); - var pieces = try fw.Pieces.init(d, o.settings.pieces); + var pieces = try fw.Pieces.init(d, o.settings.pieces, o.settings.pair); errdefer pieces.deinit(); if (o.settings.layers > c.layers) return error.BadLayers; const layers = try a.alloc(u32, if (o.settings.layers > 0) o.settings.layers else c.layers); @@ -188,6 +190,8 @@ pub const Engine = struct { errdefer e.bp.deinit(); if (o.settings.solo and rank == 0) std.log.info("glm: prompt passes of {d} rows, {d} while another stream decodes", .{ fw.solo_rows, fw.prefill_rows }); if (world > 1 and rank == 0) std.log.info("glm: a prompt chunk's exchanges in {d} row pieces", .{o.settings.pieces}); + if (o.settings.pair) e.pair = try prompt_pair.Pair.init(d, &e.w, &e.bp, &e.pieces, o.capacity); + errdefer if (e.pair) |*p| p.deinit(); e.bd = try fw.Buffers.init(d, &e.w, fw.window_rows, false, o.capacity); errdefer e.bd.deinit(); e.bm = try fw.Buffers.init(d, &e.w, fw.window_rows, false, o.capacity); @@ -235,6 +239,7 @@ pub const Engine = struct { if (e.bs) |*b| b.deinit(); e.bm.deinit(); e.bd.deinit(); + if (e.pair) |*p| p.deinit(); e.bp.deinit(); if (e.pages) |p| p.deinit(); for (0..e.seqs.items.len) |i| e.free(i); @@ -255,10 +260,11 @@ pub const Engine = struct { return .{ .nccl = .{ .lib = e.lib, .comm = e.comm, .group = e.group } }; } - fn view(e: *Engine, b: *fw.Buffers, out: *[1]fw.Rank) fw.Forward { + pub fn view(e: *Engine, b: *fw.Buffers, out: *[1]fw.Rank) fw.Forward { out.* = .{.{ .w = &e.w, .st = &e.cur.st, .b = b }}; const im: ?fw.ImageRows = if (e.images) |v| .{ .v = v, .slot = e.slot } else null; - return .{ .d = e.d, .m = &e.m, .stream = e.stream, .ranks = out, .exchange = e.exchange(), .pieces = if (b == &e.bp) &e.pieces else null, .split = e.settings.split, .ring = e.settings.ring, .kda_wide = e.settings.kda_wide, .images = im, .taps = if (e.drafter) |dr| dr.taps(b == &e.bp) else null, .l2pf = if (e.l2pf) |*p| p else null }; + const twin = if (e.pair) |*p| (if (b == &p.bp) &p.pieces else null) else null; + return .{ .d = e.d, .m = &e.m, .stream = e.stream, .ranks = out, .exchange = e.exchange(), .pieces = if (b == &e.bp) &e.pieces else twin, .split = e.settings.split, .ring = e.settings.ring, .kda_wide = e.settings.kda_wide, .images = im, .taps = if (e.drafter) |dr| dr.taps(b == &e.bp) else null, .l2pf = if (e.l2pf) |*p| p else null }; } /// Binds slot `slot`'s sequence for the steps that follow, making it (a zeroed state) on first use. @@ -351,6 +357,7 @@ pub const Engine = struct { if ((from != null or start == 0) and try e.prefillStart(len, from) != start) return error.BadChunk; if (start >= len or start % fw.prefill_rows != 0) return error.BadChunk; const rows = @min(len - start, @as(usize, if (part.len > fw.prefill_rows + 1) fw.solo_rows else fw.prefill_rows)); + if (rows > fw.prefill_rows and e.pair != null) return prompt_pair.prefill(e, part, start, len, cancel, marks); if (part.len != @min(len - start, rows + 1)) return error.BadChunk; if (rows > fw.prefill_rows and !prompt_solo.allowed(e, start, rows)) return error.BadChunk; var one: [1]fw.Rank = undefined; @@ -523,6 +530,7 @@ pub const Engine = struct { } pub const holdPrompt = @import("cuda_prompt_hold.zig").save; + pub const pairedAt = prompt_pair.pairedAt; pub const prefillMany = @import("cuda_prompts.zig").prefillMany; /// forward.commit of the last window's first `n` rows (the KDA replay when fewer than it ran). diff --git a/zig/src/families/glm/cuda_follow.zig b/zig/src/families/glm/cuda_follow.zig index 7193b7de5f..f97f6a13bc 100644 --- a/zig/src/families/glm/cuda_follow.zig +++ b/zig/src/families/glm/cuda_follow.zig @@ -223,13 +223,13 @@ pub const Leader = struct { return x.e.prefill(ids, from, cancel, marks); } - /// One prompt step on every rank (Engine.prefillChunk): a chunk, or a solo pass if `solo` and prompt_solo.at allow. - pub fn chunk(x: Leader, ids: []const u32, start: usize, solo: bool, from: ?engine.Resume, cancel: ?engine.Cancel, marks: ?engine.Marks) !Chunk { + /// One prompt step on every rank (Engine.prefillChunk): a chunk, or two if `alone` and prompt_solo.at or the pair's rule allow. + pub fn chunk(x: Leader, ids: []const u32, start: usize, alone: bool, from: ?engine.Resume, cancel: ?engine.Cancel, marks: ?engine.Marks) !Chunk { if (start >= ids.len) return error.BadChunk; try x.flush(); - const n: usize = if (solo and prompt_solo.at(x.e, start, ids.len, from, marks)) engine.solo_rows else engine.prompt_rows; + const n: usize = if (alone and (prompt_solo.at(x.e, start, ids.len, from, marks) or x.e.pairedAt(start, ids.len, marks))) engine.solo_rows else engine.prompt_rows; const part = ids[start..@min(ids.len, start + n + 1)]; - const split = sparse.splitAt(start, x.e.settings.index_split, x.e.share.world); + const split = sparse.splitAt(start, x.e.settings.index_split, x.e.share.world); // both chunks of a pair share it try x.send(.{ .kind = @backingInt(Kind.chunk), .flags = x.rows() | @as(u32, @intFromBool(split)) << 1, .a = if (from) |r| r.id else 0, .seed = if (from) |r| r.at else 0, .depth = @intCast(start), .top_k = @intCast(ids.len) }, part); x.e.index_split = split; const first = try x.e.prefillChunk(part, start, ids.len, from, cancel, marks); diff --git a/zig/src/families/glm/cuda_forward.zig b/zig/src/families/glm/cuda_forward.zig index 75c0966e05..5cdd76ba99 100644 --- a/zig/src/families/glm/cuda_forward.zig +++ b/zig/src/families/glm/cuda_forward.zig @@ -135,10 +135,21 @@ pub const Pieces = struct { gathered: [max_pieces]cuda.Event, glues: [max_pieces]cuda.Event, // split: a piece's own rows glued shared: [max_pieces]cuda.Event, // split: every rank's glued rows of a piece in + owns_side: bool = true, // a paired step's second chunk (initTwin) runs on the first's side stream - pub fn init(d: *const cuda.Driver, count: usize) !Pieces { + /// With `pair` (TF_GLM_PROMPT_PAIR) the side stream runs at the highest priority, ahead of the main stream's grids. + pub fn init(d: *const cuda.Driver, count: usize, pair: bool) !Pieces { if (count == 0 or count > max_pieces) return error.BadPieces; - var p: Pieces = .{ .count = count, .side = try cuda.Stream.init(d, false), .filled = undefined, .gathered = undefined, .glues = undefined, .shared = undefined }; + return make(d, count, if (pair) try cuda.Stream.initPriority(d, .highest) else try cuda.Stream.init(d, false), true); + } + + /// A paired step's second chunk on `of`'s side stream: every rank posts both chunks' exchanges in one order. + pub fn initTwin(d: *const cuda.Driver, of: *const Pieces) !Pieces { + return make(d, of.count, of.side, false); + } + + fn make(d: *const cuda.Driver, count: usize, side: cuda.Stream, owns_side: bool) !Pieces { + var p: Pieces = .{ .count = count, .side = side, .filled = undefined, .gathered = undefined, .glues = undefined, .shared = undefined, .owns_side = owns_side }; for (&p.filled, &p.gathered, &p.glues, &p.shared) |*a, *b, *g, *s| { a.* = try cuda.Event.init(d, false); b.* = try cuda.Event.init(d, false); @@ -155,7 +166,7 @@ pub const Pieces = struct { g.deinit(); s.deinit(); } - p.side.deinit(); + if (p.owns_side) p.side.deinit(); p.* = undefined; } }; @@ -233,7 +244,7 @@ pub const Forward = struct { } /// forward._pre: hc_pre of every row at a site into b.normed, xs, post and comb. - fn pre(f: *const Forward, r: Rank, hc: weights.Hc, norm: glue.Norm, rows: usize) !void { + pub fn pre(f: *const Forward, r: Rank, hc: weights.Hc, norm: glue.Norm, rows: usize) !void { const stamp = prof_mod.span(f.prof, f.stream, .glue); defer stamp.end(); const c = &r.w.config; @@ -249,7 +260,7 @@ pub const Forward = struct { } /// forward._mixer: the layer's KDA or DSA block on b.normed; `li`/`di`: its index among loaded KDA / DSA layers. - fn mixer(f: *const Forward, r: Rank, l: weights.Layer, li: usize, di: usize, rows: usize, nch: usize, out: ?u64) !void { + pub fn mixer(f: *const Forward, r: Rank, l: weights.Layer, li: usize, di: usize, rows: usize, nch: usize, out: ?u64) !void { const c = &r.w.config; const b = r.b; const st = r.st; @@ -288,7 +299,7 @@ pub const Forward = struct { } /// forward._ffn: the dense MLP or MoE block on b.normed into this rank's slot. - fn ffn(f: *const Forward, r: Rank, l: weights.Layer, rows: usize, out: ?u64) !void { + pub fn ffn(f: *const Forward, r: Rank, l: weights.Layer, rows: usize, out: ?u64) !void { const stamp = prof_mod.span(f.prof, f.stream, .ffn); defer stamp.end(); const c = &r.w.config; @@ -413,6 +424,8 @@ pub const Forward = struct { if (f.taps) |t| try f.shareTaps(t, cut); } + pub const pairBody = @import("cuda_forward_pair.zig").pairBody; + /// forward.forward of `ids`: logits in b.logits; a prompt chunk's KDA layers committed, a window's at `commit`. pub fn run(f: *const Forward, ids: []const i32) !void { const R = ids.len; @@ -479,7 +492,7 @@ pub const Forward = struct { } /// The final norm and this rank's rows of the head (a prompt buffer's last row only). - fn finish(f: *const Forward, R: usize) !void { + pub fn finish(f: *const Forward, R: usize) !void { prof_mod.context(f.prof, null, null, .chunk); const stamp = prof_mod.span(f.prof, f.stream, .head); defer stamp.end(); diff --git a/zig/src/families/glm/cuda_forward_pair.zig b/zig/src/families/glm/cuda_forward_pair.zig new file mode 100644 index 0000000000..eec6992e70 --- /dev/null +++ b/zig/src/families/glm/cuda_forward_pair.zig @@ -0,0 +1,85 @@ +//! A paired prompt step's pass (TF_GLM_PROMPT_PAIR): two chunks site by site, each one's exchanges under the other's. +const glue = @import("cuda_glue.zig"); +const latent = @import("cuda_latent.zig"); +const fw = @import("cuda_forward.zig"); +const pieces = @import("cuda_forward_pieces.zig"); +const Forward = fw.Forward; +const Pieces = fw.Pieces; +const Site = fw.Forward.Site; +const room = fw.Forward.room; +const share = fw.Forward.share; + +/// Two consecutive prompt chunks of one sequence (`b` a view of `a`'s at its end), each on `pieced`'s kernels and bits. +pub fn pairBody(a: *const Forward, b: *const Forward, ra: usize, rb: usize) !void { + const two = [_]*const Forward{ a, b }; + const rows = [_]usize{ ra, rb }; + var cuts: [2][fw.max_pieces][2]usize = undefined; + var n: [2]usize = undefined; + for (two, rows, 0..) |f, R, i| { + if (f.pieces == null or f.prof != null or f.taps != null) return error.BadPair; + const c = &f.ranks[0].w.config; + for (f.ranks) |r| { + if (r.segs.len > 0 or !r.b.prefill or R == 0 or R > r.b.rows) return error.BadPair; + try glue.embedB16(f.set(), f.stream, r.b.ids, r.w.embed, r.b.x, R, c.hidden, c.hc); + } + if (f.images) |im| for (f.ranks) |r| try im.v.splice(f.stream, im.slot, r.st.pos, r.st.pos + R, r.b.x, c.hc); + n[i] = fw.ranges(R, f.pieces.?.count, &cuts[i]); + if (n[i] == 1 and !f.splits(R)) return error.BadPair; // run would take this chunk whole + for (f.ranks) |r| try f.pre(r, r.w.layers[0].attn_hc, r.w.layers[0].in_norm, R); + } + const layers = a.ranks[0].w.layers; + var li: usize = 0; + var di: usize = 0; + for (layers, 0..) |l, j| { + // b's mixer follows a's on the main stream, so a's caches, KDA state and conv rows of this layer are in + for (two, rows, 0..) |f, R, i| { + if (j > 0) try awaitShared(f, f.pieces.?, n[i]); + for (f.ranks) |r| try f.mixer(r, r.w.layers[j], li, di, R, @min(latent.chunksFor(r.st.pos + R), r.b.nch), null); + try exchangeGlue(f, f.pieces.?, j, .mixer, cuts[i][0..n[i]]); + } + for (two, rows, 0..) |f, R, i| { + try awaitShared(f, f.pieces.?, n[i]); + for (f.ranks) |r| try f.ffn(r, r.w.layers[j], R, null); + try exchangeGlue(f, f.pieces.?, j, .ffn, cuts[i][0..n[i]]); + } + switch (l.mixer) { + .kda => li += 1, + .mla => di += 1, + } + } + for (two, rows, 0..) |f, R, i| { + try awaitShared(f, f.pieces.?, n[i]); + try f.finish(R); + } +} + +/// A paired site: each piece filled on the main stream, then on the side stream its exchange, glue and (split) sharing. +fn exchangeGlue(f: *const Forward, p: *const Pieces, j: usize, site: Site, cut: []const [2]usize) !void { + const D = f.ranks[0].w.config.hidden; + const world = f.ranks[0].w.share.world; + const last = site == .ffn and j + 1 == f.ranks[0].w.layers.len; + var g = f.*; // the glue's launches go on the side stream + g.stream = p.side; + for (cut, 0..) |piece, q| { + const lo = piece[0]; + const n = piece[1] - lo; + const split = f.splits(n); + for (f.ranks) |r| try f.fillRows(r, r.w.layers[j], site, lo, piece[1], if (split) r.b.part + lo * D * 4 else room(r, lo) + r.w.share.rank * n * D * 4); + try p.filled[q].record(f.stream); + try p.side.wait(p.filled[q]); + try pieces.partials(f, p, lo, n, split); + if (split) { + for (f.ranks) |r| { + const mine = share(n, world, r.w.share.rank); + try g.glueRows(r, j, site, lo + mine[0], lo + mine[1], room(r, lo)); + } + try pieces.shareGlued(f, p, lo, n, last); + } else for (f.ranks) |r| try g.glueRows(r, j, site, lo, piece[1], room(r, lo)); + try p.shared[q].record(p.side); + } +} + +/// Before a paired chunk's next site: its last site's pieces glued and shared on the side stream. +fn awaitShared(f: *const Forward, p: *const Pieces, count: usize) !void { + for (p.shared[0..count]) |e| try f.stream.wait(e); +} diff --git a/zig/src/families/glm/cuda_forward_pieces.zig b/zig/src/families/glm/cuda_forward_pieces.zig index bb41202fe3..ba803b1afb 100644 --- a/zig/src/families/glm/cuda_forward_pieces.zig +++ b/zig/src/families/glm/cuda_forward_pieces.zig @@ -13,7 +13,6 @@ const share = fw.Forward.share; pub fn exchangePieces(f: *const Forward, p: *const Pieces, j: usize, site: Site, cut: []const [2]usize) !void { defer prof.context(f.prof, j, null, if (site == .mixer) .mixer else .ffn); const D = f.ranks[0].w.config.hidden; - const world = f.ranks[0].w.share.world; for (cut, 0..) |piece, q| { prof.context(f.prof, j, q, if (site == .mixer) .mixer else .ffn); const lo = piece[0]; @@ -25,33 +24,7 @@ pub fn exchangePieces(f: *const Forward, p: *const Pieces, j: usize, site: Site, try p.side.wait(p.filled[q]); stamp.end(); stamp = prof.span(f.prof, p.side, .partials); - const bytes = n * D * 4; - if (!split) switch (f.exchange) { - .local => for (f.ranks) |to| for (f.ranks) |from| { - if (to.w.share.rank == from.w.share.rank) continue; - const off = from.w.share.rank * bytes; - try f.copy(p.side, room(to, lo) + off, room(from, lo) + off, bytes); - }, - .nccl => |x| { - const r = f.ranks[0]; - try x.lib.check(x.lib.api.ncclAllGather(room(r, lo) + r.w.share.rank * bytes, room(r, lo), n * D, .f32, x.comm, p.side.handle), "ncclAllGather"); - }, - } else switch (f.exchange) { - .local => for (f.ranks) |to| { - const mine = share(n, world, to.w.share.rank); - const own = mine[1] - mine[0]; - for (f.ranks) |from| try f.copy(p.side, room(to, lo) + from.w.share.rank * own * D * 4, from.b.part + (lo + mine[0]) * D * 4, own * D * 4); - }, - .nccl => |x| { - const r = f.ranks[0]; - const me = r.w.share.rank; - const mine = share(n, world, me); - const own = mine[1] - mine[0]; - const plan: relay.Relay = .{ .world = world, .rank = me, .n = n }; - if (f.ring) try relay.partials(x.lib, x.comm, p.side, plan, room(r, lo), r.b.part, lo, D) else try relay.splitPartials(x.lib, x.comm, p.side, plan, room(r, lo), r.b.part, lo, D); - try f.copy(p.side, room(r, lo) + me * own * D * 4, r.b.part + (lo + mine[0]) * D * 4, own * D * 4); - }, - } + try partials(f, p, lo, n, split); stamp.end(); try p.gathered[q].record(p.side); } @@ -60,8 +33,6 @@ pub fn exchangePieces(f: *const Forward, p: *const Pieces, j: usize, site: Site, /// overlap.Overlap.glue: each piece's glue once its exchange is in, its rows then shared on the second stream. pub fn gluePieces(f: *const Forward, p: *const Pieces, j: usize, site: Site, cut: []const [2]usize) !void { defer prof.context(f.prof, j, null, if (site == .mixer) .mixer else .ffn); - const c = &f.ranks[0].w.config; - const D: usize = c.hidden; const world = f.ranks[0].w.share.world; const last = site == .ffn and j + 1 == f.ranks[0].w.layers.len; var sent: usize = 0; @@ -85,23 +56,7 @@ pub fn gluePieces(f: *const Forward, p: *const Pieces, j: usize, site: Site, cut try p.side.wait(p.glues[q]); stamp.end(); stamp = prof.span(f.prof, p.side, .glued); - switch (f.exchange) { - .local => for (f.ranks) |from| { - const mine = share(n, world, from.w.share.rank); - const bytes = (mine[1] - mine[0]) * D * 2; - for (f.ranks) |to| if (to.w.share.rank != from.w.share.rank) { - const at = (lo + mine[0]) * D * 2; - try f.copy(p.side, (if (last) to.b.hidden else to.b.normed) + at, (if (last) from.b.hidden else from.b.normed) + at, bytes); - }; - }, - .nccl => |x| { - const r = f.ranks[0]; - const me = r.w.share.rank; - const rows_of = if (last) r.b.hidden else r.b.normed; - const plan: relay.Relay = .{ .world = world, .rank = me, .n = n }; - if (f.ring) try relay.glued(x.lib, x.comm, p.side, plan, rows_of, lo, D) else try relay.splitGlued(x.lib, x.comm, p.side, plan, rows_of, lo, D); - }, - } + try shareGlued(f, p, lo, n, last); stamp.end(); try p.shared[q].record(p.side); sent |= @as(usize, 1) << @intCast(q); @@ -151,3 +106,59 @@ pub fn shareTaps(f: *const Forward, t: taps.Taps, cut: []const [2]usize) !void { if (f.ring) try relay.glued(x.lib, x.comm, f.stream, plan, t.out, piece[0], dims) else try relay.splitGlued(x.lib, x.comm, f.stream, plan, t.out, piece[0], dims); } } + +/// A piece's partials on the side stream: gathered whole, or split (each rank's rows of every partial to their owner). +pub fn partials(f: *const Forward, p: *const Pieces, lo: usize, n: usize, split: bool) !void { + const D = f.ranks[0].w.config.hidden; + const world = f.ranks[0].w.share.world; + const bytes = n * D * 4; + if (!split) switch (f.exchange) { + .local => for (f.ranks) |to| for (f.ranks) |from| { + if (to.w.share.rank == from.w.share.rank) continue; + const off = from.w.share.rank * bytes; + try f.copy(p.side, room(to, lo) + off, room(from, lo) + off, bytes); + }, + .nccl => |x| { + const r = f.ranks[0]; + try x.lib.check(x.lib.api.ncclAllGather(room(r, lo) + r.w.share.rank * bytes, room(r, lo), n * D, .f32, x.comm, p.side.handle), "ncclAllGather"); + }, + } else switch (f.exchange) { + .local => for (f.ranks) |to| { + const mine = share(n, world, to.w.share.rank); + const own = mine[1] - mine[0]; + for (f.ranks) |from| try f.copy(p.side, room(to, lo) + from.w.share.rank * own * D * 4, from.b.part + (lo + mine[0]) * D * 4, own * D * 4); + }, + .nccl => |x| { + const r = f.ranks[0]; + const me = r.w.share.rank; + const mine = share(n, world, me); + const own = mine[1] - mine[0]; + const plan: relay.Relay = .{ .world = world, .rank = me, .n = n }; + if (f.ring) try relay.partials(x.lib, x.comm, p.side, plan, room(r, lo), r.b.part, lo, D) else try relay.splitPartials(x.lib, x.comm, p.side, plan, room(r, lo), r.b.part, lo, D); + try f.copy(p.side, room(r, lo) + me * own * D * 4, r.b.part + (lo + mine[0]) * D * 4, own * D * 4); + }, + } +} + +/// A split piece's glued rows (normed, or hidden after the last layer) to every rank on the side stream. +pub fn shareGlued(f: *const Forward, p: *const Pieces, lo: usize, n: usize, last: bool) !void { + const D: usize = f.ranks[0].w.config.hidden; + const world = f.ranks[0].w.share.world; + switch (f.exchange) { + .local => for (f.ranks) |from| { + const mine = share(n, world, from.w.share.rank); + const bytes = (mine[1] - mine[0]) * D * 2; + for (f.ranks) |to| if (to.w.share.rank != from.w.share.rank) { + const at = (lo + mine[0]) * D * 2; + try f.copy(p.side, (if (last) to.b.hidden else to.b.normed) + at, (if (last) from.b.hidden else from.b.normed) + at, bytes); + }; + }, + .nccl => |x| { + const r = f.ranks[0]; + const me = r.w.share.rank; + const rows_of = if (last) r.b.hidden else r.b.normed; + const plan: relay.Relay = .{ .world = world, .rank = me, .n = n }; + if (f.ring) try relay.glued(x.lib, x.comm, p.side, plan, rows_of, lo, D) else try relay.splitGlued(x.lib, x.comm, p.side, plan, rows_of, lo, D); + }, + } +} diff --git a/zig/src/families/glm/cuda_lanes_test.zig b/zig/src/families/glm/cuda_lanes_test.zig index af503aff49..1eca112ac8 100644 --- a/zig/src/families/glm/cuda_lanes_test.zig +++ b/zig/src/families/glm/cuda_lanes_test.zig @@ -1,7 +1,8 @@ -//! The lane backend's slots: chosen before a stream's map entry exists, and none left behind by a failed insertion. +//! The lane backend's slots, chosen and given back, and when a prompt step may pair. const std = @import("std"); const lanes = @import("lanes"); const engine = @import("cuda_engine.zig"); +const fw = @import("cuda_forward.zig"); const Lanes = @import("cuda_lanes.zig").Lanes; const fw = @import("cuda_forward.zig"); const solo = @import("cuda_prompt_solo.zig"); diff --git a/zig/src/families/glm/cuda_prompt_pair.zig b/zig/src/families/glm/cuda_prompt_pair.zig new file mode 100644 index 0000000000..0a1c7f9ef8 --- /dev/null +++ b/zig/src/families/glm/cuda_prompt_pair.zig @@ -0,0 +1,114 @@ +//! A prompt step of two chunks in one pass (TF_GLM_PROMPT_PAIR): each chunk's exchanges under the other's sites. +const std = @import("std"); +const cuda = @import("cuda"); +const fw = @import("cuda_forward.zig"); +const weights = @import("cuda_weights.zig"); +const snapshots = @import("cuda_snapshots.zig"); +const engine = @import("cuda_engine.zig"); +const Engine = engine.Engine; + +/// The fewest rows a paired step's second chunk takes: a shorter one's sites end before the first's exchanges do. +pub const min_rows = fw.prefill_rows / 4; + +/// A paired step's second chunk: twins of the prompt buffers, pieces on their side stream, and its first position. +pub const Pair = struct { + bp: fw.Buffers, + pieces: fw.Pieces, + pos: cuda.DeviceBuffer, // an i32: where the second chunk's view of the sequence starts + steps: u64 = 0, // paired steps so far (glm-pair-check counts them) + + pub fn init(d: *const cuda.Driver, w: *const weights.Weights, bp: *const fw.Buffers, pieces: *const fw.Pieces, capacity: usize) !Pair { + var b = try fw.Buffers.initTwin(d, w, bp, capacity); + errdefer b.deinit(); + var p = try fw.Pieces.initTwin(d, pieces); + errdefer p.deinit(); + return .{ .bp = b, .pieces = p, .pos = try cuda.DeviceBuffer.alloc(d, 4) }; + } + + pub fn deinit(p: *Pair) void { + p.pos.free(); + p.pieces.deinit(); + p.bp.deinit(); + p.* = undefined; + } +}; + +/// Whether a step at `start` of a `len`-id prompt takes the next chunk too: min_rows or more there, and no mark at it. +pub fn takesTwo(start: usize, len: usize, world: usize, marks: []const u32) bool { + const next = start + fw.prefill_rows; + if (next >= len or len - next < @max(min_rows, world)) return false; + return std.mem.indexOfScalar(u32, marks, @intCast(next)) == null; +} + +/// Rank 0's choice for the step at `start` (the others take the ids it sends); not with a drafter, prompt rows or prof. +pub fn pairedAt(e: *const Engine, start: usize, len: usize, marks: ?engine.Marks) bool { + if (e.pair == null or !e.settings.pair or e.drafter != null or e.prompt_logprobs != null or e.prompt_profs[3] != null) return false; + if (e.pieces.count == 1 and !e.settings.split) return false; // run takes such a chunk whole + return takesTwo(start, len, e.share.world, if (marks) |m| m.at else &.{}); +} + +/// Engine.prefillChunk of two chunks (`part`: both, then the next id): one pass, then each chunk's MTP rows and commit. +pub fn prefill(e: *Engine, part: []const u32, start: usize, len: usize, cancel: ?engine.Cancel, marks: ?engine.Marks) !?u32 { + const ra: usize = fw.prefill_rows; + if (len <= start + ra) return error.BadChunk; + const rb: usize = @min(len - start - ra, fw.prefill_rows); + const pair = if (e.pair) |*p| p else return error.BadChunk; + if (rb < e.share.world or part.len != @min(len - start, ra + rb + 1) or e.drafter != null or e.prompt_logprobs != null) return error.BadChunk; + if (marks) |m| if (start > 0 and std.mem.indexOfScalar(u32, m.at, @intCast(start)) != null) m.call(m.ctx, @intCast(start)); + e.gate(cancel) catch |err| { + if (err == error.Cancelled) try snapshots.interrupted(e, marks, @intCast(start)); + return err; + }; + var one: [1]fw.Rank = undefined; + var two: [1]fw.Rank = undefined; + var fa = e.view(&e.bp, &one); + var fb = e.view(&pair.bp, &two); + fa.prof = null; + fb.prof = null; + const st = &e.cur.st; + try st.ready(e.stream, ra + rb); // one rebase for both chunks, before b's view copies the rings' base + var sb = st.*; + sb.pos_dev = pair.pos.ptr; + try sb.setPos(e.stream, st.pos + ra); + two[0].st = &sb; + st.pending = ra; + st.prompt = true; + const ids: []const i32 = @ptrCast(part); + try cuda.copy_source.enqueue(e.d, e.bp.ids, std.mem.sliceAsBytes(ids[0..ra]), e.stream); + try cuda.copy_source.enqueue(e.d, pair.bp.ids, std.mem.sliceAsBytes(ids[ra .. ra + rb]), e.stream); + try fw.Forward.pairBody(&fa, &fb, ra, rb); + pair.steps += 1; + try mtpRows(e, &fa, ids[1 .. ra + 1], e.bp.fnormed); + try fa.commit(ra); + st.pending = rb; + st.prompt = true; + var first: ?u32 = null; + if (start + ra + rb == len) { + var t: [1]u32 = undefined; + try e.draw(pair.bp.logits, 1, len, 0, &t); + first = t[0]; + try e.holdPrompt(pair.bp.fnormed + (rb - 1) * e.c.hidden * 2); + } + try mtpRows(e, &fa, ids[ra + 1 ..], pair.bp.fnormed); // the MTP head on a's buffers, over b's rows + try fa.commit(rb); + return first; +} + +/// decode._absorb_rows of a chunk: the MTP head over its next ids `nxt` on its final-norm rows `hidden`. +fn mtpRows(e: *Engine, f: *fw.Forward, nxt: []const i32, hidden: u64) !void { + if (!e.mtp or nxt.len == 0) return; + try f.mtp(nxt, &.{hidden}); + try e.cur.st.setMtpLen(e.stream, e.cur.st.mtp_len + nxt.len); +} + +test "a prompt step takes the next chunk when it holds min_rows or more and no kept state starts there" { + const C = fw.prefill_rows; + try std.testing.expect(takesTwo(0, 2 * C, 3, &.{})); + try std.testing.expect(takesTwo(0, C + min_rows, 3, &.{})); + try std.testing.expect(!takesTwo(0, C + min_rows - 1, 3, &.{})); // a shorter tail goes alone + try std.testing.expect(!takesTwo(0, C, 3, &.{})); + try std.testing.expect(!takesTwo(C, 2 * C, 3, &.{})); // the prompt's last chunk + try std.testing.expect(takesTwo(2 * C, 5 * C, 3, &.{})); + try std.testing.expect(!takesTwo(0, 3 * C, 3, &.{C})); // a state kept at the next chunk's start: that chunk starts a step + try std.testing.expect(takesTwo(0, 3 * C, 3, &.{ 2 * C, 3 * C })); // marks past the pair +} diff --git a/zig/src/families/glm/cuda_ranks.zig b/zig/src/families/glm/cuda_ranks.zig index 6b17c077a5..3282877a6d 100644 --- a/zig/src/families/glm/cuda_ranks.zig +++ b/zig/src/families/glm/cuda_ranks.zig @@ -18,6 +18,7 @@ pub const Settings = struct { kda_pad: bool = false, // TF_GLM_KDA_PAD: 1 (the fused KDA input padded to eight columns at load, for glm_mm) or 0 pieces: u32 = 3, // TF_GLM_OVERLAP_PIECES: a prompt chunk's exchanges in this many row pieces (1: whole); unset, 8 under solo on the wire. solo: bool = false, // TF_GLM_CHUNK_SOLO: 4096 (prompt buffers for two chunks a pass, while no stream decodes) or 2048 + pair: bool = true, // TF_GLM_PROMPT_PAIR: 1 (while no stream decodes, a prompt step takes two chunks, each one's exchanges under the other's sites) or 0 index_split: u32 = 2048, // TF_GLM_INDEX_SPLIT: from this position on, rank 0 splits a prompt chunk's indexer rows (0: never) heat: ?heat.Bands = null, // TF_HEAT_HIGH and TF_HEAT_LOW (upstream's names): a prompt chunk waits on every rank's hottest zone layers: u32 = 0, // TF_GLM_LAYERS: the first this many layers only, a thin model for one-GPU checks (0: all) @@ -84,6 +85,12 @@ pub const Settings = struct { } } s.l2pf = try l2pf.Options.read(get, problem); + if (get("TF_GLM_PROMPT_PAIR")) |v| { + if (std.mem.eql(u8, v, "1")) s.pair = true else if (std.mem.eql(u8, v, "0")) s.pair = false else { + problem.* = "TF_GLM_PROMPT_PAIR: 0 or 1"; + return error.BadSetting; + } + } if (get("TF_GLM_OVERLAP_PIECES")) |v| { s.pieces = std.fmt.parseInt(u32, v, 10) catch 0; if (s.pieces < 1 or s.pieces > 8) { @@ -176,7 +183,7 @@ fn wireOpens(get: *const fn ([:0]const u8) ?[]const u8) bool { /// What every rank must agree on, as bytes the communicator carries. pub const Fingerprint = extern struct { - layout: u32 = 18, // both layouts: main 16 plus L2 prefetch, distinct from the PR's 17 + layout: u32 = 19, // this struct's version (18 plus the prompt pair) world: u32, rank: u32, kv8: u32 = 0, @@ -187,6 +194,7 @@ pub const Fingerprint = extern struct { kda_wide: u32, kda_pad: u32 = 0, // TF_GLM_KDA_PAD: every rank's KDA input rows carry one width pieces: u32, + pair: u32, // TF_GLM_PROMPT_PAIR: rank 0's paired steps need the second chunk's buffers on every rank layers: u32, // TF_GLM_LAYERS (0: all) mtp: u32, // 1: the MTP head loads and runs (drafts on); 2: and DFlash2 (--drafter), whose rounds gather vision: u32 = 0, // --vision's image tokens (0: off): rank 0 sends image rows every rank must hold @@ -219,6 +227,7 @@ pub const Fingerprint = extern struct { kda_wide, kda_pad, pieces, + pair, layers, mtp, vision, @@ -261,6 +270,7 @@ pub const Fingerprint = extern struct { .kda_wide = @intFromBool(s.kda_wide), .kda_pad = @intFromBool(s.kda_pad), .pieces = s.pieces, + .pair = @intFromBool(s.pair), .layers = s.layers, .mtp = @intFromBool(mtp), .roce_multi = @intFromBool(s.roce_multi), @@ -297,6 +307,7 @@ pub fn mismatch(all: []const Fingerprint) ?Mismatch { "kda_wide", "kda_pad", "pieces", + "pair", "layers", "mtp", "vision", @@ -418,6 +429,7 @@ fn explain(f: Fingerprint.Field) []const u8 { .kda_wide => "TF_GLM_KDA_DECODE_WIDE", .kda_pad => "TF_GLM_KDA_PAD", .pieces => "TF_GLM_OVERLAP_PIECES", + .pair => "TF_GLM_PROMPT_PAIR", .layers => "TF_GLM_LAYERS", .mtp => "--no-drafts given to some ranks and not others", .vision => "--vision (or its image tokens) given to some ranks and not others", diff --git a/zig/src/families/glm/cuda_ranks_test.zig b/zig/src/families/glm/cuda_ranks_test.zig index bd2e6e03ed..dc5bbbd0a5 100644 --- a/zig/src/families/glm/cuda_ranks_test.zig +++ b/zig/src/families/glm/cuda_ranks_test.zig @@ -190,7 +190,7 @@ test "every media limit is part of the ranks' agreement" { test "every scheduling and identity field that differs stops the ranks" { const zero = Fingerprint.of(2, 0, .{}, 8192, true, @splat(0), @splat(0), hashBytes("/same")); const one = Fingerprint.of(2, 1, .{}, 8192, true, @splat(0), @splat(0), hashBytes("/same")); - inline for (.{ "layout", "world", "rank", "pages", "split", "kda_wide", "kda_pad", "pieces", "solo", "layers", "mtp", "vision", "video_tokens", "max_images", "max_videos", "max_frames", "roce_multi", "kv_host", "capacity", "heat_high", "heat_low", "kv_host_gib" }) |name| { + inline for (.{ "layout", "world", "rank", "pages", "split", "kda_wide", "kda_pad", "pieces", "solo", "pair", "layers", "mtp", "vision", "video_tokens", "max_images", "max_videos", "max_frames", "roce_multi", "kv_host", "capacity", "heat_high", "heat_low", "kv_host_gib" }) |name| { var all = [_]Fingerprint{ zero, one }; @field(all[1], name) += 1; try std.testing.expectEqual(@field(Fingerprint.Field, name), mismatch(&all).?.field); @@ -388,6 +388,28 @@ test "prompt K experiment is explicit and all ranks must agree" { try std.testing.expectEqual(@as(?Mismatch, null), mismatch(&all)); } +test "TF_GLM_PROMPT_PAIR: on by default, 0 or 1, and every rank must agree (the others take rank 0's paired steps)" { + const Reader = struct { + fn absent(_: [:0]const u8) ?[]const u8 { + return null; + } + fn off(key: [:0]const u8) ?[]const u8 { + return if (std.mem.eql(u8, key, "TF_GLM_PROMPT_PAIR")) "0" else null; + } + fn bad(key: [:0]const u8) ?[]const u8 { + return if (std.mem.eql(u8, key, "TF_GLM_PROMPT_PAIR")) "2" else null; + } + }; + var problem: []const u8 = ""; + try std.testing.expect((try Settings.read(Reader.absent, &problem)).pair); + const off = try Settings.read(Reader.off, &problem); + try std.testing.expect(!off.pair); + try std.testing.expectError(error.BadSetting, Settings.read(Reader.bad, &problem)); + try std.testing.expectEqualStrings("TF_GLM_PROMPT_PAIR: 0 or 1", problem); + const all = [_]Fingerprint{ Fingerprint.of(2, 0, .{}, 8192, true, @splat(0), @splat(0), @splat(0)), Fingerprint.of(2, 1, off, 8192, true, @splat(0), @splat(0), @splat(0)) }; + try std.testing.expectEqual(Mismatch{ .field = .pair, .rank = 1 }, mismatch(&all).?); +} + test "TF_GLM_INDEX_SPLIT: 2048 by default, 0 for never, a position, and a refusal naming a value that is not one" { const Reader = struct { fn absent(_: [:0]const u8) ?[]const u8 { diff --git a/zig/src/families/glm/cuda_state.zig b/zig/src/families/glm/cuda_state.zig index ff3983f0f4..da791c2f24 100644 --- a/zig/src/families/glm/cuda_state.zig +++ b/zig/src/families/glm/cuda_state.zig @@ -177,16 +177,25 @@ pub const Buffers = struct { part: u64, // fp32 [rows][dims]: this rank's partial rows a split exchange sends from pub fn init(d: *const cuda.Driver, w: *const weights.Weights, rows: usize, prefill: bool, capacity: usize) !Buffers { - return make(d, w, rows, prefill, capacity, rows); + return make(d, w, rows, prefill, capacity, rows, null); } /// A shared window's scratch: `rows` rows of several streams, each one's attention and KDA chain on `window_rows`. pub fn initShared(d: *const cuda.Driver, w: *const weights.Weights, rows: usize, capacity: usize) !Buffers { - return make(d, w, rows, false, capacity, window_rows); + return make(d, w, rows, false, capacity, window_rows, null); } + /// A paired step's second chunk: its own rows across sites and streams, and `of`'s scratch for one site's work. + pub fn initTwin(d: *const cuda.Driver, w: *const weights.Weights, of: *const Buffers, capacity: usize) !Buffers { + return make(d, w, of.rows, of.prefill, capacity, of.rows, of); + } + + /// The fields a twin holds apart: make's sizes by index, then their names. + const twin_sizes = [_]usize{ 0, 1, 2, 3, 4, 5, 6, 7, 62, 63, 64, 65, 70 }; + const twin_fields = .{ "ids", "x", "normed", "xs", "post", "comb", "hcpart", "gathered", "hidden", "fnormed", "fxs", "logits", "part" }; + /// `seg`: the most rows one stream's attention takes (a decode window's rows, or a shared window's segment). - fn make(d: *const cuda.Driver, w: *const weights.Weights, rows: usize, prefill: bool, capacity: usize, seg: usize) !Buffers { + fn make(d: *const cuda.Driver, w: *const weights.Weights, rows: usize, prefill: bool, capacity: usize, seg: usize, twin: ?*const Buffers) !Buffers { const c = &w.config; const s = &w.share; const D: usize = c.hidden; @@ -252,8 +261,12 @@ pub const Buffers = struct { x3[0], x3[1], x3[2], x3[3], // EXL3 experts: 8-slot picks, rotated inputs, a decode window's partials if (split) rows * @max(3 * n, kw - 3 * n) * 2 else 4, // a split KDA stack's part rows }; + var held = sizes; + if (twin != null) for (&held, 0..) |*bytes, i| { // a twin allocates its own rows only + if (std.mem.indexOfScalar(usize, &twin_sizes, i) == null) bytes.* = 0; + }; var p: [sizes.len]u64 = undefined; - const a = try arena(d, &sizes, &p); + const a = try arena(d, &held, &p); var out: Buffers = .{ .arena = a, .rows = rows, @@ -306,6 +319,12 @@ pub const Buffers = struct { }; out.moe.x3 = .{ .picks = p[71], .xg = p[72], .xu = p[73], .z = p[74] }; out.kda.tmp = p[75]; + if (twin) |of| { + var t = of.*; + t.arena = out.arena; + inline for (twin_fields) |f| @field(t, f) = @field(out, f); + return t; + } return out; } diff --git a/zig/tests/cuda/glm_pair_tests.zig b/zig/tests/cuda/glm_pair_tests.zig new file mode 100644 index 0000000000..ff03559ed8 --- /dev/null +++ b/zig/tests/cuda/glm_pair_tests.zig @@ -0,0 +1,103 @@ +//! GLM-5.3-Flash's paired prompt steps (TF_GLM_PROMPT_PAIR) against twins a chunk a step, on one Engine and its ranks. + +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const Gpu = check.Gpu; + +/// `glm-pair-check`: prompts a step at a time, paired and not; first token, next window and every state byte equal. +pub fn pairCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c_int, world: c_int, id_path: []const u8) !void { + const glm = @import("glm_cuda"); + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const dc = try std.json.parseFromSliceLeaky(struct { capacity: usize = 8192 }, a, try std.Io.Dir.cwd().readFileAlloc(gpu.io, dc_path, a, .limited(1 << 26)), .{ .ignore_unknown_fields = true }); + var lib = try cuda.nccl.Library.open(); + defer lib.close(); + const comm = try @import("glm_lanes_tests.zig").openComm(gpu, &lib, rank, world, id_path); + defer _ = lib.api.ncclCommDestroy(comm); + var problem: []const u8 = ""; + const set = glm.ranks.Settings.read(&glm.ranks.environment, &problem) catch |err| { + std.debug.print("REFUSE {s}\n", .{problem}); + return err; + }; + try check.expect(set.pair, "glm-pair-check needs TF_GLM_PROMPT_PAIR on (unset or 1)", .{}); + const C = glm.engine.prompt_rows; + const capacity = @max(dc.capacity, 4 * C); + const e = try glm.engine.Engine.init(gpu.gpa, gpu.io, gpu.d, gpu.ctx, checkpoint, &lib, comm, @intCast(rank), @intCast(world), .{ .capacity = capacity, .mtp = true, .settings = set }); + defer e.deinit(); + var link = try glm.follow.Link.init(gpu.gpa, gpu.d, &lib, comm, e.stream, capacity); + defer link.deinit(); + if (rank != 0) { + try glm.follow.follow(e, &link); + check.pass("rank {d} followed rank 0's paired and single prompt steps to the end", .{rank}); + return; + } + const x: glm.follow.Leader = .{ .e = e, .link = &link }; + defer x.end() catch {}; + const V = e.share.vocab.len(); + const min = glm.prompt_pair.min_rows; + // a pair and a tail, the shortest second chunk, a pair and a chunk, a pair to the end, a tail too short to pair + const lens = [_]usize{ 2 * C + 300, C + min, 3 * C, 2 * C, C + min - 1 }; + var failed = false; + for (lens, 0..) |len, i| { + const ids = try a.alloc(u32, len); + for (ids, 0..) |*u, j| u.* = @intCast(100 + (j * 7919 + i * 31) % 5000); + var firsts: [2]u32 = undefined; + var outs: [2][4]u32 = undefined; + var logits: [2][]u16 = undefined; + var paired: [2]u64 = undefined; + for (0..2) |twin| { // twin 0 a chunk a step, twin 1 paired where it can + e.settings.pair = twin == 1; + try x.bind(@intCast(2 * i + twin)); + try x.begin(null); + const before = e.pair.?.steps; + var at: usize = 0; + const t = while (true) { + const r = x.chunk(ids, at, true, null, null, null) catch |err| { + std.debug.print("FAIL pair: prompt of {d}, twin {d}: the step at {d}: {s}\n", .{ len, twin, at, @errorName(err) }); + return err; + }; + if (r.first) |f| break f; + at += r.rows; + }; + paired[twin] = e.pair.?.steps - before; + firsts[twin] = t; + const win = [_]u32{ t, 101, 202, 303 }; + try x.verify(&win, &outs[twin], 0); + logits[twin] = try a.alloc(u16, win.len * V); + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(logits[twin].ptr, e.bd.logits, win.len * V * 2), "cuMemcpyDtoH"); + } + e.settings.pair = true; + var want: u64 = 0; // the steps glm.prompt_pair.takesTwo pairs over this prompt + var at: usize = 0; + while (at < len) { + const two = glm.prompt_pair.takesTwo(at, len, e.share.world, &.{}); + want += @intFromBool(two); + at += if (two) 2 * C else C; + } + if (paired[0] != 0 or paired[1] != want) { + std.debug.print("FAIL pair: prompt of {d}: {d} paired steps a chunk a step, {d} paired (want {d})\n", .{ len, paired[0], paired[1], want }); + failed = true; + } + if (firsts[0] != firsts[1] or !std.mem.eql(u32, &outs[0], &outs[1]) or !std.mem.eql(u16, logits[0], logits[1])) { + std.debug.print("FAIL pair: prompt of {d}: first {d}/{d}, the next window's tokens or logits differ\n", .{ len, firsts[0], firsts[1] }); + failed = true; + } + try e.stream.synchronize(); + const st = &e.seqs.items[2 * i].?.st; + const tw = &e.seqs.items[2 * i + 1].?.st; + const n = st.arena.buf.len; + const h0 = try a.alloc(u8, n); + const h1 = try a.alloc(u8, n); + try st.arena.buf.download(0, h0); + try tw.arena.buf.download(0, h1); + const diff = std.mem.indexOfDiff(u8, h0, h1); + if (st.pos != tw.pos or st.mtp_len != tw.mtp_len or !std.mem.eql(u1, st.cur, tw.cur) or diff != null) { + std.debug.print("FAIL pair: prompt of {d}: the paired state differs (pos {d}/{d}, first byte {?d} of {d})\n", .{ len, st.pos, tw.pos, diff, n }); + failed = true; + } + } + if (failed) return error.TestFailed; + check.pass("paired prompt steps at TP={d}: prompts of {any} rows equal their twins a chunk a step in the first token, the next window's logits and tokens, and every state byte", .{ world, lens }); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index f816738b4d..94bb9719ce 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -68,6 +68,7 @@ const lfm2_kernel_tests = @import("lfm2_kernel_tests.zig"); const lfm2_vision_tests = @import("lfm2_vision_tests.zig"); const radio_tests = @import("radio_tests.zig"); const glm_lanes_tests = @import("glm_lanes_tests.zig"); +const glm_pair_tests = @import("glm_pair_tests.zig"); const glm_split_tests = @import("glm_split_tests.zig"); const glm_solo_tests = @import("glm_solo_tests.zig"); const glm_paged_tests = @import("glm_paged_tests.zig"); @@ -273,6 +274,7 @@ const usage = \\ glm-lanes-cover copy drafts, cancels, discards, eos and max_new on one lanes.Engine across requests \\ glm-lanes-concurrent concurrent streams on one lanes.Engine, each against its solo run \\ glm-shared-check shared verify windows against lone twins, logits and states byte for byte + \\ glm-pair-check paired prompt steps against twins a chunk a step, logits and states byte for byte \\ glm-split-check prompt chunks with their indexer rows split across ranks against whole twins, byte for byte \\ glm-solo-check solo prompt passes against twins a chunk a step, logits and states byte for byte \\ glm-tap-split-check DFlash2 prompt taps: split and ring retained taps against gather, byte for byte, every rank @@ -530,6 +532,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "glm-lanes-cover")) return glm_lanes_tests.coverCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); if (std.mem.eql(u8, cmd, "glm-paged-check")) return glm_paged_tests.pagedCheck(gpu, try arg(rest, 0), try std.fmt.parseInt(c_int, try arg(rest, 1), 10), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try arg(rest, 3)); if (std.mem.eql(u8, cmd, "glm-shared-check")) return glm_lanes_tests.sharedCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); + if (std.mem.eql(u8, cmd, "glm-pair-check")) return glm_pair_tests.pairCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); if (std.mem.eql(u8, cmd, "glm-split-check")) return glm_split_tests.splitCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); if (std.mem.eql(u8, cmd, "glm-solo-check")) return glm_solo_tests.soloCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); if (std.mem.eql(u8, cmd, "glm-tap-split-check")) return glm_tap_tests.tapSplitCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); From 97c7d40e5e1aaa919829af0ba76a304d02f1e0f7 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Tue, 27 Oct 2026 22:57:14 +0900 Subject: [PATCH 12/36] glm cuda: a paired step takes rank 0's indexer split for both chunks and counts it once (#1404) Engine.prefillChunk reads and clears the split flag before it hands a paired step to prompt_pair.prefill, which never set it on its two forwards, so with #1404 a paired step ran the indexer whole on every rank. The step's flag now reaches both chunks: the header carries the one choice rank 0 makes at the step's start, and index_splits counts the step once. The twin buffers already share the first chunk's select scratch, whose pools hold #1404's world rows more. --- zig/src/families/glm/cuda_engine.zig | 2 +- zig/src/families/glm/cuda_prompt_pair.zig | 5 ++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index 829c5a39ff..186d61a724 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/zig/src/families/glm/cuda_engine.zig @@ -357,7 +357,7 @@ pub const Engine = struct { if ((from != null or start == 0) and try e.prefillStart(len, from) != start) return error.BadChunk; if (start >= len or start % fw.prefill_rows != 0) return error.BadChunk; const rows = @min(len - start, @as(usize, if (part.len > fw.prefill_rows + 1) fw.solo_rows else fw.prefill_rows)); - if (rows > fw.prefill_rows and e.pair != null) return prompt_pair.prefill(e, part, start, len, cancel, marks); + if (rows > fw.prefill_rows and e.pair != null) return prompt_pair.prefill(e, part, start, len, cancel, marks, split); if (part.len != @min(len - start, rows + 1)) return error.BadChunk; if (rows > fw.prefill_rows and !prompt_solo.allowed(e, start, rows)) return error.BadChunk; var one: [1]fw.Rank = undefined; diff --git a/zig/src/families/glm/cuda_prompt_pair.zig b/zig/src/families/glm/cuda_prompt_pair.zig index 0a1c7f9ef8..763bf5bc10 100644 --- a/zig/src/families/glm/cuda_prompt_pair.zig +++ b/zig/src/families/glm/cuda_prompt_pair.zig @@ -48,7 +48,7 @@ pub fn pairedAt(e: *const Engine, start: usize, len: usize, marks: ?engine.Marks } /// Engine.prefillChunk of two chunks (`part`: both, then the next id): one pass, then each chunk's MTP rows and commit. -pub fn prefill(e: *Engine, part: []const u32, start: usize, len: usize, cancel: ?engine.Cancel, marks: ?engine.Marks) !?u32 { +pub fn prefill(e: *Engine, part: []const u32, start: usize, len: usize, cancel: ?engine.Cancel, marks: ?engine.Marks, split: bool) !?u32 { const ra: usize = fw.prefill_rows; if (len <= start + ra) return error.BadChunk; const rb: usize = @min(len - start - ra, fw.prefill_rows); @@ -65,6 +65,8 @@ pub fn prefill(e: *Engine, part: []const u32, start: usize, len: usize, cancel: var fb = e.view(&pair.bp, &two); fa.prof = null; fb.prof = null; + fa.index_split = split; // the step's choice (TF_GLM_INDEX_SPLIT) for both chunks + fb.index_split = split; const st = &e.cur.st; try st.ready(e.stream, ra + rb); // one rebase for both chunks, before b's view copies the rings' base var sb = st.*; @@ -78,6 +80,7 @@ pub fn prefill(e: *Engine, part: []const u32, start: usize, len: usize, cancel: try cuda.copy_source.enqueue(e.d, pair.bp.ids, std.mem.sliceAsBytes(ids[ra .. ra + rb]), e.stream); try fw.Forward.pairBody(&fa, &fb, ra, rb); pair.steps += 1; + e.index_splits += @intFromBool(split); try mtpRows(e, &fa, ids[1 .. ra + 1], e.bp.fnormed); try fa.commit(ra); st.pending = rb; From 951ddd4ef2a0a3295480c25f967a705c7cee219a Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Tue, 27 Oct 2026 22:57:15 +0900 Subject: [PATCH 13/36] glm cuda: glm-split-check takes the pair's Leader.chunk, and a fifth prompt pairs a step past the split's position The check called Leader.chunk without the pair argument, so compile-cuda-tests failed on the stacked tree. It now steps by Chunk.rows, pairs where it can and counts the splits a step at a time; pairing leaves the first four prompts fewer split steps, so a fifth prompt of three chunks and min_rows pairs its second step past TF_GLM_INDEX_SPLIT. --- zig/tests/cuda/glm_split_tests.zig | 28 ++++++++++++++++------------ 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/zig/tests/cuda/glm_split_tests.zig b/zig/tests/cuda/glm_split_tests.zig index 4e0e90d45c..37c378c7f4 100644 --- a/zig/tests/cuda/glm_split_tests.zig +++ b/zig/tests/cuda/glm_split_tests.zig @@ -5,7 +5,7 @@ const cuda = @import("cuda"); const check = @import("check.zig"); const Gpu = check.Gpu; -/// `glm-split-check`: prompts a chunk at a time, split and not; first token, next window and every state byte equal. +/// `glm-split-check`: prompts a step at a time, split and not; first token, next window and every state byte equal. pub fn splitCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c_int, world: c_int, id_path: []const u8) !void { const glm = @import("glm_cuda"); var arena = std.heap.ArenaAllocator.init(gpu.gpa); @@ -36,10 +36,12 @@ pub fn splitCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c const x: glm.follow.Leader = .{ .e = e, .link = &link }; defer x.end() catch {}; const V = e.share.vocab.len(); - const from = C; // every chunk from the second on: each holds rows past the dense limit - // three chunks and a tail of 300, a tail of 2 (one rank without a row), a tail of 5, two chunks - const lens = [_]usize{ 3 * C + 300, 2 * C + 2, 2 * C + 5, 2 * C }; + const from = C; // every step from the second chunk on: each holds rows past the dense limit + // three chunks and a tail of 300, a tail of 2 (one rank without a row), a tail of 5, two chunks, and a second step + // that pairs past `from` (Leader.chunk pairs where it can; a paired step's two chunks take the split of its start) + const lens = [_]usize{ 3 * C + 300, 2 * C + 2, 2 * C + 5, 2 * C, 3 * C + glm.prompt_pair.min_rows }; var failed = false; + var paired: u64 = 0; // the split twins' paired steps for (lens, 0..) |len, i| { const ids = try a.alloc(u32, len); for (ids, 0..) |*u, j| u.* = @intCast(100 + (j * 7919 + i * 31) % 5000); @@ -47,20 +49,25 @@ pub fn splitCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c var outs: [2][4]u32 = undefined; var logits: [2][]u16 = undefined; var splits: [2]u64 = undefined; + var want: u64 = 0; // the steps glm.sparse.splitAt splits over this prompt for (0..2) |twin| { // twin 0 whole, twin 1 split from `from` e.settings.index_split = if (twin == 1) from else 0; try x.bind(@intCast(2 * i + twin)); try x.begin(null); const before = e.index_splits; + const steps = if (e.pair) |p| p.steps else 0; var at: usize = 0; - const t = while (true) : (at += C) { - const r = x.chunk(ids, at, false, null, null, null) catch |err| { - std.debug.print("FAIL split: prompt of {d}, twin {d}: the chunk at {d}: {s}\n", .{ len, twin, at, @errorName(err) }); + const t = while (true) { + if (twin == 1) want += @intFromBool(glm.sparse.splitAt(at, from, e.share.world)); + const r = x.chunk(ids, at, true, null, null, null) catch |err| { + std.debug.print("FAIL split: prompt of {d}, twin {d}: the step at {d}: {s}\n", .{ len, twin, at, @errorName(err) }); return err; }; if (r.first) |f| break f; + at += r.rows; }; splits[twin] = e.index_splits - before; + if (twin == 1) paired += (if (e.pair) |p| p.steps else 0) - steps; firsts[twin] = t; const win = [_]u32{ t, 101, 202, 303 }; try x.verify(&win, &outs[twin], 0); @@ -68,11 +75,8 @@ pub fn splitCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c try cuda.memory.download(gpu.d, e.bd.logits, std.mem.sliceAsBytes(logits[twin])[0 .. win.len * V * 2]); } e.settings.index_split = 0; - var want: u64 = 0; // the chunks glm.sparse.splitAt splits over this prompt - var start: usize = 0; - while (start < len) : (start += C) want += @intFromBool(glm.sparse.splitAt(start, from, e.share.world)); if (splits[0] != 0 or splits[1] != want) { - std.debug.print("FAIL split: prompt of {d}: {d} split chunks whole, {d} split (want {d})\n", .{ len, splits[0], splits[1], want }); + std.debug.print("FAIL split: prompt of {d}: {d} split steps whole, {d} split (want {d})\n", .{ len, splits[0], splits[1], want }); failed = true; } if (firsts[0] != firsts[1] or !std.mem.eql(u32, &outs[0], &outs[1]) or !std.mem.eql(u16, logits[0], logits[1])) { @@ -94,5 +98,5 @@ pub fn splitCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c } } if (failed) return error.TestFailed; - check.pass("split indexer rows at TP={d}: prompts of {any} rows equal their whole twins in the first token, the next window's logits and tokens, and every state byte", .{ world, lens }); + check.pass("split indexer rows at TP={d}: prompts of {any} rows ({d} paired steps) equal their whole twins in the first token, the next window's logits and tokens, and every state byte", .{ world, lens, paired }); } From a32a1a2f2b404c1844a5f90ef02cbdc3e0a78189 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Tue, 27 Oct 2026 23:50:24 +0900 Subject: [PATCH 14/36] glm cuda: a step pairs only where its index rings rebase as a chunk at a time's do prompt_pair.prefill readies the index rings once, at the pair's start. Where a chunk at a time rebases them only at the second chunk (from chunk 1 of rings based at 0: after a resume from a state kept at an odd chunk, or once decoders finish mid-prompt), the pair rebased at the first: the tokens agree, but the rings' bytes and base do not. pairedAt now pairs only where both rebase at the first chunk or neither rebases, judging a resumed step by its kept state's base (restored after the choice), so one unpaired chunk realigns a prompt; prefill refuses any other pair on every rank. glm-pair-check adds a prompt of four chunks with a mark at chunk 1, counts the paired steps with the rings' base followed on the host, and compares both twins' bases. --- zig/src/families/glm/cuda_follow.zig | 4 +-- zig/src/families/glm/cuda_prompt_pair.zig | 25 +++++++++++-- zig/tests/cuda/glm_pair_tests.zig | 44 +++++++++++++++-------- 3 files changed, 54 insertions(+), 19 deletions(-) diff --git a/zig/src/families/glm/cuda_follow.zig b/zig/src/families/glm/cuda_follow.zig index f97f6a13bc..942028b973 100644 --- a/zig/src/families/glm/cuda_follow.zig +++ b/zig/src/families/glm/cuda_follow.zig @@ -227,9 +227,9 @@ pub const Leader = struct { pub fn chunk(x: Leader, ids: []const u32, start: usize, alone: bool, from: ?engine.Resume, cancel: ?engine.Cancel, marks: ?engine.Marks) !Chunk { if (start >= ids.len) return error.BadChunk; try x.flush(); - const n: usize = if (alone and (prompt_solo.at(x.e, start, ids.len, from, marks) or x.e.pairedAt(start, ids.len, marks))) engine.solo_rows else engine.prompt_rows; + const n: usize = if (alone and (prompt_solo.at(x.e, start, ids.len, from, marks) or x.e.pairedAt(start, ids.len, from, marks))) engine.solo_rows else engine.prompt_rows; const part = ids[start..@min(ids.len, start + n + 1)]; - const split = sparse.splitAt(start, x.e.settings.index_split, x.e.share.world); // both chunks of a pair share it + const split = sparse.splitAt(start, x.e.settings.index_split, x.e.share.world); // a paired step's chunks share it try x.send(.{ .kind = @backingInt(Kind.chunk), .flags = x.rows() | @as(u32, @intFromBool(split)) << 1, .a = if (from) |r| r.id else 0, .seed = if (from) |r| r.at else 0, .depth = @intCast(start), .top_k = @intCast(ids.len) }, part); x.e.index_split = split; const first = try x.e.prefillChunk(part, start, ids.len, from, cancel, marks); diff --git a/zig/src/families/glm/cuda_prompt_pair.zig b/zig/src/families/glm/cuda_prompt_pair.zig index 763bf5bc10..e68367c73e 100644 --- a/zig/src/families/glm/cuda_prompt_pair.zig +++ b/zig/src/families/glm/cuda_prompt_pair.zig @@ -40,11 +40,20 @@ pub fn takesTwo(start: usize, len: usize, world: usize, marks: []const u32) bool return std.mem.indexOfScalar(u32, marks, @intCast(next)) == null; } +/// Whether a pair at `start` (b of `rb` rows) rebases the index rings where chunks one at a time do: at a, or not. +pub fn ringsAgree(base: usize, ring: usize, start: usize, rb: usize) bool { + const ra = fw.prefill_rows; + return start + ra > base + ring or start + ra + rb <= base + ring; +} + /// Rank 0's choice for the step at `start` (the others take the ids it sends); not with a drafter, prompt rows or prof. -pub fn pairedAt(e: *const Engine, start: usize, len: usize, marks: ?engine.Marks) bool { +pub fn pairedAt(e: *const Engine, start: usize, len: usize, from: ?engine.Resume, marks: ?engine.Marks) bool { if (e.pair == null or !e.settings.pair or e.drafter != null or e.prompt_logprobs != null or e.prompt_profs[3] != null) return false; if (e.pieces.count == 1 and !e.settings.split) return false; // run takes such a chunk whole - return takesTwo(start, len, e.share.world, if (marks) |m| m.at else &.{}); + if (!takesTwo(start, len, e.share.world, if (marks) |m| m.at else &.{})) return false; + // a resumed step restores its kept state, and the rings' base, after this choice + const base = if (from) |r| (e.snaps.get(r.id) orelse return false).bases[0] else e.cur.st.base; + return ringsAgree(base, e.cur.st.ring_rows, start, @min(len - start - fw.prefill_rows, fw.prefill_rows)); } /// Engine.prefillChunk of two chunks (`part`: both, then the next id): one pass, then each chunk's MTP rows and commit. @@ -54,6 +63,7 @@ pub fn prefill(e: *Engine, part: []const u32, start: usize, len: usize, cancel: const rb: usize = @min(len - start - ra, fw.prefill_rows); const pair = if (e.pair) |*p| p else return error.BadChunk; if (rb < e.share.world or part.len != @min(len - start, ra + rb + 1) or e.drafter != null or e.prompt_logprobs != null) return error.BadChunk; + if (!ringsAgree(e.cur.st.base, e.cur.st.ring_rows, start, rb)) return error.BadChunk; // pairedAt's rule, every rank if (marks) |m| if (start > 0 and std.mem.indexOfScalar(u32, m.at, @intCast(start)) != null) m.call(m.ctx, @intCast(start)); e.gate(cancel) catch |err| { if (err == error.Cancelled) try snapshots.interrupted(e, marks, @intCast(start)); @@ -68,7 +78,7 @@ pub fn prefill(e: *Engine, part: []const u32, start: usize, len: usize, cancel: fa.index_split = split; // the step's choice (TF_GLM_INDEX_SPLIT) for both chunks fb.index_split = split; const st = &e.cur.st; - try st.ready(e.stream, ra + rb); // one rebase for both chunks, before b's view copies the rings' base + try st.ready(e.stream, ra + rb); // one rebase for both chunks, at a as a chunk at a time, before b's view copies it var sb = st.*; sb.pos_dev = pair.pos.ptr; try sb.setPos(e.stream, st.pos + ra); @@ -115,3 +125,12 @@ test "a prompt step takes the next chunk when it holds min_rows or more and no k try std.testing.expect(!takesTwo(0, 3 * C, 3, &.{C})); // a state kept at the next chunk's start: that chunk starts a step try std.testing.expect(takesTwo(0, 3 * C, 3, &.{ 2 * C, 3 * C })); // marks past the pair } + +test "a pair rebases the index rings where a chunk at a time would, so a ring based at chunk 0 pairs from chunk 2" { + const C = fw.prefill_rows; + try std.testing.expect(ringsAgree(0, 2 * C, 0, C)); // neither rebases + try std.testing.expect(!ringsAgree(0, 2 * C, C, C)); // a chunk at a time rebases at b's start, a pair at a's + try std.testing.expect(!ringsAgree(0, 2 * C, C, 1)); + try std.testing.expect(ringsAgree(0, 2 * C, 2 * C, C)); // both at a's start + try std.testing.expect(ringsAgree(C, 2 * C, C, C)); +} diff --git a/zig/tests/cuda/glm_pair_tests.zig b/zig/tests/cuda/glm_pair_tests.zig index ff03559ed8..e339df1626 100644 --- a/zig/tests/cuda/glm_pair_tests.zig +++ b/zig/tests/cuda/glm_pair_tests.zig @@ -23,7 +23,7 @@ pub fn pairCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c_ }; try check.expect(set.pair, "glm-pair-check needs TF_GLM_PROMPT_PAIR on (unset or 1)", .{}); const C = glm.engine.prompt_rows; - const capacity = @max(dc.capacity, 4 * C); + const capacity = @max(dc.capacity, 5 * C); const e = try glm.engine.Engine.init(gpu.gpa, gpu.io, gpu.d, gpu.ctx, checkpoint, &lib, comm, @intCast(rank), @intCast(world), .{ .capacity = capacity, .mtp = true, .settings = set }); defer e.deinit(); var link = try glm.follow.Link.init(gpu.gpa, gpu.d, &lib, comm, e.stream, capacity); @@ -37,24 +37,37 @@ pub fn pairCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c_ defer x.end() catch {}; const V = e.share.vocab.len(); const min = glm.prompt_pair.min_rows; - // a pair and a tail, the shortest second chunk, a pair and a chunk, a pair to the end, a tail too short to pair - const lens = [_]usize{ 2 * C + 300, C + min, 3 * C, 2 * C, C + min - 1 }; + const Case = struct { len: usize, marks: []const u32 = &.{} }; + const cases = [_]Case{ + .{ .len = 2 * C + 300 }, // a pair and a tail + .{ .len = C + min }, // the shortest second chunk + .{ .len = 3 * C }, // a pair and a chunk + .{ .len = 2 * C }, // a pair to the end + .{ .len = C + min - 1 }, // a tail too short to pair + .{ .len = 4 * C, .marks = &.{C} }, // a mark at chunk 1: 0 and 1 alone (a pair at 1 rebases apart), then a pair + }; + const Kept = struct { + fn call(_: *anyopaque, _: u32) void {} // the check keeps no state: its marks only shape the steps + }; + var no_ctx: u8 = 0; var failed = false; - for (lens, 0..) |len, i| { + for (cases, 0..) |case, i| { + const len = case.len; + const marks: ?glm.engine.Marks = if (case.marks.len == 0) null else .{ .at = case.marks, .ctx = &no_ctx, .call = Kept.call }; const ids = try a.alloc(u32, len); for (ids, 0..) |*u, j| u.* = @intCast(100 + (j * 7919 + i * 31) % 5000); var firsts: [2]u32 = undefined; var outs: [2][4]u32 = undefined; var logits: [2][]u16 = undefined; var paired: [2]u64 = undefined; - for (0..2) |twin| { // twin 0 a chunk a step, twin 1 paired where it can + for (0..2) |twin| { // twin 0 a chunk a step, twin 1 paired where it can, both with the same marks e.settings.pair = twin == 1; try x.bind(@intCast(2 * i + twin)); try x.begin(null); const before = e.pair.?.steps; var at: usize = 0; const t = while (true) { - const r = x.chunk(ids, at, true, null, null, null) catch |err| { + const r = x.chunk(ids, at, true, null, null, marks) catch |err| { std.debug.print("FAIL pair: prompt of {d}, twin {d}: the step at {d}: {s}\n", .{ len, twin, at, @errorName(err) }); return err; }; @@ -69,12 +82,17 @@ pub fn pairCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c_ try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(logits[twin].ptr, e.bd.logits, win.len * V * 2), "cuMemcpyDtoH"); } e.settings.pair = true; - var want: u64 = 0; // the steps glm.prompt_pair.takesTwo pairs over this prompt + const st = &e.seqs.items[2 * i].?.st; + const tw = &e.seqs.items[2 * i + 1].?.st; + var want: u64 = 0; // the steps glm.prompt_pair.pairedAt pairs over this prompt, the rings' base followed here + var base: usize = 0; var at: usize = 0; while (at < len) { - const two = glm.prompt_pair.takesTwo(at, len, e.share.world, &.{}); + const two = glm.prompt_pair.takesTwo(at, len, e.share.world, case.marks) and glm.prompt_pair.ringsAgree(base, st.ring_rows, at, @min(len - at - C, C)); + const rows = @min(len - at, @as(usize, if (two) 2 * C else C)); + if (at + rows > base + st.ring_rows) base = at; // State.ready rebases at the step's start, a whole pool want += @intFromBool(two); - at += if (two) 2 * C else C; + at += rows; } if (paired[0] != 0 or paired[1] != want) { std.debug.print("FAIL pair: prompt of {d}: {d} paired steps a chunk a step, {d} paired (want {d})\n", .{ len, paired[0], paired[1], want }); @@ -85,19 +103,17 @@ pub fn pairCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c_ failed = true; } try e.stream.synchronize(); - const st = &e.seqs.items[2 * i].?.st; - const tw = &e.seqs.items[2 * i + 1].?.st; const n = st.arena.buf.len; const h0 = try a.alloc(u8, n); const h1 = try a.alloc(u8, n); try st.arena.buf.download(0, h0); try tw.arena.buf.download(0, h1); const diff = std.mem.indexOfDiff(u8, h0, h1); - if (st.pos != tw.pos or st.mtp_len != tw.mtp_len or !std.mem.eql(u1, st.cur, tw.cur) or diff != null) { - std.debug.print("FAIL pair: prompt of {d}: the paired state differs (pos {d}/{d}, first byte {?d} of {d})\n", .{ len, st.pos, tw.pos, diff, n }); + if (st.pos != tw.pos or st.mtp_len != tw.mtp_len or st.base != tw.base or st.mtp_base != tw.mtp_base or !std.mem.eql(u1, st.cur, tw.cur) or diff != null) { + std.debug.print("FAIL pair: prompt of {d}: the paired state differs (pos {d}/{d}, rings' base {d}/{d}, MTP base {d}/{d}, first byte {?d} of {d})\n", .{ len, st.pos, tw.pos, st.base, tw.base, st.mtp_base, tw.mtp_base, diff, n }); failed = true; } } if (failed) return error.TestFailed; - check.pass("paired prompt steps at TP={d}: prompts of {any} rows equal their twins a chunk a step in the first token, the next window's logits and tokens, and every state byte", .{ world, lens }); + check.pass("paired prompt steps at TP={d}: {d} prompts, one with a mark at chunk 1, equal their twins a chunk a step in the first token, the next window's logits and tokens, the rings' bases and every state byte", .{ world, cases.len }); } From 7649d99f4096565839b1adaa9c14f30350e0ed96 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Tue, 27 Oct 2026 23:39:21 +0900 Subject: [PATCH 15/36] glm cuda: only an engine with other ranks makes the pair's twin At TP=1 a pair has no exchanges to hide and no receipt, yet the twin was allocated and the side stream ran at top priority. prompt_pair.wanted(on, world) is now on and world > 1, for both; pairedAt already pairs only with a twin. glm-pair-check needs other ranks. --- zig/src/families/glm/cuda_engine.zig | 5 +++-- zig/src/families/glm/cuda_prompt_pair.zig | 11 +++++++++++ zig/tests/cuda/glm_pair_tests.zig | 1 + 3 files changed, 15 insertions(+), 2 deletions(-) diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index 186d61a724..66939b099a 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/zig/src/families/glm/cuda_engine.zig @@ -148,7 +148,8 @@ pub const Engine = struct { const block32 = !exl3.isExl3(gpa, text) and try @import("cuda_experts.zig").wantsBlock32(gpa, text); var m = try fw.Modules.load(d, ctx, block32); errdefer m.deinit(); - var pieces = try fw.Pieces.init(d, o.settings.pieces, o.settings.pair); + const pair = prompt_pair.wanted(o.settings.pair, world); + var pieces = try fw.Pieces.init(d, o.settings.pieces, pair); errdefer pieces.deinit(); if (o.settings.layers > c.layers) return error.BadLayers; const layers = try a.alloc(u32, if (o.settings.layers > 0) o.settings.layers else c.layers); @@ -190,7 +191,7 @@ pub const Engine = struct { errdefer e.bp.deinit(); if (o.settings.solo and rank == 0) std.log.info("glm: prompt passes of {d} rows, {d} while another stream decodes", .{ fw.solo_rows, fw.prefill_rows }); if (world > 1 and rank == 0) std.log.info("glm: a prompt chunk's exchanges in {d} row pieces", .{o.settings.pieces}); - if (o.settings.pair) e.pair = try prompt_pair.Pair.init(d, &e.w, &e.bp, &e.pieces, o.capacity); + if (pair) e.pair = try prompt_pair.Pair.init(d, &e.w, &e.bp, &e.pieces, o.capacity); errdefer if (e.pair) |*p| p.deinit(); e.bd = try fw.Buffers.init(d, &e.w, fw.window_rows, false, o.capacity); errdefer e.bd.deinit(); diff --git a/zig/src/families/glm/cuda_prompt_pair.zig b/zig/src/families/glm/cuda_prompt_pair.zig index e68367c73e..f6ce94ba55 100644 --- a/zig/src/families/glm/cuda_prompt_pair.zig +++ b/zig/src/families/glm/cuda_prompt_pair.zig @@ -33,6 +33,11 @@ pub const Pair = struct { } }; +/// Whether an engine makes the pair's twin: TF_GLM_PROMPT_PAIR on, and other ranks whose exchanges a pair hides. +pub fn wanted(on: bool, world: usize) bool { + return on and world > 1; +} + /// Whether a step at `start` of a `len`-id prompt takes the next chunk too: min_rows or more there, and no mark at it. pub fn takesTwo(start: usize, len: usize, world: usize, marks: []const u32) bool { const next = start + fw.prefill_rows; @@ -134,3 +139,9 @@ test "a pair rebases the index rings where a chunk at a time would, so a ring ba try std.testing.expect(ringsAgree(0, 2 * C, 2 * C, C)); // both at a's start try std.testing.expect(ringsAgree(C, 2 * C, C, C)); } + +test "only ranks that exchange make the twin: TP=1 has no exchanges to hide" { + try std.testing.expect(!wanted(true, 1)); + try std.testing.expect(wanted(true, 2)); + try std.testing.expect(!wanted(false, 3)); +} diff --git a/zig/tests/cuda/glm_pair_tests.zig b/zig/tests/cuda/glm_pair_tests.zig index e339df1626..633e1d9c6a 100644 --- a/zig/tests/cuda/glm_pair_tests.zig +++ b/zig/tests/cuda/glm_pair_tests.zig @@ -22,6 +22,7 @@ pub fn pairCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c_ return err; }; try check.expect(set.pair, "glm-pair-check needs TF_GLM_PROMPT_PAIR on (unset or 1)", .{}); + try check.expect(world > 1, "glm-pair-check needs other ranks: at TP=1 an engine makes no twin", .{}); const C = glm.engine.prompt_rows; const capacity = @max(dc.capacity, 5 * C); const e = try glm.engine.Engine.init(gpu.gpa, gpu.io, gpu.d, gpu.ctx, checkpoint, &lib, comm, @intCast(rank), @intCast(world), .{ .capacity = capacity, .mtp = true, .settings = set }); From 2228fd3cbea7d0d6a90d29e565e2ac8bb489e0b7 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Tue, 27 Oct 2026 23:39:21 +0900 Subject: [PATCH 16/36] glm cuda: one table gives the twin's fields their size indexes, and a size added or taken away stops the build twin_sizes tied the fields a twin holds apart to make's size indexes by hand, beside a list of their names. twin_at now maps each field to its index: make reads those fields from it, the twin's allocation and copy take it too, and make's sizes are an array of twin_checked entries, so a future size entry fails the build until the indexes are checked again. --- zig/src/families/glm/cuda_state.zig | 45 +++++++++++++++++------------ 1 file changed, 26 insertions(+), 19 deletions(-) diff --git a/zig/src/families/glm/cuda_state.zig b/zig/src/families/glm/cuda_state.zig index da791c2f24..0a5e6bf058 100644 --- a/zig/src/families/glm/cuda_state.zig +++ b/zig/src/families/glm/cuda_state.zig @@ -190,9 +190,16 @@ pub const Buffers = struct { return make(d, w, of.rows, of.prefill, capacity, of.rows, of); } - /// The fields a twin holds apart: make's sizes by index, then their names. - const twin_sizes = [_]usize{ 0, 1, 2, 3, 4, 5, 6, 7, 62, 63, 64, 65, 70 }; - const twin_fields = .{ "ids", "x", "normed", "xs", "post", "comb", "hcpart", "gathered", "hidden", "fnormed", "fxs", "logits", "part" }; + /// The fields a twin holds apart, by their index in make's sizes (make reads the fields from these indexes too). + const twin_at = .{ .ids = 0, .x = 1, .normed = 2, .xs = 3, .post = 4, .comb = 5, .hcpart = 6, .gathered = 7, .hidden = 62, .fnormed = 63, .fxs = 64, .logits = 65, .part = 70 }; + /// make's size count when twin_at's indexes were last checked: adding or removing a size fails the build. + const twin_checked = 76; + + /// Whether make's size `i` is one a twin holds apart. + fn twinHolds(i: usize) bool { + inline for (@typeInfo(@TypeOf(twin_at)).@"struct".field_names) |f| if (@field(twin_at, f) == i) return true; + return false; + } /// `seg`: the most rows one stream's attention takes (a decode window's rows, or a shared window's segment). fn make(d: *const cuda.Driver, w: *const weights.Weights, rows: usize, prefill: bool, capacity: usize, seg: usize, twin: ?*const Buffers) !Buffers { @@ -222,7 +229,7 @@ pub const Buffers = struct { var split = false; // a KDA stack in two parts moves each part's rows through one more buffer for (w.layers) |l| split = split or (l.mixer == .kda and l.mixer.kda.proj == .split); const hist_len = if (pairs > cuda.grouped.small) (pairs + cuda.grouped.small - 1) / cuda.grouped.small * (count + 1) else 1; - const sizes = [_]usize{ + const sizes: [twin_checked]usize = .{ rows * 4, rows * S * D * 2, rows * D * 2, rows * (D / 64) * 4, rows * S * 4, rows * S * S * 4, glue.hcPartBytes(rows), s.world * rows * D * 4, if (prefill or seg < rows) rows * kw * 2 else 4, rows * n * 2, rows * n * 2, // kproj, ka, kg @@ -263,7 +270,7 @@ pub const Buffers = struct { }; var held = sizes; if (twin != null) for (&held, 0..) |*bytes, i| { // a twin allocates its own rows only - if (std.mem.indexOfScalar(usize, &twin_sizes, i) == null) bytes.* = 0; + if (!twinHolds(i)) bytes.* = 0; }; var p: [sizes.len]u64 = undefined; const a = try arena(d, &held, &p); @@ -273,14 +280,14 @@ pub const Buffers = struct { .prefill = prefill, .part_rows = part_rows, .nch = nch, - .ids = p[0], - .x = p[1], - .normed = p[2], - .xs = p[3], - .post = p[4], - .comb = p[5], - .hcpart = p[6], - .gathered = p[7], + .ids = p[twin_at.ids], + .x = p[twin_at.x], + .normed = p[twin_at.normed], + .xs = p[twin_at.xs], + .post = p[twin_at.post], + .comb = p[twin_at.comb], + .hcpart = p[twin_at.hcpart], + .gathered = p[twin_at.gathered], .kproj = p[8], .kda = .{ .p = 0, .ka = p[9], .kg = p[10], .saves = .{ .out = p[11], .state_out = 0, .k = p[12], .v = p[13], .g = p[14], .b = p[15] }, .wide = .{ .q = p[16], .y = p[17] }, .mm_part = p[42] }, .dsa = .{ @@ -307,22 +314,22 @@ pub const Buffers = struct { }, .dense = .{ .tmp = p[44], .gu = p[45], .act = p[46], .xs_act = p[47], .mm_part = p[42] }, .moe = .{ .mlog = p[48], .router_part = p[26], .pick = p[49], .wts = p[50], .plan = .{ .members = p[51], .items = p[52], .counts = p[53], .rank = p[54], .hist = p[55] }, .eact = p[56], .ey = p[57], .sgu = p[58], .sact = p[59], .sxs = p[60], .sy = p[61], .mm_part = p[42] }, - .hidden = p[62], - .fnormed = p[63], - .fxs = p[64], - .logits = p[65], + .hidden = p[twin_at.hidden], + .fnormed = p[twin_at.fnormed], + .fxs = p[twin_at.fxs], + .logits = p[twin_at.logits], .hin = p[66], .me = p[67], .mcat = p[68], .mx = p[69], - .part = p[70], + .part = p[twin_at.part], }; out.moe.x3 = .{ .picks = p[71], .xg = p[72], .xu = p[73], .z = p[74] }; out.kda.tmp = p[75]; if (twin) |of| { var t = of.*; t.arena = out.arena; - inline for (twin_fields) |f| @field(t, f) = @field(out, f); + inline for (@typeInfo(@TypeOf(twin_at)).@"struct".field_names) |f| @field(t, f) = @field(out, f); return t; } return out; From 34cd3c096a0678365aace4e87e150871fddd5362 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Wed, 28 Oct 2026 07:47:40 +0900 Subject: [PATCH 17/36] glm cuda: glm-split-check's list of prompts is described on one line --- zig/tests/cuda/glm_split_tests.zig | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/zig/tests/cuda/glm_split_tests.zig b/zig/tests/cuda/glm_split_tests.zig index 37c378c7f4..a2dcf5bf73 100644 --- a/zig/tests/cuda/glm_split_tests.zig +++ b/zig/tests/cuda/glm_split_tests.zig @@ -37,8 +37,7 @@ pub fn splitCheck(gpu: Gpu, checkpoint: []const u8, dc_path: []const u8, rank: c defer x.end() catch {}; const V = e.share.vocab.len(); const from = C; // every step from the second chunk on: each holds rows past the dense limit - // three chunks and a tail of 300, a tail of 2 (one rank without a row), a tail of 5, two chunks, and a second step - // that pairs past `from` (Leader.chunk pairs where it can; a paired step's two chunks take the split of its start) + // tails of 300, 2 (one rank without a row) and 5, two whole chunks, and a second step paired past `from` const lens = [_]usize{ 3 * C + 300, 2 * C + 2, 2 * C + 5, 2 * C, 3 * C + glm.prompt_pair.min_rows }; var failed = false; var paired: u64 = 0; // the split twins' paired steps From 8cf53e4308d9e693805d43aa3edc8a382d2ebf0c Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Sun, 1 Nov 2026 09:37:03 +0900 Subject: [PATCH 18/36] glm cuda: the pair's choice and its twin's pieces sit with the pair and Cancel with the engine's types, so the engine and the lanes keep room under the line limit --- zig/src/families/glm/cuda_engine.zig | 8 ++------ zig/src/families/glm/cuda_engine_types.zig | 3 +++ zig/src/families/glm/cuda_follow.zig | 3 ++- zig/src/families/glm/cuda_prompt_pair.zig | 7 +++++++ 4 files changed, 14 insertions(+), 7 deletions(-) diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index 66939b099a..7c0b0af21b 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/zig/src/families/glm/cuda_engine.zig @@ -22,9 +22,6 @@ const l2pf_mod = @import("cuda_l2pf.zig"); const prompt_pair = @import("cuda_prompt_pair.zig"); pub const Options = @import("cuda_engine_types.zig").Options; -/// Rank 0's own word on its request (cancelled or not); the other ranks have none and hear rank 0's. -pub const Cancel = struct { ptr: ?*anyopaque, check: *const fn (?*anyopaque) bool }; - /// The most drafts a chain holds (decode.DepthPolicy's top depth). pub const max_drafts = 3; /// A greedy pick as it crosses ranks: (value, id) and padding, so the RoCE gather (16-byte packs) can carry it. @@ -48,6 +45,7 @@ pub const Kept = types.Kept; const Seq = types.Seq; const Snap = types.Snap; pub const Resume = types.Resume; +pub const Cancel = types.Cancel; pub const Marks = types.Marks; pub const Engine = struct { @@ -264,8 +262,7 @@ pub const Engine = struct { pub fn view(e: *Engine, b: *fw.Buffers, out: *[1]fw.Rank) fw.Forward { out.* = .{.{ .w = &e.w, .st = &e.cur.st, .b = b }}; const im: ?fw.ImageRows = if (e.images) |v| .{ .v = v, .slot = e.slot } else null; - const twin = if (e.pair) |*p| (if (b == &p.bp) &p.pieces else null) else null; - return .{ .d = e.d, .m = &e.m, .stream = e.stream, .ranks = out, .exchange = e.exchange(), .pieces = if (b == &e.bp) &e.pieces else twin, .split = e.settings.split, .ring = e.settings.ring, .kda_wide = e.settings.kda_wide, .images = im, .taps = if (e.drafter) |dr| dr.taps(b == &e.bp) else null, .l2pf = if (e.l2pf) |*p| p else null }; + return .{ .d = e.d, .m = &e.m, .stream = e.stream, .ranks = out, .exchange = e.exchange(), .pieces = prompt_pair.piecesOf(e, b), .split = e.settings.split, .ring = e.settings.ring, .kda_wide = e.settings.kda_wide, .images = im, .taps = if (e.drafter) |dr| dr.taps(b == &e.bp) else null, .l2pf = if (e.l2pf) |*p| p else null }; } /// Binds slot `slot`'s sequence for the steps that follow, making it (a zeroed state) on first use. @@ -531,7 +528,6 @@ pub const Engine = struct { } pub const holdPrompt = @import("cuda_prompt_hold.zig").save; - pub const pairedAt = prompt_pair.pairedAt; pub const prefillMany = @import("cuda_prompts.zig").prefillMany; /// forward.commit of the last window's first `n` rows (the KDA replay when fewer than it ran). diff --git a/zig/src/families/glm/cuda_engine_types.zig b/zig/src/families/glm/cuda_engine_types.zig index dded605953..125a6ee1ed 100644 --- a/zig/src/families/glm/cuda_engine_types.zig +++ b/zig/src/families/glm/cuda_engine_types.zig @@ -3,6 +3,9 @@ const cuda = @import("cuda"); const fw = @import("cuda_forward.zig"); const sampler = @import("cuda_sampler.zig"); +/// Rank 0's own word on its request (cancelled or not); the other ranks have none and hear rank 0's. +pub const Cancel = struct { ptr: ?*anyopaque, check: *const fn (?*anyopaque) bool }; + /// One stream's window of a shared verify: its slot, its ids, and its drawn tokens. pub const Part = struct { slot: u32, ids: []const u32, out: []u32, bits: u64 = 0 }; diff --git a/zig/src/families/glm/cuda_follow.zig b/zig/src/families/glm/cuda_follow.zig index 942028b973..15c74b3846 100644 --- a/zig/src/families/glm/cuda_follow.zig +++ b/zig/src/families/glm/cuda_follow.zig @@ -11,6 +11,7 @@ const prompt_wire = @import("cuda_prompt_wire.zig"); const dflash = @import("dflash.zig"); const pages_mod = @import("cuda_pages.zig"); const sparse = @import("cuda_sparse.zig"); +const prompt_pair = @import("cuda_prompt_pair.zig"); const Engine = engine.Engine; pub const Kind = enum(u32) { begin = 1, prefill, verify, keep, draft, end, drop, release, chunk, shared, image, block, many, redraw, idle, snapshots, pool }; @@ -227,7 +228,7 @@ pub const Leader = struct { pub fn chunk(x: Leader, ids: []const u32, start: usize, alone: bool, from: ?engine.Resume, cancel: ?engine.Cancel, marks: ?engine.Marks) !Chunk { if (start >= ids.len) return error.BadChunk; try x.flush(); - const n: usize = if (alone and (prompt_solo.at(x.e, start, ids.len, from, marks) or x.e.pairedAt(start, ids.len, from, marks))) engine.solo_rows else engine.prompt_rows; + const n: usize = if (alone and (prompt_solo.at(x.e, start, ids.len, from, marks) or prompt_pair.pairedAt(x.e, start, ids.len, from, marks))) engine.solo_rows else engine.prompt_rows; const part = ids[start..@min(ids.len, start + n + 1)]; const split = sparse.splitAt(start, x.e.settings.index_split, x.e.share.world); // a paired step's chunks share it try x.send(.{ .kind = @backingInt(Kind.chunk), .flags = x.rows() | @as(u32, @intFromBool(split)) << 1, .a = if (from) |r| r.id else 0, .seed = if (from) |r| r.at else 0, .depth = @intCast(start), .top_k = @intCast(ids.len) }, part); diff --git a/zig/src/families/glm/cuda_prompt_pair.zig b/zig/src/families/glm/cuda_prompt_pair.zig index f6ce94ba55..d35f77b5e8 100644 --- a/zig/src/families/glm/cuda_prompt_pair.zig +++ b/zig/src/families/glm/cuda_prompt_pair.zig @@ -51,6 +51,13 @@ pub fn ringsAgree(base: usize, ring: usize, start: usize, rb: usize) bool { return start + ra > base + ring or start + ra + rb <= base + ring; } +/// The pieces a forward over `b` exchanges in: the engine's, or the pair's for its twin; none for other buffers. +pub fn piecesOf(e: *const Engine, b: *const fw.Buffers) ?*const fw.Pieces { + if (b == &e.bp) return &e.pieces; + const p = if (e.pair) |*p| p else return null; + return if (b == &p.bp) &p.pieces else null; +} + /// Rank 0's choice for the step at `start` (the others take the ids it sends); not with a drafter, prompt rows or prof. pub fn pairedAt(e: *const Engine, start: usize, len: usize, from: ?engine.Resume, marks: ?engine.Marks) bool { if (e.pair == null or !e.settings.pair or e.drafter != null or e.prompt_logprobs != null or e.prompt_profs[3] != null) return false; From 7aa33609cf24a8106b1e8bdd6c21389a5fcde805 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Sun, 1 Nov 2026 20:19:00 +0900 Subject: [PATCH 19/36] glm cuda: TF_GLM_CHUNK_SOLO=4096 takes the pair's place, and its server holds no twin A solo pass and a paired step both reach the other ranks as a part of two chunks, so one server cannot mix them. With the wider buffers set, the pair stands aside and makes no twin. --- docs/recipes/glm-5.3-flash-cuda.md | 3 ++- zig/src/families/glm/cuda_engine.zig | 2 +- zig/src/families/glm/cuda_follow.zig | 2 +- zig/src/families/glm/cuda_lanes_test.zig | 3 +-- zig/src/families/glm/cuda_prompt_pair.zig | 17 +++++++++++------ 5 files changed, 16 insertions(+), 11 deletions(-) diff --git a/docs/recipes/glm-5.3-flash-cuda.md b/docs/recipes/glm-5.3-flash-cuda.md index f37e540de6..2848a34c7f 100644 --- a/docs/recipes/glm-5.3-flash-cuda.md +++ b/docs/recipes/glm-5.3-flash-cuda.md @@ -100,7 +100,8 @@ for each pass to end. A pass that takes more rows reads each routed expert's wei comes sooner. `TF_GLM_CHUNK_SOLO=4096` lets a pass take two of those chunks when no other stream is decoding as it starts. With one decoding, or with `prompt_logprobs` or image rows, the pass takes 2,048. Beside `--drafter`, DFlash2 reads a solo pass's taps whole and keeps the prompt's last 2,048 rows as before; its tap buffer holds 4,096 rows, 80 MiB -more on each rank. +more on each rank. Solo passes take the place of `TF_GLM_PROMPT_PAIR`'s paired steps, whose second set of buffers they +do not hold. The grid does not move, so warm turns resume where they do unset, and streams decoding beside a prompt see the same passes. Replies are the same bytes either way. Two costs remain: diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index 7c0b0af21b..2b8ad93f77 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/zig/src/families/glm/cuda_engine.zig @@ -146,7 +146,7 @@ pub const Engine = struct { const block32 = !exl3.isExl3(gpa, text) and try @import("cuda_experts.zig").wantsBlock32(gpa, text); var m = try fw.Modules.load(d, ctx, block32); errdefer m.deinit(); - const pair = prompt_pair.wanted(o.settings.pair, world); + const pair = prompt_pair.wanted(o.settings.pair, o.settings.solo, world); var pieces = try fw.Pieces.init(d, o.settings.pieces, pair); errdefer pieces.deinit(); if (o.settings.layers > c.layers) return error.BadLayers; diff --git a/zig/src/families/glm/cuda_follow.zig b/zig/src/families/glm/cuda_follow.zig index 15c74b3846..ef0715e288 100644 --- a/zig/src/families/glm/cuda_follow.zig +++ b/zig/src/families/glm/cuda_follow.zig @@ -224,7 +224,7 @@ pub const Leader = struct { return x.e.prefill(ids, from, cancel, marks); } - /// One prompt step on every rank (Engine.prefillChunk): a chunk, or two if `alone` and prompt_solo.at or the pair's rule allow. + /// One prompt step on every rank (Engine.prefillChunk): a chunk, or two if `alone` and solo or pair rules allow. pub fn chunk(x: Leader, ids: []const u32, start: usize, alone: bool, from: ?engine.Resume, cancel: ?engine.Cancel, marks: ?engine.Marks) !Chunk { if (start >= ids.len) return error.BadChunk; try x.flush(); diff --git a/zig/src/families/glm/cuda_lanes_test.zig b/zig/src/families/glm/cuda_lanes_test.zig index 1eca112ac8..af503aff49 100644 --- a/zig/src/families/glm/cuda_lanes_test.zig +++ b/zig/src/families/glm/cuda_lanes_test.zig @@ -1,8 +1,7 @@ -//! The lane backend's slots, chosen and given back, and when a prompt step may pair. +//! The lane backend's slots: chosen before a stream's map entry exists, and none left behind by a failed insertion. const std = @import("std"); const lanes = @import("lanes"); const engine = @import("cuda_engine.zig"); -const fw = @import("cuda_forward.zig"); const Lanes = @import("cuda_lanes.zig").Lanes; const fw = @import("cuda_forward.zig"); const solo = @import("cuda_prompt_solo.zig"); diff --git a/zig/src/families/glm/cuda_prompt_pair.zig b/zig/src/families/glm/cuda_prompt_pair.zig index d35f77b5e8..c864e93338 100644 --- a/zig/src/families/glm/cuda_prompt_pair.zig +++ b/zig/src/families/glm/cuda_prompt_pair.zig @@ -33,9 +33,9 @@ pub const Pair = struct { } }; -/// Whether an engine makes the pair's twin: TF_GLM_PROMPT_PAIR on, and other ranks whose exchanges a pair hides. -pub fn wanted(on: bool, world: usize) bool { - return on and world > 1; +/// Whether an engine makes the pair's twin: TF_GLM_PROMPT_PAIR on, no solo passes, and ranks to hide exchanges from. +pub fn wanted(on: bool, solo: bool, world: usize) bool { + return on and !solo and world > 1; } /// Whether a step at `start` of a `len`-id prompt takes the next chunk too: min_rows or more there, and no mark at it. @@ -148,7 +148,12 @@ test "a pair rebases the index rings where a chunk at a time would, so a ring ba } test "only ranks that exchange make the twin: TP=1 has no exchanges to hide" { - try std.testing.expect(!wanted(true, 1)); - try std.testing.expect(wanted(true, 2)); - try std.testing.expect(!wanted(false, 3)); + try std.testing.expect(!wanted(true, false, 1)); + try std.testing.expect(wanted(true, false, 2)); + try std.testing.expect(!wanted(false, false, 3)); +} + +test "a server with solo passes holds no twin: they take the pair's steps" { + try std.testing.expect(!wanted(true, true, 2)); + try std.testing.expect(!wanted(true, true, 3)); } From 6d69c5dde138c9a3ca032a251cb8cf82ebd29aa5 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Sat, 7 Nov 2026 06:00:04 +0900 Subject: [PATCH 20/36] docs: the ledger's item 20 lists GLM's paired prompt steps, about 2 % sooner at TP=2 --- docs/optimization-ledger.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/optimization-ledger.md b/docs/optimization-ledger.md index 808c99bcb6..ddf3cf2873 100644 --- a/docs/optimization-ledger.md +++ b/docs/optimization-ledger.md @@ -371,6 +371,7 @@ below; these are estimates for planning a port, not results. | glm (C) | cold first token −6..−7 % at TP=2, −7..−8 % at TP=3 | [#1823] | | flashnext (M) | 1,024-token reply beside a 64K prompt 48 → 8 s, prompt's first token 90 → 54 s | [#2167] | | flashnext (C), one GPU | prompts 15–23 % faster at `--parallel 1`, up to 8,192 rows a pass | [#2246] | + | glm (C), TP=2 | prompts about 2 % sooner at 8K–128K, two chunks a step, replies unchanged | [#1423] | `TF_GLM_CHUNK_SOLO=4096` is slower with GLM's NVFP4 checkpoint on two GB10s ([#2288]). @@ -460,6 +461,7 @@ below; these are estimates for planning a port, not results. [#1394]: https://github.com/ashhart/TensorFold/pull/1394 [#1411]: https://github.com/ashhart/TensorFold/pull/1411 [#1413]: https://github.com/ashhart/TensorFold/pull/1413 +[#1423]: https://github.com/ashhart/TensorFold/pull/1423 [#1426]: https://github.com/ashhart/TensorFold/pull/1426 [#1434]: https://github.com/ashhart/TensorFold/pull/1434 [#1439]: https://github.com/ashhart/TensorFold/pull/1439 From 8026e4533804b37cb7134a6a468ecf4f3fdad37d Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Sat, 7 Nov 2026 06:20:01 +0900 Subject: [PATCH 21/36] changelog: #1423 under Unreleased --- CHANGELOG.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 356b45fdf5..583887edda 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -65,6 +65,9 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Internal: 9 more lean findings fixed in Metal kernel and server comments, with no change in behaviour (#2329). - Native DeepSeek V4 and V4.1 prompts follow the recorded encoder outputs, except where those outputs drop a tool's namespace, which we keep. Assistant thinking markers, blank-line text joins and tool schemas change; minimal keeps thinking on at budget 50 and xhigh uses 75. Developer instructions render as system, strict tools keep their existing grammar handling, listings and tool-call history keep qualified function names, and schemas retain namespace descriptions. Invalid user null content, malformed tool-call lists or calls without string id/name/arguments or type function, names outside 1-64 ASCII letters/digits/underscore/hyphen, tool replies without tool_call_id and non-function tool envelopes now return HTTP 400. Ultra keeps its existing refusal because no recorded case establishes a budget (#2141). +- GLM-5.3-Flash on CUDA: a prompt step takes two chunks while no other stream decodes, each one's exchanges under the + other's sites, so prompts of 8K to 128K finish about 2% sooner at TP=2, with replies unchanged; + `TF_GLM_PROMPT_PAIR=0` turns it off (#1423). - TENSORFOLD_PHASE_METRICS=1 adds a bounded request queue-wait histogram and its vLLM alias to native metrics; it is off by default (#2139). From 1c22fa344c7854e26a1bda26279722a4dece788e Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Sat, 7 Nov 2026 00:57:12 +0000 Subject: [PATCH 22/36] gemma4 host: pin affine_mm.metal's bytes after cf10a6b7d's comment edit Step: test-gemma4-affine-host. Error: "landed Metal source bytes are pinned without shader compilation" fails expectEqualStrings: expected d8f47cbb..., found a91a3539... Cause: cf10a6b7d rewrapped one comment in core/affine_mm.metal to fit 120 columns and left the pin at the old bytes (the file at 6d412e724 hashes to the old pin; that comment is the only change since). Not macOS-only: the step is built on every OS but is not in zig build test, so no CI job runs it. The lane_projection pin already matches. --- zig/gemma4_shared_affine_test.zig | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zig/gemma4_shared_affine_test.zig b/zig/gemma4_shared_affine_test.zig index 5c2b25871b..2f52f224c3 100644 --- a/zig/gemma4_shared_affine_test.zig +++ b/zig/gemma4_shared_affine_test.zig @@ -103,7 +103,7 @@ test "landed Metal source bytes are pinned without shader compilation" { { var hash: [32]u8 = undefined; std.crypto.hash.sha2.Sha256.hash(ks.core_affine_mm, &hash, .{}); - try std.testing.expectEqualStrings("d8f47cbbddf24502ca2c16cfdef0672731cc2dbe59ad027f47a3e662bcf8f627", &std.fmt.bytesToHex(hash, .lower)); + try std.testing.expectEqualStrings("a91a353981c4e24142673ab9c97e9a449b973751bd5df7ad7c980757046cf61e", &std.fmt.bytesToHex(hash, .lower)); } { var hash: [32]u8 = undefined; From a569b24c23c0fc188d699b3718781f19b64f94ba Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Sat, 7 Nov 2026 00:58:03 +0000 Subject: [PATCH 23/36] build: the compile-only Metal contract runners import masked_sources too Steps: compile-fsr-metal-contract-host and compile-gemma4-contract-host. Error: "no module named 'masked_sources' available within module 'contract_main'" ('runner' for Gemma), from tests/contract/masked_attention_more.zig:5. Cause: 17e6242a7 made metal_main.zig import masked_sources and gave it to tf-kernel-contract (kernel_contract.zig) but not to these two compile-only copies of the same runner. Not macOS-only by logic: both steps exist only on macOS, and no CI job builds them. --- zig/build/fsr.zig | 1 + zig/build/gemma4_host.zig | 1 + 2 files changed, 2 insertions(+) diff --git a/zig/build/fsr.zig b/zig/build/fsr.zig index 02855d5a06..06dfc5c8e1 100644 --- a/zig/build/fsr.zig +++ b/zig/build/fsr.zig @@ -25,6 +25,7 @@ pub fn deviceSteps(b: *std.Build, metal: *std.Build.Module, engine: *std.Build.M b.step("test-fsr-metal", "GPU: uncached four-slot synthetic MoE equals resident at every node and logit").dependOn(&run.step); b.step("compile-fsr-metal-gate-host", "Compile the real FSR device gate without shader compilation or execution").dependOn(&device_test.step); const contract = b.createModule(.{ .root_source_file = b.path("zig/tests/contract/metal_main.zig"), .target = b.graph.host, .link_libc = true, .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "tensorfold", .module = engine }, .{ .name = "kernel_sources", .module = sources }, .{ .name = "contract", .module = @import("kernel_contract.zig").module(b, b.graph.host, .debug) }, .{ .name = "h3_kernels", .module = b.createModule(.{ .root_source_file = b.path("zig/kernels/metal/h3/kernels.zig") }) }, .{ .name = "ling_metal", .module = b.createModule(.{ .root_source_file = b.path("zig/kernels/metal/ling/sources.zig") }) }, .{ .name = "ling_route", .module = b.createModule(.{ .root_source_file = b.path("zig/src/families/ling/route.zig") }) } } }); + contract.addImport("masked_sources", b.createModule(.{ .root_source_file = b.path("zig/masked_value_sources.zig"), .imports = &.{.{ .name = "kernel_sources", .module = sources }} })); const compile_contract = b.addObject(.{ .name = "fsr-metal-contract", .root_module = b.createModule(.{ .root_source_file = b.path("zig/fsr_contract_compile.zig"), .target = b.graph.host, .link_libc = true, .imports = &.{.{ .name = "contract_main", .module = contract }} }) }); _ = compile_contract.getEmittedBin(); b.step("compile-fsr-metal-contract-host", "Compile the actual Metal contract entry with inert embedded bytes and no execution").dependOn(&compile_contract.step); diff --git a/zig/build/gemma4_host.zig b/zig/build/gemma4_host.zig index 36958cb157..8fde58faae 100644 --- a/zig/build/gemma4_host.zig +++ b/zig/build/gemma4_host.zig @@ -40,6 +40,7 @@ pub fn compileHost(b: *std.Build, metal: *std.Build.Module, engine: *std.Build.M runner.addImport("h3_kernels", b.createModule(.{ .root_source_file = b.path("zig/kernels/metal/h3/kernels.zig") })); runner.addAnonymousImport("ling_metal", .{ .root_source_file = b.path("zig/kernels/metal/ling/sources.zig") }); runner.addAnonymousImport("ling_route", .{ .root_source_file = b.path("zig/src/families/ling/route.zig") }); + runner.addImport("masked_sources", b.createModule(.{ .root_source_file = b.path("zig/masked_value_sources.zig"), .imports = &.{.{ .name = "kernel_sources", .module = sources }} })); const contract_root = b.createModule(.{ .root_source_file = b.path("zig/gemma4_contract_compile.zig"), .target = b.graph.host, .link_libc = true, .imports = &.{.{ .name = "runner", .module = runner }} }); const contract_object = b.addObject(.{ .name = "gemma4-contract-compile", .root_module = contract_root }); _ = contract_object.getEmittedBin(); From bc4506b710df7eeead53e2f0f9815c7ae55e00c9 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Sat, 7 Nov 2026 00:59:28 +0000 Subject: [PATCH 24/36] build: compile-metal-residency-host's native module imports kimi_k3 Step: compile-metal-residency-host. Error: "no module named 'kimi_k3' available within module 'host'" at native/kimi_metal.zig:4. Cause: 2c4fa0c1a made native/metal.zig open Kimi K3 through kimi_metal.zig, and 57ff96d51 wired kimi_k3 and its inert metallib into slide_host.zig's copy of that module but not into this one. The inert bytes are now one module shared by Drama and Kimi (a file roots one module). Not macOS-only by logic: the step exists only on macOS and no CI job builds it. test-metal-formats-host, on the same module, still passes. --- zig/build/residency_host.zig | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/zig/build/residency_host.zig b/zig/build/residency_host.zig index 55378f3021..e184514552 100644 --- a/zig/build/residency_host.zig +++ b/zig/build/residency_host.zig @@ -23,18 +23,20 @@ pub fn steps(b: *std.Build) void { pub fn compileHost(b: *std.Build, metal: *std.Build.Module, engine: *std.Build.Module, lanes: *std.Build.Module) void { const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = b.graph.host, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} }); const host = b.createModule(.{ .root_source_file = b.path("zig/src/native/metal.zig"), .target = b.graph.host, .link_libc = true, .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "tensorfold", .module = engine }, .{ .name = "engine_api", .module = api } } }); - const bytes = b.addWriteFiles().add("residency-drama-empty.zig", "pub const bytes: []const u8 = &.{};\n"); + const bytes = b.createModule(.{ .root_source_file = b.addWriteFiles().add("residency-drama-empty.zig", "pub const bytes: []const u8 = &.{};\n") }); const tokenizer = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = b.graph.host, .link_libc = true }); const rounds = b.createModule(.{ .root_source_file = b.path("zig/src/core/frame_rounds.zig"), .target = b.graph.host }); const shim = b.addWriteFiles().add("residency-drama-core.zig", "pub const safetensors = @import(\"tensorfold\").safetensors;\npub const tokenizer = @import(\"tokenizer\");\npub const frame_rounds = @import(\"frame_rounds\");\n"); const core = b.createModule(.{ .root_source_file = shim, .target = b.graph.host, .imports = &.{ .{ .name = "tensorfold", .module = engine }, .{ .name = "tokenizer", .module = tokenizer }, .{ .name = "frame_rounds", .module = rounds } } }); - const drama = b.createModule(.{ .root_source_file = b.path("zig/src/families/drama3/drama3_metal.zig"), .target = b.graph.host, .link_libc = true, .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "core", .module = core }, .{ .name = "drama3_metallib", .module = b.createModule(.{ .root_source_file = bytes }) } } }); + const drama = b.createModule(.{ .root_source_file = b.path("zig/src/families/drama3/drama3_metal.zig"), .target = b.graph.host, .link_libc = true, .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "core", .module = core }, .{ .name = "drama3_metallib", .module = bytes } } }); host.addImport("core", core); host.addImport("tokenizer", tokenizer); host.addImport("drama3", drama); // native/metal_slide.zig imports living: the engine module's own (a file belongs to one module) host.addImport("living", engine.import_table.get("living") orelse @panic("the engine has no living import")); @import("qwen4.zig").metalNative(b, b.graph.host, host); + host.addImport("kimi_k3", @import("metal_targets.zig").kimiModule(b, b.graph.host, .Debug, metal, engine, lanes)); + host.addImport("kimi_metallib", bytes); b.step("test-metal-formats-host", "Actual Metal format offers with inert images and no device").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = host, .filters = &.{"Metal capabilities:"} })).step); const root = b.createModule(.{ .root_source_file = b.path("zig/residency_native_compile.zig"), .target = b.graph.host, .link_libc = true, .imports = &.{.{ .name = "host", .module = host }} }); const object = b.addObject(.{ .name = "metal-residency-native", .root_module = root }); From 13bdba9c00a89243aa37ba932f19d3b820592390 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Sat, 7 Nov 2026 01:03:21 +0000 Subject: [PATCH 25/36] glm metal: the backend fact tests give facts() a host device object Step: test-shared-head-cost-metal-host. Error: "shared head costs: GLM Metal MTP retains costed copy grants" terminated with signal ABRT, "Segmentation fault at address 0x10" in objc_msgSend from Device.responds <- generation <- tensorUnits <- Backend.facts. Cause: ae320aafb made facts() ask e.device.tensorUnits(), but this test and "zero grant: GLM Metal MTP resumes after a long priced skip" (test-drafter-skip-metal-host) leave the Engine undefined, so facts() messages an uninitialized id. The outcome changes from process to process: on origin/main the shared-head-cost binary aborted in 15 of 20 direct runs, and the zero-grant binary segfaulted in 20 of 20 direct runs while passing under the build runner. Both tests now set the NSObject host object that draft_sources_metal_test.zig already uses: it names no Metal architecture, so tensorUnits() is false without a GPU. Fixed binaries: 20 of 20 each. Not macOS-only by logic: the tests build only on macOS. --- zig/src/families/glm/backend.zig | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/zig/src/families/glm/backend.zig b/zig/src/families/glm/backend.zig index 70609d6f99..0365621332 100644 --- a/zig/src/families/glm/backend.zig +++ b/zig/src/families/glm/backend.zig @@ -441,6 +441,10 @@ test "zero grant: GLM Metal MTP resumes after a long priced skip" { var e: Engine = undefined; e.w = &w; e.ep = null; + const objc = @import("metal").objc; + const no_arch = objc.msg(objc.Id, objc.class("NSObject"), "new", .{}); // facts() asks it for tensor units, no GPU + defer objc.release(no_arch); + e.device = .{ .id = no_arch }; var slots: [16]slots_mod.Slot = undefined; var sl: slots_mod.Slots = undefined; sl.e = &e; @@ -457,6 +461,10 @@ test "shared head costs: GLM Metal MTP retains costed copy grants" { var e: Engine = undefined; e.w = &w; e.ep = null; + const objc = @import("metal").objc; + const no_arch = objc.msg(objc.Id, objc.class("NSObject"), "new", .{}); // facts() asks it for tensor units, no GPU + defer objc.release(no_arch); + e.device = .{ .id = no_arch }; var slots: [16]slots_mod.Slot = undefined; var sl: slots_mod.Slots = undefined; sl.e = &e; From f29e504a67b8884e6e22dd4af5a250706661c20e Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Sat, 7 Nov 2026 04:36:57 +0000 Subject: [PATCH 26/36] changelog: #2409 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 583887edda..681951936b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Tests: four host steps that no CI job runs build and pass again: `test-gemma4-affine-host` pins `affine_mm.metal`'s current bytes, the compile-only Metal contract runners import `masked_sources`, `compile-metal-residency-host` imports `kimi_k3`, and the GLM Metal fact tests give `facts()` a host device object instead of an undefined one, which crashed them at random (#2409, part of #2031). + - Flash Next on CUDA with `TF_FLASHNEXT_KV_PAGES`: a shared round's grouped draft head reads and writes each stream's own pages, so paged streams' replies equal their runs alone again, and paged decode runs as fast as unpaged; default settings are unchanged (#2244). Thanks to @Arminova. - Tests: `zig build test-glm-snapshots-host` builds again: its copies module imports `core`, which `cuda_snaps.zig` reads for `TF_GLM_PROMPT_TAIL` since #2226 (#2403). Thanks to @Bizuayeu. From 8539b7703f2e49af815ea1d6ae188d6efd336c9e Mon Sep 17 00:00:00 2001 From: elderorb <23638+ElderOrb@users.noreply.github.com> Date: Thu, 5 Nov 2026 22:39:07 +0100 Subject: [PATCH 27/36] cli: run the pull tests' fake hub on std.Io.net so Windows builds them --- .github/workflows/no-gpu.yml | 5 +- zig/src/cli/pull_test.zig | 165 ++++++++++++++++++++--------------- 2 files changed, 97 insertions(+), 73 deletions(-) diff --git a/.github/workflows/no-gpu.yml b/.github/workflows/no-gpu.yml index 99114cfd67..7409c0acf7 100644 --- a/.github/workflows/no-gpu.yml +++ b/.github/workflows/no-gpu.yml @@ -1,4 +1,4 @@ -name: No GPU +name: No GPU on: push: @@ -176,6 +176,9 @@ jobs: run: zig build -Dfixed-cpu test-golden - name: zig build fuzz-cortheon run: zig build -Dfixed-cpu fuzz-cortheon + # The pull client's framing and path checks: its fake hub runs on std.Io.net, so Windows builds and runs them. + - name: zig build test-pull-security + run: zig build -Dfixed-cpu test-pull-security - name: Save the build cache if: ${{ !cancelled() && github.event_name != 'pull_request' && github.ref == 'refs/heads/main' }} uses: actions/cache/save@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 diff --git a/zig/src/cli/pull_test.zig b/zig/src/cli/pull_test.zig index 43616c4045..64d01586a5 100644 --- a/zig/src/cli/pull_test.zig +++ b/zig/src/cli/pull_test.zig @@ -8,8 +8,12 @@ const exlr8_pull = @import("pull_exlr8.zig"); /// A fake hub on 127.0.0.1: revision and tree APIs, ranged resolve endpoints, a body cut once, a repo ignoring ranges. const FakeHub = struct { - server_fd: std.posix.socket_t, + server: std.Io.net.Server, + io: std.Io, + threaded: *std.Io.Threaded, stop: std.atomic.Value(bool) = .init(false), + live: std.ArrayList(std.Io.net.Stream) = .empty, + live_lock: std.Io.Mutex = .init, thread: ?std.Thread = null, weight_requests: std.atomic.Value(u32) = .init(0), flaky_requests: std.atomic.Value(u32) = .init(0), @@ -33,20 +37,13 @@ const FakeHub = struct { const nested_body: []const u8 = &nested_body_storage; fn open() !FakeHub { - const posix = std.posix; - const rc = posix.system.socket(posix.AF.INET, posix.SOCK.STREAM, 0); - if (posix.errno(rc) != .SUCCESS) return error.SocketFailed; - const fd: posix.socket_t = @intCast(rc); - errdefer _ = posix.system.close(fd); - const one: c_int = 1; - try posix.setsockopt(fd, posix.SOL.SOCKET, posix.SO.REUSEADDR, std.mem.asBytes(&one)); - var addr: posix.sockaddr.in = .{ .family = posix.AF.INET, .port = std.mem.nativeToBig(u16, 0), .addr = std.mem.nativeToBig(u32, 0x7f000001) }; - if (posix.errno(posix.system.bind(fd, @ptrCast(&addr), @sizeOf(posix.sockaddr.in))) != .SUCCESS) return error.BindFailed; - if (posix.errno(posix.system.listen(fd, 32)) != .SUCCESS) return error.ListenFailed; - var bound: posix.sockaddr.in = undefined; - var len: posix.socklen_t = @sizeOf(posix.sockaddr.in); - if (posix.errno(posix.system.getsockname(fd, @ptrCast(&bound), &len)) != .SUCCESS) return error.BindFailed; - const port = std.mem.bigToNative(u16, bound.port); + const threaded = try std.testing.allocator.create(std.Io.Threaded); + threaded.* = std.Io.Threaded.init(std.testing.allocator, .{}); + errdefer std.testing.allocator.destroy(threaded); + const io = threaded.io(); + const loopback: std.Io.net.IpAddress = .{ .ip4 = .loopback(0) }; + var server = try loopback.listen(io, .{}); + errdefer server.deinit(io); var hash = std.crypto.hash.sha2.Sha256.init(.{}); hash.update(weights_body); var digest: [32]u8 = undefined; @@ -56,7 +53,9 @@ const FakeHub = struct { var nested_digest: [32]u8 = undefined; nested_hash.final(&nested_digest); return .{ - .server_fd = fd, + .server = server, + .io = io, + .threaded = threaded, .weights = weights_body, .config = config_body, .nested = nested_body, @@ -64,7 +63,7 @@ const FakeHub = struct { .nested_sha_hex = try std.fmt.allocPrint(std.testing.allocator, "{s}", .{std.fmt.bytesToHex(nested_digest, .lower)}), .config_git_hex = gitHex(config_body), .flaky_git_hex = gitHex(flaky_body), - .port = port, + .port = server.socket.address.getPort(), }; } @@ -85,43 +84,46 @@ const FakeHub = struct { var n: usize = 0; defer for (conns[0..n]) |t| t.join(); while (!fake.stop.load(.acquire)) { - var fds = [_]std.posix.pollfd{.{ .fd = fake.server_fd, .events = std.posix.POLL.IN, .revents = 0 }}; - const ready = std.posix.poll(&fds, 50) catch return; - if (ready == 0) continue; - const rc = std.posix.system.accept(fake.server_fd, null, null); - if (std.posix.errno(rc) != .SUCCESS) return; - const conn_fd: std.posix.socket_t = @intCast(rc); + const conn = fake.server.accept(fake.io) catch return; + if (fake.stop.load(.acquire)) { + conn.close(fake.io); + return; + } if (n == conns.len) { - _ = std.posix.system.close(conn_fd); + conn.close(fake.io); continue; } - conns[n] = std.Thread.spawn(.{}, handle, .{ fake, conn_fd }) catch { - _ = std.posix.system.close(conn_fd); + fake.live_lock.lockUncancelable(fake.io); + fake.live.append(std.testing.allocator, conn) catch {}; + fake.live_lock.unlock(fake.io); + conns[n] = std.Thread.spawn(.{}, handle, .{ fake, conn }) catch { + fake.forget(conn); + conn.close(fake.io); continue; }; n += 1; } } - fn handle(fake: *FakeHub, fd: std.posix.socket_t) void { - defer _ = std.posix.system.close(fd); + fn handle(fake: *FakeHub, conn: std.Io.net.Stream) void { + defer { + fake.forget(conn); + conn.close(fake.io); + } var buf: [8192]u8 = undefined; var end: usize = 0; + var read_buf: [4096]u8 = undefined; + var reader = conn.reader(fake.io, &read_buf); while (true) { const head_end = std.mem.indexOf(u8, buf[0..end], "\r\n\r\n") orelse { if (end == buf.len) return; - var fds = [_]std.posix.pollfd{.{ .fd = fd, .events = std.posix.POLL.IN, .revents = 0 }}; - const ready = std.posix.poll(&fds, 50) catch return; - if (ready == 0) { - if (fake.stop.load(.acquire)) return; - continue; - } - const n = std.posix.read(fd, buf[end..]) catch return; + var vecs = [_][]u8{buf[end..]}; + const n = reader.interface.readVec(&vecs) catch return; if (n == 0) return; end += n; continue; }; - if (!fake.respondOne(fd, buf[0..head_end])) return; + if (!fake.respondOne(conn, buf[0..head_end])) return; const rest = end - (head_end + 4); std.mem.copyForwards(u8, buf[0..rest], buf[head_end + 4 .. end]); end = rest; @@ -129,13 +131,13 @@ const FakeHub = struct { } /// Answers one request; false closes the connection. - fn respondOne(fake: *FakeHub, fd: std.posix.socket_t, head: []const u8) bool { + fn respondOne(fake: *FakeHub, conn: std.Io.net.Stream, head: []const u8) bool { var tree_buf: [256]u8 = undefined; const line_end = std.mem.indexOf(u8, head, "\r\n") orelse return false; var parts = std.mem.tokenizeScalar(u8, head[0..line_end], ' '); const method = parts.next() orelse return false; if (std.mem.eql(u8, method, "HEAD")) { - sendAll(fd, "HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: keep-alive\r\n\r\n"); // no redirect: the hub's link is the storage link + fake.sendAll(conn, "HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: keep-alive\r\n\r\n"); // no redirect: the hub's link is the storage link return true; } if (!std.mem.eql(u8, method, "GET")) return false; @@ -157,86 +159,86 @@ const FakeHub = struct { } } if (std.mem.eql(u8, path, "/api/models/Org/Flash/revision/main")) { - respond(fd, "{\"sha\": \"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa1\"}"); + fake.respond(conn, "{\"sha\": \"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa1\"}"); } else if (std.mem.eql(u8, path, "/api/models/Org/Whole/revision/main")) { - respond(fd, "{\"sha\": \"bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb2\"}"); + fake.respond(conn, "{\"sha\": \"bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb2\"}"); } else if (std.mem.eql(u8, path, "/api/models/Org/Draft/revision/main")) { - respond(fd, "{\"sha\": \"ccccccccccccccccccccccccccccccccccccccc3\"}"); + fake.respond(conn, "{\"sha\": \"ccccccccccccccccccccccccccccccccccccccc3\"}"); } else if (std.mem.eql(u8, path, "/api/models/Org/Draft/tree/ccccccccccccccccccccccccccccccccccccccc3")) { - respond(fd, "[{\"type\": \"file\", \"path\": \"config.json\", \"size\": 27}]"); + fake.respond(conn, "[{\"type\": \"file\", \"path\": \"config.json\", \"size\": 27}]"); } else if (std.mem.eql(u8, path, "/Org/Draft/resolve/ccccccccccccccccccccccccccccccccccccccc3/config.json")) { - respond(fd, "{\"model_type\": \"gemma3n\"}"); + fake.respond(conn, "{\"model_type\": \"gemma3n\"}"); } else if (std.mem.eql(u8, path, "/api/models/Org/DFlash/revision/main")) { - respond(fd, "{\"sha\": \"ddddddddddddddddddddddddddddddddddddddd4\"}"); + fake.respond(conn, "{\"sha\": \"ddddddddddddddddddddddddddddddddddddddd4\"}"); } else if (std.mem.eql(u8, path, "/api/models/Org/DFlash/tree/ddddddddddddddddddddddddddddddddddddddd4")) { - respond(fd, std.fmt.bufPrint(&tree_buf, "[{{\"type\": \"file\", \"path\": \"config.json\", \"size\": {d}}}]", .{drafter_config.len}) catch return false); + fake.respond(conn, std.fmt.bufPrint(&tree_buf, "[{{\"type\": \"file\", \"path\": \"config.json\", \"size\": {d}}}]", .{drafter_config.len}) catch return false); } else if (std.mem.eql(u8, path, "/Org/DFlash/resolve/ddddddddddddddddddddddddddddddddddddddd4/config.json")) { - respond(fd, drafter_config); + fake.respond(conn, drafter_config); } else if (std.mem.eql(u8, path, "/api/models/Org/Flash/tree/aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa1") or std.mem.eql(u8, path, "/api/models/Org/Whole/tree/bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb2")) { // The hub's tree sends LFS oids as bare hex. const tree = std.fmt.allocPrint(std.testing.allocator, "[{{\"type\": \"file\", \"oid\": \"{s}\", \"path\": \"config.json\", \"size\": {d}}}, {{\"type\": \"file\", \"oid\": \"1111111111111111111111111111111111111111\", \"path\": \"weights.safetensors\", \"size\": {d}, \"lfs\": {{\"oid\": \"{s}\", \"size\": {d}}}}}, {{\"type\": \"file\", \"oid\": \"{s}\", \"path\": \"flaky.json\", \"size\": {d}}}, {{\"type\": \"file\", \"oid\": \"3333333333333333333333333333333333333333\", \"path\": \"tensors/model-00001-of-00001.safetensors\", \"size\": {d}, \"lfs\": {{\"oid\": \"{s}\", \"size\": {d}}}}}]", .{ &fake.config_git_hex, fake.config.len, fake.weights.len, fake.weights_sha_hex, fake.weights.len, &fake.flaky_git_hex, flaky_body.len, fake.nested.len, fake.nested_sha_hex, fake.nested.len }) catch return false; defer std.testing.allocator.free(tree); - respond(fd, tree); + fake.respond(conn, tree); } else if (std.mem.eql(u8, path, config_path) or std.mem.eql(u8, path, "/Org/Whole/resolve/bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb2/config.json")) { - respond(fd, fake.config); + fake.respond(conn, fake.config); } else if (std.mem.endsWith(u8, path, "/flaky.json")) { if (fake.flaky_requests.fetchAdd(1, .monotonic) > 0) { - respond(fd, flaky_body); + fake.respond(conn, flaky_body); return true; } var head_buf: [128]u8 = undefined; // the first answer stops halfway and drops the connection - sendAll(fd, std.fmt.bufPrint(&head_buf, "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {d}\r\n\r\n", .{flaky_body.len}) catch return false); - sendAll(fd, flaky_body[0 .. flaky_body.len / 2]); + fake.sendAll(conn, std.fmt.bufPrint(&head_buf, "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {d}\r\n\r\n", .{flaky_body.len}) catch return false); + fake.sendAll(conn, flaky_body[0 .. flaky_body.len / 2]); return false; } else if (std.mem.eql(u8, path, "/Org/Whole/resolve/bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb2/weights.safetensors")) { _ = fake.weight_requests.fetchAdd(1, .monotonic); - whole(fd, fake.weights); + fake.whole(conn, fake.weights); + return false; } else if (std.mem.eql(u8, path, nested_path) or std.mem.eql(u8, path, "/Org/Whole/resolve/bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb2/tensors/model-00001-of-00001.safetensors")) { - whole(fd, fake.nested); + fake.whole(conn, fake.nested); + return false; } else if (std.mem.eql(u8, path, weights_path)) { _ = fake.weight_requests.fetchAdd(1, .monotonic); const from = range_start orelse { - whole(fd, fake.weights); + fake.whole(conn, fake.weights); return true; }; fake.range_starts.store(from, .monotonic); if (from >= fake.weights.len) { - respond(fd, ""); + fake.respond(conn, ""); return true; } const stop = if (range_end) |last| @min(last + 1, fake.weights.len) else fake.weights.len; const body = fake.weights[from..stop]; var head_buf: [256]u8 = undefined; const head_text = std.fmt.bufPrint(&head_buf, "HTTP/1.1 206 Partial Content\r\nContent-Type: application/octet-stream\r\nContent-Length: {d}\r\nContent-Range: bytes {d}-{d}/{d}\r\nConnection: keep-alive\r\n\r\n", .{ body.len, from, stop - 1, fake.weights.len }) catch return false; - sendAll(fd, head_text); - sendAll(fd, body); + fake.sendAll(conn, head_text); + fake.sendAll(conn, body); } else { - respond(fd, "{}"); + fake.respond(conn, "{}"); } return true; } - fn whole(fd: std.posix.socket_t, body: []const u8) void { + fn whole(fake: *FakeHub, conn: std.Io.net.Stream, body: []const u8) void { var head_buf: [128]u8 = undefined; const head_text = std.fmt.bufPrint(&head_buf, "HTTP/1.1 200 OK\r\nContent-Type: application/octet-stream\r\nContent-Length: {d}\r\nConnection: keep-alive\r\n\r\n", .{body.len}) catch return; - sendAll(fd, head_text); - sendAll(fd, body); + fake.sendAll(conn, head_text); + fake.sendAll(conn, body); } - fn respond(fd: std.posix.socket_t, body: []const u8) void { + fn respond(fake: *FakeHub, conn: std.Io.net.Stream, body: []const u8) void { var head_buf: [128]u8 = undefined; const head_text = std.fmt.bufPrint(&head_buf, "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {d}\r\nConnection: keep-alive\r\n\r\n", .{body.len}) catch return; - sendAll(fd, head_text); - sendAll(fd, body); + fake.sendAll(conn, head_text); + fake.sendAll(conn, body); } - fn sendAll(fd: std.posix.socket_t, bytes: []const u8) void { - var sent: usize = 0; - while (sent < bytes.len) { - const n = std.posix.system.write(fd, bytes[sent..].ptr, bytes.len - sent); - if (std.posix.errno(n) != .SUCCESS) return; - sent += @intCast(n); - } + fn sendAll(fake: *FakeHub, conn: std.Io.net.Stream, bytes: []const u8) void { + var buf: [4096]u8 = undefined; + var w = conn.writer(fake.io, &buf); + w.interface.writeAll(bytes) catch return; + w.interface.flush() catch return; } fn start(fake: *FakeHub) void { @@ -247,10 +249,29 @@ const FakeHub = struct { return std.fmt.allocPrint(a, "http://127.0.0.1:{d}", .{fake.port}); } + fn forget(fake: *FakeHub, conn: std.Io.net.Stream) void { + fake.live_lock.lockUncancelable(fake.io); + defer fake.live_lock.unlock(fake.io); + for (fake.live.items, 0..) |c, i| { + if (c.socket.handle == conn.socket.handle) { + _ = fake.live.swapRemove(i); + return; + } + } + } + fn stopServer(fake: *FakeHub) void { fake.stop.store(true, .release); + fake.live_lock.lockUncancelable(fake.io); + for (fake.live.items) |conn| conn.shutdown(fake.io, .both) catch {}; + fake.live_lock.unlock(fake.io); + // one connection of our own unblocks accept, and is closed on arrival + if (fake.server.socket.address.connect(fake.io, .{ .mode = .stream }) catch null) |wake| wake.close(fake.io); if (fake.thread) |t| t.join(); - _ = std.posix.system.close(fake.server_fd); + fake.server.deinit(fake.io); + fake.live.deinit(std.testing.allocator); + fake.threaded.deinit(); + std.testing.allocator.destroy(fake.threaded); std.testing.allocator.free(fake.weights_sha_hex); std.testing.allocator.free(fake.nested_sha_hex); } From c9713f0f17222deab67e8b47982ea1fc3078a23a Mon Sep 17 00:00:00 2001 From: elderorb <23638+ElderOrb@users.noreply.github.com> Date: Fri, 6 Nov 2026 00:25:18 +0100 Subject: [PATCH 28/36] ci: drop the workflow's byte-order mark and note the Windows pull tests --- .github/workflows/no-gpu.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/no-gpu.yml b/.github/workflows/no-gpu.yml index 7409c0acf7..7188bd4dcd 100644 --- a/.github/workflows/no-gpu.yml +++ b/.github/workflows/no-gpu.yml @@ -1,4 +1,4 @@ -name: No GPU +name: No GPU on: push: From f275d991882e78e2ab88f35c0f2cd0a2ab42d6b5 Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Fri, 6 Nov 2026 08:37:45 +0000 Subject: [PATCH 29/36] changelog: #2284 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 681951936b..df775c2343 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -138,6 +138,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Qwen3.8-27B and the other dense Qwen3.5 models on CUDA: streams run on KV pool pages of 256 positions by default, the engine's own stream too, so `--parallel` is now a cap rather than a reservation: a request can wait for pages instead of every stream holding a full-context window. `TF_QWEN_KV_PAGES=0` keeps a window each, as before. Qwen3.6-35B-A3B and the other MoE checkpoints stay on windows by default, with pages opt-in (`TF_QWEN_KV_PAGES=256`). Paged prompt attention stages a chunk's page ids in shared memory, a decode thread looks its page up once a page, and a paged prompt-cache save or restore copies one 2D block a page. Replies are the same tokens on pages and windows (#2011). Thanks to @jschmied. - Qwen3.5 and Qwen3.6 on CUDA: KV pool pages stay opt-in (`TF_QWEN_KV_PAGES=256`), dense and MoE alike; unset keeps a window each, as before, and `TF_QWEN_KV_PAGES_ALL` unset still leaves the engine's own stream a window. With pages set, `--parallel` is a cap rather than a reservation: a request can wait for pages instead of every stream holding a full-context window. Paged prompt attention stages a chunk's page ids in shared memory, a decode thread looks its page up once a page, and a paged prompt-cache save or restore copies one 2D block a page. Replies are the same tokens on pages and windows (#2011). Thanks to @jschmied. - Qwen3.5 and Qwen3.6 on CUDA: KV pool pages stay opt-in (`TF_QWEN_KV_PAGES=256`), dense and MoE alike; unset keeps a window each, as before, and `TF_QWEN_KV_PAGES_ALL` unset still leaves the engine's own stream a window. With pages set, `--parallel` is a cap rather than a reservation: a request can wait for pages instead of every stream holding a full-context window. Paged prompt attention stages a chunk's page ids in shared memory, a decode thread looks its page up once a page, and a paged prompt-cache save or restore copies one 2D block a page. Replies are the same tokens on pages and windows (#2011, part of #1943). Thanks to @jschmied. +- Tests: the pull client's security tests build and run on Windows, on a portable fake hub, and CI runs them there + (#2284). Thanks to @ElderOrb. - Metal: a request's finished stats list its prompt pass widths (`tensorfold.prefill_widths`, one entry per completed target pass) on every Metal text family except Flash Next, whose engine reports none. Telemetry only (#2293). From 39fd1ab9f93a5faf6c678b08d0c146d5296b96bf Mon Sep 17 00:00:00 2001 From: Christopher Maher <7034124+Defilan@users.noreply.github.com> Date: Fri, 6 Nov 2026 13:48:53 -0800 Subject: [PATCH 30/36] build: the DeepSeek server gets the lanes module its /v1/systemone route reads --- zig/build/dsv41_draft.zig | 1 + 1 file changed, 1 insertion(+) diff --git a/zig/build/dsv41_draft.zig b/zig/build/dsv41_draft.zig index 274fd7ce7c..f8470e0a13 100644 --- a/zig/build/dsv41_draft.zig +++ b/zig/build/dsv41_draft.zig @@ -68,6 +68,7 @@ pub fn serverOn(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b .link_libc = true, .imports = &.{ .{ .name = "core", .module = engine.import_table.get("core").? }, .{ .name = "engine_api", .module = api }, .{ .name = "tokenizer", .module = tokenizer }, .{ .name = "template", .module = template }, .{ .name = "native_engines", .module = engine }, .{ .name = "checkpoint_cli", .module = cli } }, }); + m.addImport("lanes", engine.import_table.get("lanes").?); m.addImport("living", b.createModule(.{ .root_source_file = b.path("zig/src/living/living.zig"), .target = target, .optimize = optimize, .link_libc = true })); m.addOptions("build_options", build_options); // the server's json and the family's text side share the engine's serving module (one `dsv41_serve`) From bcd846173bdaf1a0bce800715f2c0dadc1d837fd Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Sat, 7 Nov 2026 04:37:00 +0000 Subject: [PATCH 31/36] changelog: #2412 under Unreleased --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index df775c2343..c923f019c8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel ## Unreleased +- Build: `zig build compile-dsv41-server` and `dsv41-serve` build again: the DeepSeek server module gets the `lanes` module its `/v1/systemone` route reads (#2412). Thanks to @Defilan. + - Tests: four host steps that no CI job runs build and pass again: `test-gemma4-affine-host` pins `affine_mm.metal`'s current bytes, the compile-only Metal contract runners import `masked_sources`, `compile-metal-residency-host` imports `kimi_k3`, and the GLM Metal fact tests give `facts()` a host device object instead of an undefined one, which crashed them at random (#2409, part of #2031). - Flash Next on CUDA with `TF_FLASHNEXT_KV_PAGES`: a shared round's grouped draft head reads and writes each stream's own pages, so paged streams' replies equal their runs alone again, and paged decode runs as fast as unpaged; default settings are unchanged (#2244). Thanks to @Arminova. From 91377a4392b5e9a32595d2464f246e6459bd19ce Mon Sep 17 00:00:00 2001 From: Whq02 <178854749+Whq02@users.noreply.github.com> Date: Fri, 6 Nov 2026 15:43:40 +0000 Subject: [PATCH 32/36] cli: tensorfold models --json, the listed checkpoints as one document `tensorfold models` prints one line a checkpoint and takes no argument, so a program that wants the inventory parses the prose or reads the Hugging Face cache itself; `models --json` printed the usage and exited 2. `models --json` now prints the same checkpoints as one document, {"schema": 1, "models": [...]}, each entry the line's repo, family and model_type; a cache that does not exist yet is an empty list. The lines are unchanged, and any other word after `models` is still the usage. The README's Models on disk names the flag. From the review: the CHANGELOG line under Unreleased, `models [--json]` in the native binary's usage text (both lines in main.zig), and the README says that `family` is a display title and `model_type` the key to match on. --- CHANGELOG.md | 1 + README.md | 2 ++ zig/src/cli/cli.zig | 7 +++-- zig/src/cli/models.zig | 69 +++++++++++++++++++++++++++++++++-------- zig/src/server/main.zig | 4 +-- 5 files changed, 65 insertions(+), 18 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c923f019c8..4a40563961 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -98,6 +98,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Docs: the pull request pre-flight asks a PR that adds, finishes or drops a family's optimization to update that family's cell in the optimization ledger (#2243). Thanks to @jschmied. - Server: a refused `--api-key-file` logs the path it was given, the mode the file system reported and which test failed, once at startup and once when a reload is refused (#2365). Thanks to @Whq02. +- CLI: `tensorfold models --json` prints the cached checkpoints as one document (`schema` 1, each checkpoint's `repo`, `family` and `model_type`); the lines are unchanged (#2364). Thanks to @Whq02. - Server: tool calls completed just before the thinking-close marker are returned when no answer follows, whole and streamed; calls followed by more reasoning or an answer stay in reasoning (#2313, fixes #2310). Thanks to @AdrianBinDC. diff --git a/README.md b/README.md index 270d356107..01b5b03311 100644 --- a/README.md +++ b/README.md @@ -234,12 +234,14 @@ On Metal, each of up to 128 streams in the 27B's shared rounds drafts its own DF ```sh tensorfold pull TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit tensorfold models +tensorfold models --json tensorfold info TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit ``` `pull` downloads a checkpoint into the Hugging Face cache and checks what it downloads against the Hub: large files by sha256, small files by their git blob sha1. Large files come down in pieces over up to 16 connections, and an interrupted pull resumes with the pieces it lacks. It refuses a checkpoint whose `model_type` has no native family, and fetches the 27B's DFlash2 drafter for `--drafter`. `models` lists the cached checkpoints with a model type it knows, and `info` shows one checkpoint's family, format, context and memory floor. +`models --json` prints the same list as one document for programs: `schema` 1, and each checkpoint's `repo`, `family` and `model_type`; `family` is a display title and `model_type` is the key to match on. `tensorfold serve REPO` serves a cached checkpoint by its repository name. ## Context compaction diff --git a/zig/src/cli/cli.zig b/zig/src/cli/cli.zig index 41c443f508..e8aeffa579 100644 --- a/zig/src/cli/cli.zig +++ b/zig/src/cli/cli.zig @@ -8,7 +8,7 @@ const pull_exlr8 = @import("pull_exlr8.zig"); const lora_merge = @import("core").lora_merge; const usage = - \\usage: tensorfold models + \\usage: tensorfold models [--json] \\ tensorfold info MODEL \\ tensorfold pull REPO[@REVISION] \\ tensorfold pull REPO[@REVISION] --exlr8 PROFILE --node R/N --dir DIR @@ -38,11 +38,12 @@ pub fn main(init: std.process.Init, argv: []const [:0]const u8) !u8 { for (args, argv) |*x, y| x.* = y; const env = init.environ_map; if (std.mem.eql(u8, args[0], "models")) { - if (args.len != 1) { + const json = args.len == 2 and std.mem.eql(u8, args[1], "--json"); + if (args.len != 1 and !json) { try err.interface.writeAll(usage); return 2; } - return models.run(a, init.io, &out.interface, env, null); + return models.run(a, init.io, &out.interface, env, null, json); } if (std.mem.eql(u8, args[0], "pull")) { if (args.len < 2) { diff --git a/zig/src/cli/models.zig b/zig/src/cli/models.zig index 8d0f1040f7..4a7329d3d4 100644 --- a/zig/src/cli/models.zig +++ b/zig/src/cli/models.zig @@ -3,15 +3,19 @@ const std = @import("std"); const Allocator = std.mem.Allocator; const hub = @import("hub.zig"); -/// Lists each cache checkpoint whose model_type a Zig family serves. Returns the process exit code. -pub fn run(a: Allocator, io: std.Io, out: *std.Io.Writer, env: ?*const std.process.Environ.Map, override: ?[]const u8) !u8 { +/// One listed checkpoint: the facts a line prints, and the fields of a `--json` entry. +const Listed = struct { repo: []const u8, family: []const u8, model_type: []const u8 }; + +/// The cache checkpoints a Zig family serves, as lines or as one JSON document. Returns the process exit code. +pub fn run(a: Allocator, io: std.Io, out: *std.Io.Writer, env: ?*const std.process.Environ.Map, override: ?[]const u8, json: bool) !u8 { const root = try hub.cacheDir(a, env, override); var dir = std.Io.Dir.cwd().openDir(io, root, .{ .iterate = true }) catch { + if (json) return printJson(out, &.{}); try out.print("no checkpoints on disk: the Hugging Face cache at {s} does not exist yet; run: tensorfold pull OWNER/NAME\n", .{root}); return 0; }; defer dir.close(io); - var found: usize = 0; + var listed: std.ArrayList(Listed) = .empty; var it = dir.iterate(); while (try it.next(io)) |entry| { if (entry.kind != .directory) continue; @@ -19,10 +23,18 @@ pub fn run(a: Allocator, io: std.Io, out: *std.Io.Writer, env: ?*const std.proce const snapshot = (try hub.cachedSnapshot(a, io, root, repo)) orelse continue; const model_type = hub.modelType(a, io, snapshot); const family = hub.family(model_type) orelse continue; - try out.print("{s}: {s} ({s})\n", .{ repo, family.title, family.model_type }); - found += 1; + try listed.append(a, .{ .repo = repo, .family = family.title, .model_type = family.model_type }); } - if (found == 0) try out.print("no checkpoints on disk that a registered Zig family can serve; run: tensorfold pull OWNER/NAME\n", .{}); + if (json) return printJson(out, listed.items); + for (listed.items) |m| try out.print("{s}: {s} ({s})\n", .{ m.repo, m.family, m.model_type }); + if (listed.items.len == 0) try out.print("no checkpoints on disk that a registered Zig family can serve; run: tensorfold pull OWNER/NAME\n", .{}); + return 0; +} + +/// `{"schema": 1, "models": [...]}`: the same checkpoints the lines print, one object each, for programs. +fn printJson(out: *std.Io.Writer, listed: []const Listed) !u8 { + try std.json.Stringify.value(.{ .schema = 1, .models = listed }, .{}, out); + try out.writeByte('\n'); return 0; } @@ -36,13 +48,9 @@ fn repoFromDirName(a: Allocator, name: []const u8) !?[]const u8 { return try std.fmt.allocPrint(a, "{s}/{s}", .{ rest[0..sep], rest[sep + 2 ..] }); } -test "models lists serve-able checkpoints and skips families Zig cannot serve" { - var arena_state = std.heap.ArenaAllocator.init(std.testing.allocator); - defer arena_state.deinit(); - const a = arena_state.allocator(); - const io = std.testing.io; +/// A cache with one checkpoint a Zig family serves (Org/Flash) and one none does (Org/Draft); the caller deletes it. +fn syntheticCache(a: Allocator, io: std.Io) ![]const u8 { const root = try std.fmt.allocPrint(a, ".tf-models-test-{d}", .{std.Io.Clock.awake.now(io).toNanoseconds()}); - defer std.Io.Dir.cwd().deleteTree(io, root) catch {}; const w = std.Io.Dir.cwd(); try w.createDirPath(io, try std.fs.path.join(a, &.{ root, "models--Org--Flash/snapshots/rev1" })); try w.createDirPath(io, try std.fs.path.join(a, &.{ root, "models--Org--Flash/refs" })); @@ -50,15 +58,50 @@ test "models lists serve-able checkpoints and skips families Zig cannot serve" { try w.writeFile(io, .{ .sub_path = try std.fs.path.join(a, &.{ root, "models--Org--Flash/snapshots/rev1/config.json" }), .data = "{\"model_type\": \"qwen4_exp\", \"quantization\": {\"bits\": 6, \"group_size\": 32}}" }); try w.createDirPath(io, try std.fs.path.join(a, &.{ root, "models--Org--Draft/snapshots/rev2" })); try w.writeFile(io, .{ .sub_path = try std.fs.path.join(a, &.{ root, "models--Org--Draft/snapshots/rev2/config.json" }), .data = "{\"model_type\": \"gemma3n\"}" }); + return root; +} + +test "models lists serve-able checkpoints and skips families Zig cannot serve" { + var arena_state = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena_state.deinit(); + const a = arena_state.allocator(); + const io = std.testing.io; + const root = try syntheticCache(a, io); + defer std.Io.Dir.cwd().deleteTree(io, root) catch {}; var out: std.Io.Writer.Allocating = .init(a); - const code = try run(a, io, &out.writer, null, root); + const code = try run(a, io, &out.writer, null, root, false); try std.testing.expectEqual(@as(u8, 0), code); const listed = out.written(); try std.testing.expect(std.mem.indexOf(u8, listed, "Org/Flash: Qwen3.8 Flash Next (qwen4_exp)") != null); try std.testing.expect(std.mem.indexOf(u8, listed, "Org/Draft") == null); } +test "models --json prints the listed checkpoints as one document, and an absent cache as none" { + var arena_state = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena_state.deinit(); + const a = arena_state.allocator(); + const io = std.testing.io; + const root = try syntheticCache(a, io); + defer std.Io.Dir.cwd().deleteTree(io, root) catch {}; + + var out: std.Io.Writer.Allocating = .init(a); + try std.testing.expectEqual(@as(u8, 0), try run(a, io, &out.writer, null, root, true)); + var parsed = try std.json.parseFromSlice(std.json.Value, a, out.written(), .{}); + defer parsed.deinit(); + try std.testing.expectEqual(@as(i64, 1), parsed.value.object.get("schema").?.integer); + const models = parsed.value.object.get("models").?.array.items; + try std.testing.expectEqual(@as(usize, 1), models.len); + try std.testing.expectEqualStrings("Org/Flash", models[0].object.get("repo").?.string); + try std.testing.expectEqualStrings("Qwen3.8 Flash Next", models[0].object.get("family").?.string); + try std.testing.expectEqualStrings("qwen4_exp", models[0].object.get("model_type").?.string); + + var none: std.Io.Writer.Allocating = .init(a); + const missing = try std.fs.path.join(a, &.{ root, "missing" }); + try std.testing.expectEqual(@as(u8, 0), try run(a, io, &none.writer, null, missing, true)); + try std.testing.expectEqualStrings("{\"schema\":1,\"models\":[]}\n", none.written()); +} + fn has(text: []const u8, needle: []const u8) bool { if (std.mem.indexOf(u8, text, needle) != null) return true; std.debug.print("missing \"{s}\" in:\n{s}\n", .{ needle, text }); diff --git a/zig/src/server/main.zig b/zig/src/server/main.zig index ee42dc2b2a..5332e0ea07 100644 --- a/zig/src/server/main.zig +++ b/zig/src/server/main.zig @@ -57,7 +57,7 @@ pub fn main(init: std.process.Init) !u8 { return 0; } if (argv.len == 2 and (std.mem.eql(u8, argv[1], "--help") or std.mem.eql(u8, argv[1], "-h"))) { - try std.Io.File.stdout().writeStreamingAll(io, "usage: tensorfold-native --version | capabilities --json | models | info MODEL | pull REPO[@REVISION] | slide rollback FOLDER | slide audit [flags] | serve MODEL [flags] | benchmark MODEL [flags]\n" ++ usage_line); + try std.Io.File.stdout().writeStreamingAll(io, "usage: tensorfold-native --version | capabilities --json | models [--json] | info MODEL | pull REPO[@REVISION] | slide rollback FOLDER | slide audit [flags] | serve MODEL [flags] | benchmark MODEL [flags]\n" ++ usage_line); return 0; } if (argv.len >= 3 and std.mem.eql(u8, argv[1], "serve")) { @@ -87,7 +87,7 @@ pub fn main(init: std.process.Init) !u8 { } else return benchmark.main(init, argv[2..], engines.capabilities(a)); } if (argv.len < 2 or !std.mem.eql(u8, argv[1], "serve")) { - std.debug.print("usage: tensorfold-native capabilities --json | models | info MODEL | pull REPO[@REVISION] | slide rollback FOLDER | slide audit [flags] | serve MODEL [flags] | benchmark MODEL [flags]\n", .{}); + std.debug.print("usage: tensorfold-native capabilities --json | models [--json] | info MODEL | pull REPO[@REVISION] | slide rollback FOLDER | slide audit [flags] | serve MODEL [flags] | benchmark MODEL [flags]\n", .{}); return 2; } const started = std.Io.Clock.awake.now(io).toNanoseconds(); From 1fb70c1355b27cac756caf37e972aefd7751ee15 Mon Sep 17 00:00:00 2001 From: devopsfortunato <57877156+devopsfortunato@users.noreply.github.com> Date: Thu, 5 Nov 2026 12:17:09 +0100 Subject: [PATCH 33/36] engine api: an optional image capability beside speech (#1573) --- CHANGELOG.md | 1 + zig/src/core/engine_api.zig | 27 ++-- zig/src/core/engine_api_test.zig | 1 + zig/src/core/image_api.zig | 24 ++++ zig/src/core/image_api_test.zig | 219 +++++++++++++++++++++++++++++++ 5 files changed, 260 insertions(+), 12 deletions(-) create mode 100644 zig/src/core/image_api.zig create mode 100644 zig/src/core/image_api_test.zig diff --git a/CHANGELOG.md b/CHANGELOG.md index 4a40563961..b2856928bf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -255,6 +255,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Flash Next on Metal reports a filling prompt's progress: the dashboard shows what's left, and `TENSORFOLD_PREFILL_ORDER=srpt` works from a server's first long prompt (#2194). - CI: `zig` and `zig-windows` keep main's build cache instead of rebuilding from nothing, every check still reruns, and the quick checks run first. +- The engine API carries an optional image capability beside speech: `Info.image` states what an engine's image host serves and a defaulted `VTable.submit_image` hands back one encoded PNG a request through the caller's own list and allocator, so an engine without the entry is unchanged (#1573). - HIP: the `tf_sum_partials` records follow ROCm 10.1's hipcc, taking 8 limit problems off the contract (#2233). Thanks to @AbrahamCete. diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index ad99495d00..70ad410097 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -27,10 +27,7 @@ pub const seedFor = lanes.sampling.seedFor; pub const LogprobRow = lanes.LogprobRow; /// The request's stop strings, checked on the engine's thread after each committed token (the server decodes the tail). -pub const Stop = struct { - ctx: *anyopaque, - check: *const fn (ctx: *anyopaque, emitted: []const u32) bool, -}; +pub const Stop = struct { ctx: *anyopaque, check: *const fn (ctx: *anyopaque, emitted: []const u32) bool }; pub const CallGate = @import("engine_request_controls.zig").CallGate; pub const Structure = @import("engine_request_controls.zig").Structure; @@ -124,6 +121,10 @@ pub const SpeechRequest = @import("speech_api.zig").SpeechRequest; pub const SpeechResult = @import("speech_api.zig").SpeechResult; pub const SpeechInfo = @import("speech_api.zig").SpeechInfo; pub const SpeechSubmitError = error{ Closed, Busy, InvalidRequest, Unsupported }; +pub const ImageRequest = @import("image_api.zig").ImageRequest; +pub const ImageResult = @import("image_api.zig").ImageResult; +pub const ImageInfo = @import("image_api.zig").ImageInfo; +pub const ImageSubmitError = error{ Closed, Busy, InvalidRequest, Unsupported }; pub const Reason = enum { stop, length, cancelled, failed, silence }; pub const speech_guard = @import("speech_guard.zig"); @@ -165,6 +166,8 @@ pub const Info = struct { embeddings: ?lanes.embedding.Info = null, readouts: ?lanes.readout.Info = null, // non-null: the engine answers readout requests and nothing else speech: ?SpeechInfo = null, + /// The engine serves one picture a request through `submit_image`; null: no image capability. + image: ?ImageInfo = null, lanes: u32 = 1, /// Server preflight prompt-plus-reply limit (0: the engine owns admission). context_window: u32 = 0, @@ -241,10 +244,7 @@ pub const Info = struct { } }; -pub const Admit = struct { - ctx: *anyopaque, - fits: *const fn (ctx: *anyopaque, prompt: u32, max_tokens: u32) bool, -}; +pub const Admit = struct { ctx: *anyopaque, fits: *const fn (ctx: *anyopaque, prompt: u32, max_tokens: u32) bool }; /// The bytes the lane host may promise at once, and what one request's caches cost at most (core/kv_ledger.zig). pub const Ledger = struct { @@ -376,10 +376,7 @@ pub const Opened = struct { pub const Memory = struct { active: u64 = 0, cache: u64 = 0, peak: u64 = 0, pinned: u64 = 0, limit: ?u64 = null }; /// A backend's words for its own refusals (a request it cannot serve); null: the error's name. -pub const Explain = struct { - ctx: ?*anyopaque = null, - text: *const fn (ctx: ?*anyopaque, err: anyerror) ?[]const u8, -}; +pub const Explain = struct { ctx: ?*anyopaque = null, text: *const fn (ctx: ?*anyopaque, err: anyerror) ?[]const u8 }; /// A backend's own memory counts for LaneHost; read from HTTP threads, so it must not touch the GPU. pub const MemorySource = struct { @@ -424,6 +421,7 @@ pub const Engine = struct { info: *const fn (ctx: *anyopaque) Info, submit: *const fn (ctx: *anyopaque, id: Id, request: *const Request, sink: Sink) SubmitError!void, submit_speech: ?*const fn (*anyopaque, Id, *const SpeechRequest, Sink) SpeechSubmitError!void = null, + submit_image: ?*const fn (*anyopaque, Id, *const ImageRequest, Sink) ImageSubmitError!void = null, /// Ends a request between rounds; its ``finished`` (cancelled) still arrives. Unknown ids are ignored. cancel: *const fn (ctx: *anyopaque, id: Id) void, status: *const fn (ctx: *anyopaque, out: *Status, stream_tokens: []u32) void, @@ -471,6 +469,11 @@ pub const Engine = struct { const f = e.vtable.submit_speech orelse return error.Unsupported; return f(e.ctx, id, request, sink); } + /// The image hook, or Unsupported when this engine serves no picture. + pub fn submitImage(e: Engine, id: Id, request: *const ImageRequest, sink: Sink) ImageSubmitError!void { + const f = e.vtable.submit_image orelse return error.Unsupported; + return f(e.ctx, id, request, sink); + } pub fn cancel(e: Engine, id: Id) void { e.vtable.cancel(e.ctx, id); } diff --git a/zig/src/core/engine_api_test.zig b/zig/src/core/engine_api_test.zig index 030fd852a2..74a4aadc51 100644 --- a/zig/src/core/engine_api_test.zig +++ b/zig/src/core/engine_api_test.zig @@ -14,4 +14,5 @@ test { _ = @import("prompt_imprint.zig"); _ = @import("keepalive.zig"); _ = @import("originals_test.zig"); + _ = @import("image_api_test.zig"); } diff --git a/zig/src/core/image_api.zig b/zig/src/core/image_api.zig new file mode 100644 index 0000000000..e67fee21c7 --- /dev/null +++ b/zig/src/core/image_api.zig @@ -0,0 +1,24 @@ +//! The image model's own request and reply types, beside the speech ones in engine_api.zig. +const std = @import("std"); + +/// One picture: the encoded PNG's byte count and why the pass ended; the bytes are in the request's `png`. +pub const ImageResult = struct { bytes: usize = 0, reason: @import("engine_api.zig").Reason = .stop }; + +/// What the engine's image host serves, fixed once it has loaded: the host's stated limits, each one non-zero. +pub const ImageInfo = struct { max_side: u16, max_pixels: u32, max_steps: u32, max_prompt_tokens: u32 }; + +/// The one picture a request carries, caller-owned until `finished`. +pub const ImageRequest = struct { + prompt: []const u32, + negative: []const u32 = &.{}, + width: u16, + height: u16, + steps: u32, + guidance: f32, + seed: u64, + result: *ImageResult, + /// The engine appends its one PNG here with `alloc`; the caller frees it after `finished`. + png: *std.ArrayList(u8), + /// The allocator `png` lives in: the engine appends with it, the caller frees with it. + alloc: std.mem.Allocator, +}; diff --git a/zig/src/core/image_api_test.zig b/zig/src/core/image_api_test.zig new file mode 100644 index 0000000000..406299ae04 --- /dev/null +++ b/zig/src/core/image_api_test.zig @@ -0,0 +1,219 @@ +//! The engine API's image capability, against a fake host that records what it was handed and hands back one PNG. +const std = @import("std"); +const api = @import("engine_api.zig"); +const t = std.testing; +const a = t.allocator; + +/// The bytes a fake host appends: a PNG signature, so the caller can tell its own list from an untouched one. +const png = "\x89PNG\r\n\x1a\nIDAT"; + +/// What a fake host's entry does when it is reached. +const Mode = enum { ok, busy, invalid }; + +/// One fake image host: `serve` says whether it has the entry at all, `mode` what that entry answers. +const Stub = struct { + serve: bool = true, + mode: Mode = .ok, + reason: api.Reason = .stop, + limits: ?api.ImageInfo = null, + /// Every field the entry was handed, as the caller passed it. + seen: Seen = .{}, + /// The table `engine` hands out: `serve` decides whether it carries the image entry. + vtable: api.Engine.VTable = undefined, + + pub const Seen = struct { + calls: usize = 0, + id: api.Id = 0, + prompt: []const u32 = &.{}, + negative: []const u32 = &.{}, + width: u16 = 0, + height: u16 = 0, + steps: u32 = 0, + guidance: f32 = 0, + seed: u64 = 0, + }; + + fn engine(s: *Stub) api.Engine { + s.vtable = .{ + .info = info, + .submit = submit, + .cancel = cancel, + .status = status, + .memory = memory, + .submit_image = if (s.serve) submitImage else null, + }; + return .{ .ctx = s, .vtable = &s.vtable }; + } + + fn info(ctx: *anyopaque) api.Info { + const s: *Stub = @ptrCast(@alignCast(ctx)); + return .{ .text_generation = false, .image = s.limits }; + } + + fn submitImage(ctx: *anyopaque, id: api.Id, request: *const api.ImageRequest, sink: api.Sink) api.ImageSubmitError!void { + const s: *Stub = @ptrCast(@alignCast(ctx)); + s.seen = .{ + .calls = s.seen.calls + 1, + .id = id, + .prompt = request.prompt, + .negative = request.negative, + .width = request.width, + .height = request.height, + .steps = request.steps, + .guidance = request.guidance, + .seed = request.seed, + }; + switch (s.mode) { + .ok => {}, + .busy => return error.Busy, + .invalid => return error.InvalidRequest, + } + request.png.appendSlice(request.alloc, png) catch return error.Busy; // no room: the caller may come back + request.result.* = .{ .bytes = png.len, .reason = s.reason }; + sink.event(sink.ctx, id, &.{ .finished = .{ .reason = s.reason } }); + } + + fn submit(_: *anyopaque, _: api.Id, _: *const api.Request, _: api.Sink) api.SubmitError!void { + unreachable; + } + fn cancel(_: *anyopaque, _: api.Id) void {} + fn status(_: *anyopaque, _: *api.Status, _: []u32) void {} + fn memory(_: *anyopaque, _: bool) ?api.Memory { + return null; + } +}; + +/// The caller's end of the sink: one `finished` a picture, with the reason the engine stated. +const Record = struct { + events: usize = 0, + id: api.Id = 0, + reason: api.Reason = .stop, + + fn sink(r: *Record) api.Sink { + return .{ .ctx = r, .event = event }; + } + + fn event(ctx: *anyopaque, id: api.Id, e: *const api.Event) void { + const r: *Record = @ptrCast(@alignCast(ctx)); + switch (e.*) { + .finished => |f| { + r.events += 1; + r.id = id; + r.reason = f.reason; + }, + else => {}, + } + } +}; + +/// One caller-owned request: its prompt ids, its result, and the list its PNG lands in. +const Case = struct { + prompt: [3]u32 = .{ 11, 12, 13 }, + negative: [2]u32 = .{ 7, 8 }, + result: api.ImageResult = .{}, + list: std.ArrayList(u8) = .empty, + + fn request(c: *Case) api.ImageRequest { + return .{ + .prompt = &c.prompt, + .negative = &c.negative, + .width = 1024, + .height = 768, + .steps = 8, + .guidance = 3.5, + .seed = 42, + .result = &c.result, + .png = &c.list, + .alloc = a, + }; + } +}; + +test "the image entry is reached with exactly the caller's fields" { + var stub = Stub{}; + var case = Case{}; + var record = Record{}; + const engine = stub.engine(); + var request = case.request(); + try engine.submitImage(9, &request, record.sink()); + defer case.list.deinit(a); + try t.expectEqual(@as(usize, 1), stub.seen.calls); + try t.expectEqual(@as(api.Id, 9), stub.seen.id); + try t.expectEqualSlices(u32, &.{ 11, 12, 13 }, stub.seen.prompt); + try t.expectEqualSlices(u32, &.{ 7, 8 }, stub.seen.negative); + try t.expectEqual(@as(u16, 1024), stub.seen.width); + try t.expectEqual(@as(u16, 768), stub.seen.height); + try t.expectEqual(@as(u32, 8), stub.seen.steps); + try t.expectEqual(@as(f32, 3.5), stub.seen.guidance); + try t.expectEqual(@as(u64, 42), stub.seen.seed); +} + +test "the PNG the engine appends is the caller's to free with its own allocator" { + var stub = Stub{}; + var case = Case{}; + var record = Record{}; + const engine = stub.engine(); + var request = case.request(); + try engine.submitImage(1, &request, record.sink()); + defer case.list.deinit(a); // the allocator the engine appended with: a mismatch or a leak fails here + try t.expectEqualSlices(u8, png, case.list.items); + try t.expectEqual(@as(usize, png.len), case.result.bytes); + try t.expectEqual(api.Reason.stop, case.result.reason); +} + +test "an engine without the image entry is unsupported" { + var stub = Stub{ .serve = false }; + var case = Case{}; + var record = Record{}; + const engine = stub.engine(); + var request = case.request(); + try t.expectError(error.Unsupported, engine.submitImage(1, &request, record.sink())); + try t.expectEqual(@as(usize, 0), stub.seen.calls); + try t.expectEqual(@as(usize, 0), case.list.items.len); + try t.expectEqual(@as(usize, 0), record.events); + try t.expect(engine.info().image == null); +} + +test "an engine error comes back to the caller unchanged" { + inline for (.{ Mode.busy, Mode.invalid }) |mode| { + var stub = Stub{ .mode = mode }; + var case = Case{}; + var record = Record{}; + const engine = stub.engine(); + var request = case.request(); + const want: api.ImageSubmitError = switch (mode) { + .busy => error.Busy, + .invalid => error.InvalidRequest, + .ok => unreachable, + }; + try t.expectError(want, engine.submitImage(3, &request, record.sink())); + try t.expectEqual(@as(usize, 0), case.list.items.len); + try t.expectEqual(@as(usize, 0), record.events); + } +} + +test "Info carries the engine's image limits exactly, and none by default" { + const limits = api.ImageInfo{ .max_side = 2048, .max_pixels = 4_194_304, .max_steps = 50, .max_prompt_tokens = 512 }; + var stub = Stub{ .limits = limits }; + const engine = stub.engine(); + try t.expectEqual(limits, engine.info().image.?); + try t.expectEqual(@as(u16, 2048), engine.info().image.?.max_side); + try t.expectEqual(@as(u32, 4_194_304), engine.info().image.?.max_pixels); + try t.expectEqual(@as(u32, 50), engine.info().image.?.max_steps); + try t.expectEqual(@as(u32, 512), engine.info().image.?.max_prompt_tokens); + try t.expect((api.Info{}).image == null); +} + +test "the finished event the engine sends arrives at the caller's sink" { + var stub = Stub{ .reason = .length }; + var case = Case{}; + var record = Record{}; + const engine = stub.engine(); + var request = case.request(); + try engine.submitImage(12, &request, record.sink()); + defer case.list.deinit(a); + try t.expectEqual(@as(usize, 1), record.events); + try t.expectEqual(@as(api.Id, 12), record.id); + try t.expectEqual(api.Reason.length, record.reason); + try t.expectEqual(api.Reason.length, case.result.reason); +} From 9c3dbb29eca17797cc160b8ba771ed45476b603a Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 5 Nov 2026 18:39:41 -0600 Subject: [PATCH 34/36] moss_image: the transformer gated block by block against float32, on Metal and CUDA tf-moss-image-blocks runs one step's branch through Model.steps (on both backends): the embedding, each kept block alone on a float32 run's stream rounded to bf16, then the final layer. Each must be within twice torch bf16's distance from float32 in cosine error and worst row, and a second run must repeat the first byte for byte. moss_image_case.py blocks records the float32 stream and torch bf16's steps alone, on the case's device and on the CPU; the limit is the farther of the two (reference.wider), since torch's own bf16 runs disagree at a massive-activation row (block 0 of the fox's step 1, negative branch: worst row 0.00687 on MPS, 0.00960 on the CPU). tf-moss-image's whole forwards are reported beside torch bf16's and not gated, as the encoder's whole prompt is not; a repeat that differs still fails it. Seven steps' branches: 305 of 305 gates pass on Metal (at most 0.50 of the limit) and on an RTX 5090 (at most 0.60), each run repeated byte for byte; a row planted 5% off in block 20 fails at blocks 20 and 21. Model.forward's modulation and embedding moved into helpers that Model.steps shares; both cases' latents are byte-identical before and after on both backends. --- docs/recipes/moss-image-cuda.md | 23 ++++ docs/recipes/moss-image.md | 69 +++++++++-- tools/zig/moss_image_case.py | 80 ++++++++++++- zig/build/cuda.zig | 4 +- zig/build/metal_targets.zig | 1 + zig/src/families/moss_image/cuda.zig | 1 + zig/src/families/moss_image/cuda_model.zig | 65 +++++++++-- zig/src/families/moss_image/model.zig | 60 ++++++++-- zig/src/families/moss_image/reference.zig | 27 +++++ zig/tests/moss_image.zig | 2 +- zig/tests/moss_image_blocks.zig | 6 + zig/tests/moss_image_blocks_cuda.zig | 6 + zig/tests/moss_image_blocks_run.zig | 129 +++++++++++++++++++++ zig/tests/moss_image_cuda.zig | 2 +- zig/tests/moss_image_run.zig | 1 + 15 files changed, 443 insertions(+), 33 deletions(-) create mode 100644 zig/tests/moss_image_blocks.zig create mode 100644 zig/tests/moss_image_blocks_cuda.zig create mode 100644 zig/tests/moss_image_blocks_run.zig diff --git a/docs/recipes/moss-image-cuda.md b/docs/recipes/moss-image-cuda.md index fca601b390..de104a8f69 100644 --- a/docs/recipes/moss-image-cuda.md +++ b/docs/recipes/moss-image-cuda.md @@ -77,6 +77,16 @@ python3 tools/zig/moss_image_case.py compare case/reference.png case/zig.png cas The case tool needs diffusers 0.38 or later: the checkpoint's decoder config names `decoder_block_out_channels`. +`tf-moss-image` gates each whole forward element by element against the case's float32 run, and +`tf-moss-image-blocks` gates each block alone against the farther of torch bf16's two runs, both as on Metal +([moss-image.md](moss-image.md), "Run a case"). The block cases measured below were made on a Mac by +`moss_image_case.py blocks` (torch on MPS and on the CPU); a block case runs the whole transformer in float32, which +does not fit an RTX 5090. + +```sh +zig-out/bin/tf-moss-image-blocks MODEL_DIR/transformer case/blocks-23-negative.safetensors +``` + ## Measured Setup: @@ -114,6 +124,19 @@ velocities and float32 runs as recorded on the GB10): - **Image:** 44.01 dB PSNR from the pipeline's picture for the fox, 34.51 dB for the sticky note. The latents differ from the GB10's in their low bits: cuBLASLt picks different GEMM kernels on each GPU. - **Repeat:** identical bytes, as on the GB10. +- **Each block alone**, on the Metal doc's seven block cases (made on the Mac: torch bf16 on MPS and on the CPU): all + 305 gates pass, at a mean of 0.50 of the limit and at most 0.60, and a second run repeats the first byte for byte. + At block 0 of the fox's step 1, negative branch, ours is 5.70e-6 from float32 in cosine error and 0.00960 in its + worst row (the negative prompt's seventh row), as torch bf16 on the CPU is (6.58e-6, 0.00960); torch on MPS is + 2.28e-6 and 0.00687. Each of our operations there, run from our own inputs, matches torch's bf16 code to within the + order of its sums; that row turns small rounding differences into larger ones, on torch's devices as on ours. +- **Each whole forward**, over each case's 100: our cosine error from float32 is 1.00 times torch bf16's on average + for the fox (0.87 to 1.13) and 1.01 times for the sticky note (0.85 to 1.55); our worst row 1.01 times (0.34 to 1.63) + and 1.02 times (0.66 to 1.98). +- **The teacher gate** (`tf-moss-image`, exit 1, its bound from the GB10's torch bf16 run) fails all 100 branch-steps + of both cases, as on Metal ([moss-image.md](moss-image.md), Measured, where torch's own bf16 run on the CPU fails it + too). Worst element over its bound, min / median / max: fox 161 / 432 / 2946, sticky note 115 / 459 / 7921. Two + runs give identical latents. - **Speed**, the GPU otherwise idle: - torch bf16 (torch 2.13.0, the pipeline's own forward on this card): 0.307 s a forward for the fox, 0.309 s for the sticky note. diff --git a/docs/recipes/moss-image.md b/docs/recipes/moss-image.md index bfc0f407a8..94cf4c6114 100644 --- a/docs/recipes/moss-image.md +++ b/docs/recipes/moss-image.md @@ -179,14 +179,19 @@ The case holds: - the timesteps and noise levels; - the starting latents; - the publisher's pipeline's input and both velocities at every step; -- with `--fp32-steps N`, a float32 run of the same transformer for the first N steps. - -`tf-moss-image` first runs both branches of each step on the pipeline's own input and compares the velocities. It then -runs the whole guided schedule from the starting latents, times it, frees the transformer, and decodes the final -latents with the checkpoint's `vae/` decoder into `picture.png`. A case made by `start` holds only the prompts, the -schedule and the noise; `tf-moss-image` then just makes the picture. With `--prompt` (and `--negative`, else the empty -negative prompt), it encodes both prompts natively in place of the case's rows; `start --no-encoder` leaves them out of -the case, and torch never loads the encoder. +- a float32 run of the same transformer at every step (`--fp32-steps N` keeps only the first N, which leaves the + teacher gate below inconclusive). + +`tf-moss-image` first runs both branches of each step on the pipeline's own input and compares the velocities with torch +bf16's (max |diff|, cosine, worst row) and with the case's float32 ones. The teacher gate checks every element: its +distance from float32 must be within twice the larger of torch bf16's own distance there and half a bf16 ULP at the +float32 value. A case without float32 velocities for both branches at every step is inconclusive (exit 2), and a step +run twice must repeat byte for byte. On this doc's cases the gate fails every step, as torch's own bf16 run on the CPU +does (see Measured); the bound's form is still being settled in #1529. It then runs the whole guided schedule from the +starting latents, times it, frees the transformer, and decodes the final latents with the checkpoint's `vae/` decoder +into `picture.png`. A case made by `start` holds only the prompts, the schedule and the noise; `tf-moss-image` then just +makes the picture. With `--prompt` (and `--negative`, else the empty negative prompt), it encodes both prompts natively +in place of the case's rows; `start --no-encoder` leaves them out of the case, and torch never loads the encoder. To check the decoder alone against a float64 decode of the same latents: @@ -200,6 +205,22 @@ zig-out/bin/tf-moss-image-decoder MODEL_DIR case/latents.f32 1024 1024 case/deco Ours must be within twice torch float32's distance from float64, in cosine error and in its worst row of pixels, and a second decode must repeat the first byte for byte; it fails otherwise. +To gate one step's transformer block by block (as the encoder's are), from a `case`: + +```sh +python3 tools/zig/moss_image_case.py blocks MODEL_DIR case/ --step 23 --branch negative +zig build tf-moss-image-blocks -Dcpu=apple_m1 +zig-out/bin/tf-moss-image-blocks MODEL_DIR/transformer case/blocks-23-negative.safetensors +``` + +The case tool runs that step's branch in float32, keeping the stream into each block and after the last, then in bf16 +with every block and the final layer handed the float32 stream rounded to bf16, on the case's device and again on the +CPU. `Model.steps` runs ours the same way: the embedding, each block but the skipped one alone, and the final layer +alone. Each must be within twice torch bf16's distance from float32 (the farther of its two runs, or bf16 rounding's, +when larger), in cosine error and in its worst row, and a second run must repeat the first byte for byte; it fails +otherwise. Two torch runs set the limit because one device's bf16 can land well inside torch's own spread (see +Measured). + On MPS, which has no float64, the case tool computes the time sinusoid on the host in float64 and moves its float32 result to the GPU; every other step of the publisher's code runs unchanged. @@ -237,6 +258,38 @@ Results: - Against a float32 run of the same transformer on the same inputs (the second case, every step): mean 1 - cosine 1.602e-4 for the prompt branch, as for torch bf16's own (1.602e-4); 1.735e-4 against torch's 1.715e-4 for the negative branch. Ours is the closer of the two at 25 and 27 of the 50 steps. +- **Each block alone** (`tf-moss-image-blocks`), on seven steps' branches: the fox's steps 1 and 25 and the sticky + note's step 23, both branches each, and the sticky note's step 50, the prompt's. Each gives the embedding, every + block but the skipped one and the final layer the float32 run's stream in bf16; torch bf16 ran each on MPS and on + the CPU: + - All 305 gates pass, at a mean of 0.50 of the limit and at most 0.50. Our worst row is 0.99 to 1.00 times the + farther torch run's on average in each of the seven, and a second run repeats the first byte for byte. + - torch's two runs are within 0.84 to 1.20 times each other in worst row, and 0.95 to 1.05 in cosine error, at all + but one step: block 0 of the fox's step 1, negative branch, where the negative prompt's seventh row becomes a + massive activation (one channel near 16,500). There torch on the CPU is 6.58e-6 from float32 in cosine error and + 0.00960 in that row; torch on MPS 2.28e-6 and 0.00687, and ours on Metal the same as MPS. A gate on one device's + run alone would fail a result as close as the other device's. + - A row planted 5% off in block 20's output, in the float32 stream and in both torch runs, fails at blocks 20 and + 21, with an identical repeat. +- **Each whole forward**, over each case's 100: our cosine error from float32 is 1.00 times torch bf16's on average + for the fox (0.92 to 1.08) and 1.01 times for the sticky note (0.76 to 1.26); our worst row 1.01 times (0.64 to + 1.48) and 1.03 times (0.37 to 2.55). At the sticky note's step 23, negative branch, torch's own bf16 run on the CPU + has a worst row 1.37 times its MPS run's (0.0664 and 0.0486) and ours 2.55 times (0.1238, row 313): the worst row + of 42 bf16 blocks in a row moves by more than torch's own bf16 runs agree. +- **The teacher gate** (`tf-moss-image`, exit 1) fails all 100 branch-steps of both cases. Worst element over its + bound, min / median / max: fox 146 / 509 / 7206, sticky note 168 / 435 / 6018. A step run twice repeats byte for + byte. #2291's head without this commit gives the same latents and the same worst element at every step. At the + sticky note's step 23, negative branch, ours is 0.0173 from float32 against a bound of 6.1e-5 there. torch's own + bf16 run of that step on the CPU, scored against the same bound (from its MPS run), fails too: + + | step 23, negative branch, against float32 | worst error / bound | elements over the bound | median error / bound | + |---|---:|---:|---:| + | torch bf16, CPU | 353 | 16.6% | 0.50 | + | ours, Metal | 284 | 22.3% | 0.50 | + + Each element's bound is twice the MPS run's own error at that element, so it is tight wherever MPS happened to + round close to float32; any other bf16 run of 42 blocks, torch's on the CPU included, lands outside it at a sizeable + share of the elements. - **Image:** our latents through the pipeline's decoder are 48.52 dB PSNR from the pipeline's own picture for the fox and 28.93 dB for the sticky note. Fifty guided steps magnify bf16 rounding into different fine detail, so the ruler is a float32 run of the whole schedule. There are two rulers, over the same 8 cases (both prompts, 4 seeds each): diff --git a/tools/zig/moss_image_case.py b/tools/zig/moss_image_case.py index 2ed153c51f..471d0e25d1 100644 --- a/tools/zig/moss_image_case.py +++ b/tools/zig/moss_image_case.py @@ -26,6 +26,10 @@ encode MODEL_DIR OUT_DIR [--prompt-file F] the prompt encoder alone: token ids, embedding rows, and every layer's output and the final rows in bf16 and in float32, and each step alone in bf16 on the float32 run's input + blocks MODEL_DIR CASE_DIR --step N --branch B one step of a `case` in float32, the stream into every block + kept, and torch bf16's embedding, each block and the final layer + alone on that stream in bf16, on --device and on the CPU: + CASE_DIR/blocks-N-B.safetensors, for tf-moss-image-blocks compare A.png B.png [SIDE.png] PSNR and mean absolute difference of two pictures, optionally side by side noise OUT.json torch.randn's starting latents for many seeds and sizes, as sha256s (x86) schedule MODEL_DIR OUT.json MOSSImageScheduler's timesteps and sigmas for many step counts and configs @@ -270,6 +274,74 @@ def vae(args) -> int: return 0 +def blocks(args) -> int: + """One step's branch of a `case` in float32, keeping each block's input; then each step alone in bf16, two ways.""" + import torch + from safetensors.torch import load_file, save_file + + case = load_file(str(Path(args.case_dir) / "case.safetensors")) + if "inputs" not in case: + sys.exit("blocks needs a case made by `case`, which holds each step's input") + latents = case["inputs"][args.step - 1].permute(2, 0, 1)[None] + timestep = case["timesteps"][args.step - 1] + text = case[args.branch][None] + skip = [22] if args.branch == "negative" else [] + pipe, _ = pipeline(args.model_dir, args.device, text=False) + tr = pipe.transformer + kept = [layer for i, layer in enumerate(tr.layers) if i not in skip] + image = latents.shape[2] * latents.shape[3] + + def forward(device, dtype, pre, post): + handles = [m.register_forward_pre_hook(fn, with_kwargs=True) for m, fn in pre] + handles += [m.register_forward_hook(fn) for m, fn in post] + with torch.no_grad(): + out = tr(latents.to(device), timestep, text.to(device=device, dtype=dtype), skip_layers=skip).sample + for h in handles: + h.remove() + return rows(out) + + def record(into): + def hook(module, a, kw): + into.append(a[0][0].float().cpu()) + return hook + + def output(into): + return lambda module, a, out: into.append(out[0].float().cpu()) + + # float32 throughout: the stream into each block kept, then the last block's output + tr.float() + stream = [] + velocity_fp32 = forward(args.device, torch.float32, [(m, record(stream)) for m in kept], [(kept[-1], output(stream))]) + + def alone(device): + """bf16, each step alone: the embedding, every block and the final layer, given the float32 stream in bf16.""" + tr.to(device=device, dtype=torch.bfloat16) + embedded, outs = [], [] + + def feed(k): + def hook(module, a, kw): + if k == 0: + embedded.append(a[0][0].float().cpu()) + x = stream[k].to(device=device, dtype=torch.bfloat16)[None] + return (x[:, :image] if k == len(kept) else x,) + tuple(a[1:]), kw + return hook + + velocity = forward(device, torch.bfloat16, [(m, feed(k)) for k, m in enumerate(kept + [tr.final_norm_video])], + [(m, output(outs)) for m in kept]) + return torch.stack(embedded + outs).to(torch.bfloat16).contiguous(), velocity + + steps, velocity_alone = alone(args.device) + steps_cpu, velocity_alone_cpu = alone("cpu") + out = Path(args.case_dir) / f"blocks-{args.step}-{args.branch}.safetensors" + save_file({"step": torch.tensor([args.step], dtype=torch.int64), "skip": torch.tensor(skip or [-1], dtype=torch.int64), + "timestep": timestep.reshape(1).float(), "latents": case["inputs"][args.step - 1].contiguous(), + "prompt": text[0].contiguous(), "stream": torch.stack(stream).contiguous(), "alone": steps, + "alone_cpu": steps_cpu, "velocity_fp32": velocity_fp32, "velocity_alone": velocity_alone, + "velocity_alone_cpu": velocity_alone_cpu}, str(out)) + print(json.dumps({"out": str(out), "blocks": len(kept), "rows": int(stream[0].shape[0])})) + return 0 + + def compare(args) -> int: import numpy as np from PIL import Image @@ -475,6 +547,12 @@ def main() -> int: v.add_argument("--device", default="mps", help="where the float32 decode runs, as the pipeline would") v.add_argument("--width", type=int, default=1024) v.add_argument("--height", type=int, default=1024) + b = sub.add_parser("blocks") + b.add_argument("model_dir") + b.add_argument("case_dir") + b.add_argument("--step", type=int, required=True, help="the step, from 1") + b.add_argument("--branch", choices=("prompt", "negative"), required=True) + b.add_argument("--device", default="mps") p = sub.add_parser("compare") p.add_argument("a") p.add_argument("b") @@ -486,7 +564,7 @@ def main() -> int: q.add_argument("out") args = parser.parse_args() modes = {"case": case, "start": start, "exact": exact, "decode": decode, "encode": encode, "vae": vae, - "compare": compare, "noise": noise, "schedule": schedule} + "blocks": blocks, "compare": compare, "noise": noise, "schedule": schedule} return modes[args.mode](args) diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index d7a9fefb46..89e4bdd6e5 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -136,10 +136,12 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu const q35_install = b.addInstallArtifact(b.addExecutable(.{ .name = "tf-qwen35-cuda", .root_module = q35 }), .{}); b.getInstallStep().dependOn(&q35_install.step); b.step("qwen35-cuda", "Build the Qwen CUDA forward and raw state receipt tool").dependOn(&q35_install.step); - // tf-moss-image: MOSS-Image's transformer on a reference case (tools/zig/moss_image_case.py), compared and timed + // tf-moss-image (a case, gated and timed) and tf-moss-image-blocks (each block alone): tools/zig/moss_image_case.py const moss = b.createModule(.{ .root_source_file = b.path("zig/src/families/moss_image/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core } } }); const moss_tool = b.createModule(.{ .root_source_file = b.path("zig/tests/moss_image_cuda.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core }, .{ .name = "moss_image", .module = moss } } }); b.installArtifact(b.addExecutable(.{ .name = "tf-moss-image", .root_module = moss_tool })); + const moss_blocks = b.createModule(.{ .root_source_file = b.path("zig/tests/moss_image_blocks_cuda.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core }, .{ .name = "moss_image", .module = moss } } }); + b.installArtifact(b.addExecutable(.{ .name = "tf-moss-image-blocks", .root_module = moss_blocks })); const d3_tool = b.createModule(.{ .root_source_file = b.path("zig/src/families/drama3/generate.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "core", .module = mods.core }, .{ .name = "drama3", .module = @import("drama3_host.zig").module(b, target, optimize, cuda, mods.core) } } }); const d3_exe = b.addExecutable(.{ .name = "drama3-generate", .root_module = d3_tool }); b.step("drama3-generate", "Drama generation and replay controls on CUDA").dependOn(&b.addInstallArtifact(d3_exe, .{}).step); diff --git a/zig/build/metal_targets.zig b/zig/build/metal_targets.zig index 4126269650..da02394280 100644 --- a/zig/build/metal_targets.zig +++ b/zig/build/metal_targets.zig @@ -336,6 +336,7 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s .{ .name = "tf-moss-image", .path = "zig/tests/moss_image.zig", .about = "MOSS-Image on a reference case: each step's two branches compared, then the guided schedule timed and decoded" }, .{ .name = "tf-moss-image-encoder", .path = "zig/tests/moss_image_encoder.zig", .about = "MOSS-Image's prompt encoder on a reference case: the ids checked, each step gated against float32, then timed" }, .{ .name = "tf-moss-image-decoder", .path = "zig/tests/moss_image_decoder.zig", .about = "MOSS-Image's image decoder on latents, gated against float64 beside torch float32, then timed" }, + .{ .name = "tf-moss-image-blocks", .path = "zig/tests/moss_image_blocks.zig", .about = "MOSS-Image's transformer on one step, each block alone gated against float32, repeated byte for byte" }, .{ .name = "tf-glm-attn-probe", .path = "zig/tests/glm_attn_probe.zig", .about = "GLM-5.3-Flash's dense latent attention kernels on given inputs (scores, probabilities, outputs)" }, .{ .name = "tf-moe-bench", .path = "zig/tests/moe_bench.zig", .about = "Routed-expert prompt kernels from .metal files on a skewed routing: chained speed, bits against the first" }, .{ .name = "tf-mm-bench", .path = "zig/tests/mm_bench.zig", .about = "Prompt matmul kernels from .metal files: chained speed at one shape, bits against the first, error against fp64" }, diff --git a/zig/src/families/moss_image/cuda.zig b/zig/src/families/moss_image/cuda.zig index 48724e0675..1faba72cd0 100644 --- a/zig/src/families/moss_image/cuda.zig +++ b/zig/src/families/moss_image/cuda.zig @@ -1,6 +1,7 @@ //! MOSS-Image's transformer on CUDA: a denoiser driven by tf-moss-image, outside the lane engine. pub const config = @import("config.zig"); pub const host = @import("host.zig"); +pub const reference = @import("reference.zig"); pub const Model = @import("cuda_model.zig").Model; pub const Prompt = @import("cuda_model.zig").Prompt; pub const seconds = @import("cuda_model.zig").seconds; diff --git a/zig/src/families/moss_image/cuda_model.zig b/zig/src/families/moss_image/cuda_model.zig index 46476082b9..6c794d70a7 100644 --- a/zig/src/families/moss_image/cuda_model.zig +++ b/zig/src/families/moss_image/cuda_model.zig @@ -220,19 +220,10 @@ pub const Model = struct { if (latents.len != R * cfg.channels or velocity.len != latents.len) return error.LatentRows; if (skip) |s| if (s >= cfg.layers) return error.BadSkipLayer; const started = seconds(); - var spent: f64 = 0; - if (t.timestep != timestep) { - var vectors: [host.vectors][H]u16 = undefined; - t.timestep = null; - try host.modulation(&vectors, t.gpa, t.w.modulation, timestep); - try t.upload(t.mods, std.mem.sliceAsBytes(&vectors)); - t.timestep = timestep; - spent = seconds() - started; - } + const spent = try t.modulate(timestep); const halves = try t.gpa.alloc(u16, latents.len); defer t.gpa.free(halves); - for (halves, latents) |*o, v| o.* = host.bf16(v); - try t.upload(t.lat, std.mem.sliceAsBytes(halves)); + try t.embed(latents, halves); const N = R + p.rows; try t.marks[0].record(t.stream); try t.d.check(t.d.api.cuMemcpyDtoDAsync_v2(t.x + R * H * 2, p.text.ptr, p.rows * H * 2, t.stream.handle), "cuMemcpyDtoDAsync"); @@ -247,6 +238,58 @@ pub const Model = struct { return .{ .wall = seconds() - started, .gpu = @as(f64, try cuda.Event.elapsedMs(t.marks[0], t.marks[1])) / 1000, .host = spent }; } + /// For checks: the embedding, each kept block alone on its block of `inputs`, then the final layer on the last. + pub fn steps(t: *Model, p: *const Prompt, latents: []const f32, timestep: f32, skip: ?usize, inputs: []const u16, rows: []u16, velocity: []f32) !void { + const R = t.image; + const N = R + p.rows; + const kept: usize = cfg.layers - @as(usize, @intFromBool(skip != null)); + if (latents.len != R * cfg.channels or velocity.len != latents.len) return error.LatentRows; + if (skip) |s| if (s >= cfg.layers) return error.BadSkipLayer; + if (inputs.len != (kept + 1) * N * H or rows.len != inputs.len) return error.StreamRows; + _ = try t.modulate(timestep); + const halves = try t.gpa.alloc(u16, latents.len); + defer t.gpa.free(halves); + try t.embed(latents, halves); + try t.d.check(t.d.api.cuMemcpyDtoDAsync_v2(t.x + R * H * 2, p.text.ptr, p.rows * H * 2, t.stream.handle), "cuMemcpyDtoDAsync"); + try t.linearBias(t.lat, t.w.video_in, t.w.video_in_b, t.x, R, H, cfg.channels); + try t.download(rows[0 .. N * H], t.x); + var k: usize = 0; + for (t.w.blocks, 0..) |b, i| if (skip != i) { + try t.upload(t.x, std.mem.sliceAsBytes(inputs[k * N * H ..][0 .. N * H])); + try t.block(b, N, p.cos(), p.sin(R)); + k += 1; + try t.download(rows[k * N * H ..][0 .. N * H], t.x); + }; + try t.upload(t.x, std.mem.sliceAsBytes(inputs[kept * N * H ..][0 .. N * H])); + try t.norm(R, t.w.final_norm, .{ .scale = .scale_out, .shift = .shift_out }); + try t.linear(t.xn, t.w.final_out, t.vel, R, cfg.channels, H); + try t.download(halves, t.vel); + for (velocity, halves) |*o, v| o.* = host.widen(v); + } + + /// The shared modulation for `timestep` on the device, made on the host unless it holds it; the seconds it took. + fn modulate(t: *Model, timestep: f32) !f64 { + if (t.timestep == timestep) return 0; + const started = seconds(); + var vectors: [host.vectors][H]u16 = undefined; + t.timestep = null; + try host.modulation(&vectors, t.gpa, t.w.modulation, timestep); + try t.upload(t.mods, std.mem.sliceAsBytes(&vectors)); + t.timestep = timestep; + return seconds() - started; + } + + /// The latents in bf16 (through `halves`) on the device, ready for the input projection. + fn embed(t: *const Model, latents: []const f32, halves: []u16) !void { + for (halves, latents) |*o, v| o.* = host.bf16(v); + try t.upload(t.lat, std.mem.sliceAsBytes(halves)); + } + + fn download(t: *const Model, into: []u16, at: u64) !void { + try t.d.check(t.d.api.cuMemcpyDtoHAsync_v2(into.ptr, at, into.len * 2, t.stream.handle), "cuMemcpyDtoHAsync"); + try t.stream.synchronize(); + } + const Modulated = struct { scale: host.Vector, shift: host.Vector }; /// xn = the RMSNorm of the stream's `rows` rows with weight `w`, then modulated when given. diff --git a/zig/src/families/moss_image/model.zig b/zig/src/families/moss_image/model.zig index 374d8cb0c3..ddbc90dab5 100644 --- a/zig/src/families/moss_image/model.zig +++ b/zig/src/families/moss_image/model.zig @@ -189,17 +189,9 @@ pub const Model = struct { if (latents.len != R * cfg.channels or velocity.len != latents.len) return error.LatentRows; if (skip) |s| if (s >= cfg.layers) return error.BadSkipLayer; const started = mtl.clock.seconds(); - var spent: f64 = 0; - if (t.timestep != timestep) { - const vectors: *[host.vectors][H]u16 = @ptrCast(@alignCast(t.mods.addr())); - t.timestep = null; - try host.modulation(vectors, t.gpa, t.w.modulation, timestep); - t.timestep = timestep; - spent = mtl.clock.seconds() - started; - } - for (halves(t.lat, latents.len), latents) |*o, v| o.* = host.bf16(v); + const spent = try t.modulate(timestep); + t.embed(p, latents); const N = R + p.rows; - @memcpy(halves(t.x.at(R * H * 2), p.rows * H), halves(.{ .buf = p.text }, p.rows * H)); const pool = mtl.objc.Pool.push(); defer pool.pop(); const cb = t.queue.commandBuffer(); @@ -213,6 +205,54 @@ pub const Model = struct { return timing; } + /// For checks: the embedding, each kept block alone on its block of `inputs`, then the final layer on the last. + pub fn steps(t: *Model, p: *const Prompt, latents: []const f32, timestep: f32, skip: ?usize, inputs: []const u16, rows: []u16, velocity: []f32) !void { + const R = t.image; + const N = R + p.rows; + const kept: usize = cfg.layers - @as(usize, @intFromBool(skip != null)); + if (latents.len != R * cfg.channels or velocity.len != latents.len) return error.LatentRows; + if (skip) |s| if (s >= cfg.layers) return error.BadSkipLayer; + if (inputs.len != (kept + 1) * N * H or rows.len != inputs.len) return error.StreamRows; + _ = try t.modulate(timestep); + t.embed(p, latents); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + var k: usize = 0; + while (k <= kept + 1) : (k += 1) { + if (k > 0) @memcpy(halves(t.x, N * H), inputs[(k - 1) * N * H ..][0 .. N * H]); + const cb = t.queue.commandBuffer(); + const e = cb.compute(.serial); + if (k == 0) { + t.k.linearBias(e, t.lat, t.w.video_in, t.w.video_in_b, t.sums, t.x, R, H, cfg.channels); + } else if (k <= kept) { + t.block(e, t.w.blocks[if (skip) |s| (if (k - 1 < s) k - 1 else k) else k - 1], N, p.cos(), p.sin(R)); + } else { + t.norm(e, R, t.w.final_norm, .{ .scale = .scale_out, .shift = .shift_out }); + t.k.linear(e, t.xn, t.w.final_out, t.vel, R, cfg.channels, H); + } + _ = try finish(cb, e, 0, 0); + if (k <= kept) @memcpy(rows[k * N * H ..][0 .. N * H], halves(t.x, N * H)); + } + for (velocity, halves(t.vel, velocity.len)) |*o, v| o.* = host.widen(v); + } + + /// The shared modulation for `timestep`, made on the host unless it already holds it; the seconds it took. + fn modulate(t: *Model, timestep: f32) !f64 { + if (t.timestep == timestep) return 0; + const started = mtl.clock.seconds(); + const vectors: *[host.vectors][H]u16 = @ptrCast(@alignCast(t.mods.addr())); + t.timestep = null; + try host.modulation(vectors, t.gpa, t.w.modulation, timestep); + t.timestep = timestep; + return mtl.clock.seconds() - started; + } + + /// The latents in bf16 and the prompt's rows after the image rows, ready for the input projection. + fn embed(t: *Model, p: *const Prompt, latents: []const f32) void { + for (halves(t.lat, latents.len), latents) |*o, v| o.* = host.bf16(v); + @memcpy(halves(t.x.at(t.image * H * 2), p.rows * H), halves(.{ .buf = p.text }, p.rows * H)); + } + fn finish(cb: mtl.CommandBuffer, e: mtl.ComputeEncoder, started: f64, spent: f64) !Timing { e.end(); cb.commit(); diff --git a/zig/src/families/moss_image/reference.zig b/zig/src/families/moss_image/reference.zig index 0e80be3236..9f5cfc187d 100644 --- a/zig/src/families/moss_image/reference.zig +++ b/zig/src/families/moss_image/reference.zig @@ -53,6 +53,12 @@ pub fn gate(ours: Match, reference: Match, rounded: Match) Verdict { return within(ours, .{ .max_abs = 0, .cosine = @min(reference.cosine, rounded.cosine), .worst_row = @max(reference.worst_row, rounded.worst_row), .worst_at = 0 }); } +/// The farther of two reference runs from exact, in each measure: their spread, as one reference. +pub fn wider(a: Match, b: Match) Match { + const row = if (b.worst_row > a.worst_row) b else a; + return .{ .max_abs = @max(a.max_abs, b.max_abs), .cosine = @min(a.cosine, b.cosine), .worst_row = row.worst_row, .worst_at = row.worst_at }; +} + /// Ours against exact rows: its cosine error and worst row each within `factor` of the reference's. pub fn within(ours: Match, reference: Match) Verdict { const cosine_limit = factor * (1 - reference.cosine); @@ -140,3 +146,24 @@ test "against float64, a row further off than the reference allows fails" { try std.testing.expect(!v.pass); try std.testing.expectEqual(@as(usize, 3), v.row_at); } + +test "two reference runs: the gate is their spread, and a result outside it still fails" { + var exact: [cols * n]f32 = undefined; + synthetic(&exact, 4); + // one reference run near rounding, the other with row 2 off by 0.5%: as torch bf16 on two devices + var near: [cols * n]f32 = undefined; + for (&near, exact) |*r, x| r.* = host.round(x); + var far: [cols * n]f32 = near; + for (far[2 * cols ..][0..cols], exact[2 * cols ..][0..cols]) |*r, x| r.* = host.round(x * 1.005); + const spread = wider(match(f32, &near, &exact, cols), match(f32, &far, &exact, cols)); + try std.testing.expectEqual(@as(usize, 2), spread.worst_at); + var ours: [cols * n]u16 = undefined; + bf16Of(&ours, &exact); + for (ours[2 * cols ..][0..cols], exact[2 * cols ..][0..cols]) |*o, x| o.* = host.bf16(x * 1.005); + const rounded = rounding(&exact, cols); + // as far as the farther run: outside the nearer run's gate, inside the spread's + try std.testing.expect(!gate(match(u16, &ours, &exact, cols), match(f32, &near, &exact, cols), rounded).pass); + try std.testing.expect(gate(match(u16, &ours, &exact, cols), spread, rounded).pass); + for (ours[2 * cols ..][0..cols], exact[2 * cols ..][0..cols]) |*o, x| o.* = host.bf16(x * 1.05); + try std.testing.expect(!gate(match(u16, &ours, &exact, cols), spread, rounded).pass); +} diff --git a/zig/tests/moss_image.zig b/zig/tests/moss_image.zig index 66cf5ad535..62ae95edf7 100644 --- a/zig/tests/moss_image.zig +++ b/zig/tests/moss_image.zig @@ -4,7 +4,7 @@ const mtl = @import("metal"); const tf = @import("tensorfold"); const tokenizer = @import("tokenizer"); -const Metal = struct { +pub const Metal = struct { pub const mi = tf.moss_image; pub const st = tf.safetensors; pub const seconds = mtl.clock.seconds; diff --git a/zig/tests/moss_image_blocks.zig b/zig/tests/moss_image_blocks.zig new file mode 100644 index 0000000000..a92327f382 --- /dev/null +++ b/zig/tests/moss_image_blocks.zig @@ -0,0 +1,6 @@ +//! tf-moss-image-blocks on Metal: moss_image_blocks_run.zig's gate on the system GPU. +const std = @import("std"); + +pub fn main(init: std.process.Init) !u8 { + return @import("moss_image_blocks_run.zig").run(@import("moss_image.zig").Metal, init); +} diff --git a/zig/tests/moss_image_blocks_cuda.zig b/zig/tests/moss_image_blocks_cuda.zig new file mode 100644 index 0000000000..750e302b49 --- /dev/null +++ b/zig/tests/moss_image_blocks_cuda.zig @@ -0,0 +1,6 @@ +//! tf-moss-image-blocks on CUDA: moss_image_blocks_run.zig's gate on the first GPU. +const std = @import("std"); + +pub fn main(init: std.process.Init) !u8 { + return @import("moss_image_blocks_run.zig").run(@import("moss_image_cuda.zig").Cuda, init); +} diff --git a/zig/tests/moss_image_blocks_run.zig b/zig/tests/moss_image_blocks_run.zig new file mode 100644 index 0000000000..3e63badd09 --- /dev/null +++ b/zig/tests/moss_image_blocks_run.zig @@ -0,0 +1,129 @@ +//! tf-moss-image-blocks' driver for any backend `B` (as moss_image_run.zig's): one step, each block alone, gated. +const std = @import("std"); + +const usage = + \\tf-moss-image-blocks TRANSFORMER_DIR BLOCKS + \\BLOCKS (moss_image_case.py blocks) holds one step's branch: its input latents, timestep and prompt rows, a float32 + \\run's stream into each block and after the last, its velocity, and torch bf16's embedding, each block and the final + \\layer alone on that stream rounded to bf16, on the case's device and on the CPU. Ours runs the same steps + \\(Model.steps); each must be within twice torch bf16's distance from float32 (the farther of its two runs, or bf16 + \\rounding's, when larger) in cosine error and in its worst row, and a second run must repeat the first byte for + \\byte; else it fails. +; + +fn tensor(comptime st: type, f: *const st.File, name: []const u8, dtype: st.DType) ![]const u8 { + const t = f.get(name) orelse { + std.debug.print("BLOCKS has no {s}\n", .{name}); + return error.MissingTensor; + }; + if (t.dtype != dtype) return error.BadBlocksTensor; + return t.bytes; +} + +fn take(comptime T: type, comptime st: type, a: std.mem.Allocator, f: *const st.File, name: []const u8, dtype: st.DType) ![]T { + const bytes = try tensor(st, f, name, dtype); + const out = try a.alloc(T, bytes.len / @sizeOf(T)); + @memcpy(std.mem.sliceAsBytes(out), bytes); + return out; +} + +fn scalar(comptime st: type, f: *const st.File, name: []const u8) !i64 { + const bytes = try tensor(st, f, name, .i64); + if (bytes.len != 8) return error.BadBlocksTensor; + return std.mem.readInt(i64, bytes[0..8], .little); +} + +pub fn run(comptime B: type, init: std.process.Init) !u8 { + const mi = B.mi; + const st = B.st; + const check = mi.reference; + const H = mi.config.hidden; + const a = init.arena.allocator(); + const io = init.io; + const args = try init.minimal.args.toSlice(a); + if (args.len != 3) { + std.debug.print("{s}\n", .{usage}); + return 2; + } + var f = try st.File.open(a, io, args[2]); + defer f.close(io); + const step = try scalar(st, &f, "step"); + const skip_at = try scalar(st, &f, "skip"); + const skip: ?usize = if (skip_at < 0) null else std.math.cast(usize, skip_at) orelse return error.BadBlocksTensor; + const timestep = (try take(f32, st, a, &f, "timestep", .f32))[0]; + const latents = try take(f32, st, a, &f, "latents", .f32); + const prompt = try take(u16, st, a, &f, "prompt", .bf16); + const stream = try take(f32, st, a, &f, "stream", .f32); + const alone = [2][]u16{ try take(u16, st, a, &f, "alone", .bf16), try take(u16, st, a, &f, "alone_cpu", .bf16) }; + const velocity_fp32 = try take(f32, st, a, &f, "velocity_fp32", .f32); + const velocity_alone = [2][]f32{ try take(f32, st, a, &f, "velocity_alone", .f32), try take(f32, st, a, &f, "velocity_alone_cpu", .f32) }; + const e = f.names.getPtr("latents") orelse return error.MissingTensor; + if (e.rank != 3) return error.BadBlocksTensor; + const grid = e.shape[0..3]; + const image = grid[0] * grid[1]; + const text = prompt.len / mi.config.context; + const N = image + text; + const kept: usize = mi.config.layers - @as(usize, @intFromBool(skip != null)); + if (stream.len != (kept + 1) * N * H or alone[0].len != stream.len or alone[1].len != stream.len or latents.len != image * mi.config.channels or velocity_fp32.len != latents.len or velocity_alone[0].len != latents.len or velocity_alone[1].len != latents.len) { + std.debug.print("BLOCKS' tensors do not fit {d} image rows, {d} prompt rows and {d} blocks\n", .{ image, text, kept }); + return 2; + } + + var gpu: B.Device = undefined; + try gpu.init(); + defer gpu.deinit(); + const model = try gpu.load(std.heap.page_allocator, io, args[1], text, grid[1] * 16, grid[0] * 16); + defer model.deinit(); + var p = try model.prompt(prompt); + defer p.deinit(); + const inputs = try a.alloc(u16, stream.len); + for (inputs, stream) |*o, x| o.* = mi.host.bf16(x); + const rows = try a.alloc(u16, stream.len); + const velocity = try a.alloc(f32, latents.len); + try model.steps(&p, latents, timestep, skip, inputs, rows, velocity); + std.debug.print("{s}: step {d}, {s} branch, {d} image and {d} prompt rows, {d} blocks\n", .{ gpu.name(), step, if (skip == null) "prompt" else "negative", image, text, kept }); + + // block k's float32 output is the stream into block k + 1 (the embedding's, into block 0) + const verdicts = try a.alloc(check.Verdict, kept + 2); + const theirs = [2][]f32{ try a.alloc(f32, N * H), try a.alloc(f32, N * H) }; + var total: f64 = 0; + var most: f64 = 0; + var most_at: usize = 0; + var ratio: f64 = 0; + for (0..kept + 2) |k| { + const last = k == kept + 1; + const exact = if (last) velocity_fp32 else stream[k * N * H ..][0 .. N * H]; + const width: usize = if (last) mi.config.channels else H; + var runs: [2]check.Match = undefined; + for (&runs, theirs, alone, velocity_alone) |*r, t, al, va| { + if (!last) for (t, al[k * N * H ..][0 .. N * H]) |*o, v| { + o.* = mi.host.widen(v); + }; + r.* = check.match(f32, if (last) va else t, exact, width); + } + const ours = if (last) check.match(f32, velocity, exact, width) else check.match(u16, rows[k * N * H ..][0 .. N * H], exact, width); + const reference = check.wider(runs[0], runs[1]); + const v = check.gate(ours, reference, check.rounding(exact, width)); + verdicts[k] = v; + const used = @max(v.cosine_err / v.cosine_limit, v.row / v.row_limit); + total += used; + if (used > most) { + most = used; + most_at = k; + } + ratio += ours.worst_row / @max(reference.worst_row, 1e-30); + const name = if (k == 0) "the embedding" else if (last) "the final layer" else try std.fmt.allocPrint(a, "block {d}", .{if (skip) |s| (if (k - 1 < s) k - 1 else k) else k - 1}); + std.debug.print("step {d}, {s}: 1-cosine {e:.2} of {e:.2}, worst row {d:.5} (row {d}) of {d:.5}; torch bf16's 1-cosine {e:.2} and {e:.2}, worst row {d:.5} and {d:.5} (device, CPU): {s}\n", .{ k, name, v.cosine_err, v.cosine_limit, v.row, v.row_at, v.row_limit, 1 - runs[0].cosine, 1 - runs[1].cosine, runs[0].worst_row, runs[1].worst_row, if (v.pass) "pass" else "FAIL" }); + } + const steps_n: f64 = @floatFromInt(kept + 2); + std.debug.print("{d} steps: a mean of {d:.2} of the limit, at most {d:.2} (step {d}); our worst row {d:.2} times the farther torch bf16 run's on average\n", .{ kept + 2, total / steps_n, most, most_at, ratio / steps_n }); + + const rows_again = try a.alloc(u16, stream.len); + const velocity_again = try a.alloc(f32, latents.len); + try model.steps(&p, latents, timestep, skip, inputs, rows_again, velocity_again); + const same = std.mem.eql(u16, rows, rows_again) and std.mem.eql(u8, std.mem.sliceAsBytes(velocity), std.mem.sliceAsBytes(velocity_again)); + std.debug.print("run again, the same bytes: {}\n", .{same}); + const pass = check.accept(verdicts, same); + std.debug.print("{s}\n", .{if (pass) "PASS every step within its gate, repeated byte for byte" else "FAIL"}); + return if (pass) 0 else 1; +} diff --git a/zig/tests/moss_image_cuda.zig b/zig/tests/moss_image_cuda.zig index f633a8e13c..f5689c11cb 100644 --- a/zig/tests/moss_image_cuda.zig +++ b/zig/tests/moss_image_cuda.zig @@ -3,7 +3,7 @@ const std = @import("std"); const cuda = @import("cuda"); const core = @import("core"); -const Cuda = struct { +pub const Cuda = struct { pub const mi = @import("moss_image"); pub const st = core.safetensors; pub const seconds = mi.seconds; diff --git a/zig/tests/moss_image_run.zig b/zig/tests/moss_image_run.zig index ad1779aa2f..ddea817bef 100644 --- a/zig/tests/moss_image_run.zig +++ b/zig/tests/moss_image_run.zig @@ -15,6 +15,7 @@ const usage = \\build has the encoder. Exit 0 means every teacher step and the repeat passed, 1 failed, 2 is INCONCLUSIVE. \\Per element the bound is 2 * max(|bf16 teacher - fp32 teacher|, half a bf16 ULP at the fp32 value). \\Missing or partial fp32 teacher coverage cannot pass. A start-only case still produces a picture with exit 2. + \\tf-moss-image-blocks gates each transformer block alone against the same publisher float32 run. ; /// The card's sampling defaults: the guidance scale, and the block the negative branch skips. From 91fdbe8213e486183453841cad4e7d2695eb8106 Mon Sep 17 00:00:00 2001 From: BobClawblaw <288991054+BobClawblaw@users.noreply.github.com> Date: Thu, 5 Nov 2026 20:36:08 -0600 Subject: [PATCH 35/36] changelog: #2294 under Unreleased --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index b2856928bf..7f8f86d6e6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,7 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel - Internal: the kernel contract's nvcc 13.0 resource records match what nvcc 13.0.88 builds for 57 CUDA kernels (226 per-SM slots), so the five-SM limits gate passes on main; no kernel, flag or launch changes (#2390). - Flash Next on CUDA: a prompt-cache mark inside a prompt chunk is kept on the way through the pass instead of cutting the pass there, so a new short prompt's first token comes 19% sooner on a GB10, with replies unchanged; `TF_FLASHNEXT_KEEP_INPASS=0` cuts the pass at the mark as before (#1866). Thanks to @Arminova. +- Tests: `tf-moss-image-blocks` gates MOSS-Image's transformer block by block against the publisher's float32 run, within twice the spread of torch bf16's device and CPU runs, on Metal and CUDA (#2294, part of #1529). - Cluster: MCDMA bring-up on real hosts, step 2a of #1594 (#2088). Thanks to @ThinkOffApp. From bee665d6f4b3394d38c9d5c8b2b3164a7f07426f Mon Sep 17 00:00:00 2001 From: ashhart <71890694+ashhart@users.noreply.github.com> Date: Sat, 7 Nov 2026 03:27:38 +0000 Subject: [PATCH 36/36] build: the MOSS-Image CUDA tools install from moss_image_tools.zig, so zig/build/cuda.zig stays under 600 lines With main at 600 lines, this branch's tf-moss-image-blocks took zig/build/cuda.zig to 602. Both MOSS-Image tools move, unchanged, to one file that cuda.zig calls once. --- zig/build/cuda.zig | 7 +------ zig/build/moss_image_tools.zig | 11 +++++++++++ 2 files changed, 12 insertions(+), 6 deletions(-) create mode 100644 zig/build/moss_image_tools.zig diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 89e4bdd6e5..bc7411717d 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -136,12 +136,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu const q35_install = b.addInstallArtifact(b.addExecutable(.{ .name = "tf-qwen35-cuda", .root_module = q35 }), .{}); b.getInstallStep().dependOn(&q35_install.step); b.step("qwen35-cuda", "Build the Qwen CUDA forward and raw state receipt tool").dependOn(&q35_install.step); - // tf-moss-image (a case, gated and timed) and tf-moss-image-blocks (each block alone): tools/zig/moss_image_case.py - const moss = b.createModule(.{ .root_source_file = b.path("zig/src/families/moss_image/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core } } }); - const moss_tool = b.createModule(.{ .root_source_file = b.path("zig/tests/moss_image_cuda.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core }, .{ .name = "moss_image", .module = moss } } }); - b.installArtifact(b.addExecutable(.{ .name = "tf-moss-image", .root_module = moss_tool })); - const moss_blocks = b.createModule(.{ .root_source_file = b.path("zig/tests/moss_image_blocks_cuda.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = mods.core }, .{ .name = "moss_image", .module = moss } } }); - b.installArtifact(b.addExecutable(.{ .name = "tf-moss-image-blocks", .root_module = moss_blocks })); + @import("moss_image_tools.zig").install(b, target, optimize, strip, cuda, mods.core); const d3_tool = b.createModule(.{ .root_source_file = b.path("zig/src/families/drama3/generate.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "core", .module = mods.core }, .{ .name = "drama3", .module = @import("drama3_host.zig").module(b, target, optimize, cuda, mods.core) } } }); const d3_exe = b.addExecutable(.{ .name = "drama3-generate", .root_module = d3_tool }); b.step("drama3-generate", "Drama generation and replay controls on CUDA").dependOn(&b.addInstallArtifact(d3_exe, .{}).step); diff --git a/zig/build/moss_image_tools.zig b/zig/build/moss_image_tools.zig new file mode 100644 index 0000000000..4021d6f513 --- /dev/null +++ b/zig/build/moss_image_tools.zig @@ -0,0 +1,11 @@ +//! MOSS-Image's CUDA tools: tf-moss-image (a case, gated and timed) and tf-moss-image-blocks (each block alone). +const std = @import("std"); + +/// Both tools run tools/zig/moss_image_case.py's reference cases; cuda.zig installs them with its CUDA targets. +pub fn install(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, strip: bool, cuda: *std.Build.Module, core: *std.Build.Module) void { + const moss = b.createModule(.{ .root_source_file = b.path("zig/src/families/moss_image/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core } } }); + const moss_tool = b.createModule(.{ .root_source_file = b.path("zig/tests/moss_image_cuda.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core }, .{ .name = "moss_image", .module = moss } } }); + b.installArtifact(b.addExecutable(.{ .name = "tf-moss-image", .root_module = moss_tool })); + const moss_blocks = b.createModule(.{ .root_source_file = b.path("zig/tests/moss_image_blocks_cuda.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core }, .{ .name = "moss_image", .module = moss } } }); + b.installArtifact(b.addExecutable(.{ .name = "tf-moss-image-blocks", .root_module = moss_blocks })); +}