From 97b782d3017c7c816149b7f8f78336ebfd911dfd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 21 Oct 2026 01:42:50 +0200 Subject: [PATCH 1/2] Flash Next CUDA: no zero fills the one-rank ModelOpt pass reads back unchanged Two fills a layer that a one-rank pass over ModelOpt weights never needed: - the MoE output y [rows, 11, dims] fp32 was cleared before the experts, though the down kernel writes every routed slot and sharedScatter the shared one; a rank that skips the shared slot (tensor parallel) still clears it; - the native out projections (DeltaNet and attention) reduce their own K slices, then cleared the out_sk - 1 slices the writeback sums after the first. The writeback now adds one zero slice kept at part's cap-row stride instead: x + 0 repeated is x + 0 once (-0 to +0 the same), so its bits are the same. The zero slice is filled at allocation; a legacy lane that writes slices into part, or a native one past cap rows, marks it spent and the old path runs. Kept branches (training) and collectives keep the old path. At 1,536 rows that is 173 MB plus 110 MB of fills and 94 MB of reads a layer: the 8k prompt takes 5.77-6.74 s (main 6.09-6.96), 30,922 tokens 25.04 s (25.74), every reply byte-identical. --- zig/src/families/flashnext/cuda_batch.zig | 18 ++++++++++++++++-- .../families/flashnext/cuda_batch_state.zig | 5 +++++ 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index 5e79eee832..c1b66192af 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -107,6 +107,8 @@ pub fn Pass(comptime Tri: type) type { prepared: bool = false, /// The pending write-back's branch sits in fp32 in moe_f32 (something added to it, or it is kept). pending_f32: bool = false, + /// This layer's out projection leaves part's other slices alone; its writeback adds the zero slice once. + zero_slices: bool = false, // ---- launches ---------------------------------------------------------------------------------------- @@ -115,11 +117,15 @@ pub fn Pass(comptime Tri: type) type { if (w.native) |native| { if (sk == 1) return self.k.native_lane.matmul(self.stream, x, w.k, m, native, out, self.s.native_part); // Legacy callers consume SK slices; the native lane has already reduced its own shape-driven slices. + if (part == self.s.part and m * w.n > self.s.cap * dims) self.s.part_zero = false; try self.k.native_lane.matmulF32(self.stream, x, w.k, m, native, part, self.s.native_part); + // the writeback adds part's cap-stride zero slice in their place (zero_slices) + if (part == self.s.part and self.zero_slices) return; const bytes = m * w.n * 4; if (sk > 1) try self.stream.d.check(self.stream.d.api.cuMemsetD8Async(part + bytes, 0, bytes * (sk - 1), self.stream.handle), "cuMemsetD8Async"); return; } + if (part == self.s.part and sk > 1) self.s.part_zero = false; const bm: usize = 16; const bn: usize = 64; const k = w.k; @@ -435,7 +441,9 @@ pub fn Pass(comptime Tri: type) type { const plan: cuda.grouped.Plan = .{ .members = s.members, .items = s.items, .counts = s.counts, .rank = s.rank, .hist = s.hist }; if (lw.native_experts) |*native| { const owned = &native.experts; - try self.stream.d.check(self.stream.d.api.cuMemsetD8Async(s.y, 0, pairs * dims * 4, self.stream.handle), "cuMemsetD8Async"); + // one rank writes every slot (routed pairs by the down kernel, the shared one by sharedScatter), so only + // a rank that skips the shared slot needs the zeros its sum over ranks adds + if (native.world != 1 or native.rank != 0) try self.stream.d.check(self.stream.d.api.cuMemsetD8Async(s.y, 0, pairs * dims * 4, self.stream.handle), "cuMemsetD8Async"); try self.k.router.route(self.stream, s.picks, pairs, e, 16, plan); const items = cuda.grouped.maxItems(pairs, e, 16); try self.k.nvfp4.run(self.stream, .up, .{ .x = s.mixed, .stride = dims, .slots = slots }, owned.view(), plan, s.moe_act, items, experts_n); @@ -527,13 +535,19 @@ pub fn Pass(comptime Tri: type) type { } try self.writeback(0, n); try self.readoutKeep(lw.attn_hc, 0, n, keep.attn_hc); + // one rank, nothing kept from the branch, and the zero slice intact: its native out projection's + // writeback reads one zero slice where it read out_sk - 1 zeroed ones (x + 0 repeated is x + 0 once) + self.zero_slices = self.s.part_zero and keep.mix_branch == 0 and self.m.collective == null and lw.out.native != null; if (lw.linear) { try self.gdnBlock(lw, wt.linearIndex(layer), keep); } else { try self.attnBlockKeep(lw, wt.attnIndex(layer), keep); } if (keep.mix_branch != 0) try self.sumSlices(s.part, keep.mix_branch, true, n * dims, out_sk); - try self.tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", n * dims, n, dims, streams, 4, 1, 1, out_sk); + if (self.zero_slices) { + try self.tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", s.cap * dims, n, dims, streams, 4, 1, 1, 2); + } else try self.tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", n * dims, n, dims, streams, 4, 1, 1, out_sk); + self.zero_slices = false; try self.readoutKeep(lw.mlp_hc, 0, n, keep.mlp_hc); try self.moeBlockKeep(lw, keep); try self.moeAdds(layer, keep); diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index dc1749c657..6d1d248734 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -266,6 +266,9 @@ pub const Scratch = struct { xsm: u64 = 0, proj: u64 = 0, part: u64 = 0, + /// part's second slice at a cap-row stride ([cap * dims, 2 * cap * dims) fp32) holds zeros: native lanes write only + /// the first, so a slice writeback can add these zeros once where it would add out_sk - 1 zeroed slices. + part_zero: bool = false, gout: u64 = 0, gxs: u64 = 0, ys: u64 = 0, @@ -411,6 +414,8 @@ pub const Scratch = struct { s.xsm = try s.add(gpa, d, cap * (dims / 32) * 4); s.proj = try s.add(gpa, d, cap * proj_row * 2); s.part = try s.add(gpa, d, out_sk * cap * dims * 4); + try s.bufs.items[s.bufs.items.len - 1].fill8(0, stream.handle); + s.part_zero = true; s.gout = try s.add(gpa, d, cap * gdn.value_dim * 2); s.gxs = try s.add(gpa, d, cap * (gdn.value_dim / 32) * 4); s.ys = try s.add(gpa, d, cap * gdn.value_dim * 4); From 3e050cf9301e57f95a68fe025b8db350c170c35d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=BCrgen=20Schmied?= <600316+jschmied@users.noreply.github.com> Date: Wed, 21 Oct 2026 15:13:26 +0200 Subject: [PATCH 2/2] Flash Next CUDA: one-line comments, cuda_batch.zig back under 600 lines The lean gate's new findings on this branch: the three multi-line comments folded to one line within 120 columns, which also brings cuda_batch.zig to 599 lines. Comments only; lean_new.py: 0 new. --- zig/src/families/flashnext/cuda_batch.zig | 6 ++---- zig/src/families/flashnext/cuda_batch_state.zig | 3 +-- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index c1b66192af..4566e1696e 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -441,8 +441,7 @@ pub fn Pass(comptime Tri: type) type { const plan: cuda.grouped.Plan = .{ .members = s.members, .items = s.items, .counts = s.counts, .rank = s.rank, .hist = s.hist }; if (lw.native_experts) |*native| { const owned = &native.experts; - // one rank writes every slot (routed pairs by the down kernel, the shared one by sharedScatter), so only - // a rank that skips the shared slot needs the zeros its sum over ranks adds + // one rank writes every slot; only a rank that skips the shared one needs zeros for the sum over ranks if (native.world != 1 or native.rank != 0) try self.stream.d.check(self.stream.d.api.cuMemsetD8Async(s.y, 0, pairs * dims * 4, self.stream.handle), "cuMemsetD8Async"); try self.k.router.route(self.stream, s.picks, pairs, e, 16, plan); const items = cuda.grouped.maxItems(pairs, e, 16); @@ -535,8 +534,7 @@ pub fn Pass(comptime Tri: type) type { } try self.writeback(0, n); try self.readoutKeep(lw.attn_hc, 0, n, keep.attn_hc); - // one rank, nothing kept from the branch, and the zero slice intact: its native out projection's - // writeback reads one zero slice where it read out_sk - 1 zeroed ones (x + 0 repeated is x + 0 once) + // one rank, nothing kept, zero slice intact: the writeback adds one zero slice, not out_sk - 1 zeroed self.zero_slices = self.s.part_zero and keep.mix_branch == 0 and self.m.collective == null and lw.out.native != null; if (lw.linear) { try self.gdnBlock(lw, wt.linearIndex(layer), keep); diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index 6d1d248734..b24d77fc43 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -266,8 +266,7 @@ pub const Scratch = struct { xsm: u64 = 0, proj: u64 = 0, part: u64 = 0, - /// part's second slice at a cap-row stride ([cap * dims, 2 * cap * dims) fp32) holds zeros: native lanes write only - /// the first, so a slice writeback can add these zeros once where it would add out_sk - 1 zeroed slices. + /// part's second cap-stride slice is zeros (native lanes write only the first): a writeback adds it once. part_zero: bool = false, gout: u64 = 0, gxs: u64 = 0,