exl3: zero the split-K counters once as they grow, not before every call

CHANGELOG.md

@@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel
 
 ## Unreleased
 
+- CUDA EXL3: the dense projections zero their split-K counters only when the buffer grows, on the calls' stream, instead of before every 128-row block of every call; each block's last program already resets its own counter. Replies are token-identical, and on a GB10 with the Qwen3.8-27B 3.50 bpw pack decode is 0.4% to 0.9% faster (#1475). Thanks to @diwakergupta.
+
 - Qwen3.6-35B-A3B on CUDA serves NVIDIA's NVFP4 checkpoint (`nvidia/Qwen3.6-35B-A3B-NVFP4`) as downloaded: W4A16 NVFP4 routed and shared experts, FP8 attention and DeltaNet projections and the NVFP4 head run as stored, and `--drafter` naming the same checkpoint drafts with its bf16 MTP layer, quantized to affine-4 at load. One GB10 decodes 66 tok/s without drafts and 112 to 125 tok/s with them; drafted, concurrent and resumed replies equal their plain runs. The rules come from @gilby's #345.
 
 - DeepSeek-V4.1 on CUDA no longer stays up unable to answer after a sticky CUDA context fault: a host monitor sends the fault to fail-fast at once instead of waiting for a health poll, and with `TF_TP_FAILFAST=0` the process exits with status 70 rather than waiting on communicator teardown. Under fail-fast, a round error the ranks cannot recover from (such as `CudaFailed` or `Timeout`) now takes every rank down instead of leaving them out of step. A stall watchdog, off by default, withdraws readiness and exits when a round, snapshot, KV union or follower step makes no progress for `TF_DSV41_STALL_EXIT_S` seconds; off, it costs one null check per launch, and replies are unchanged (#1646). Thanks to @jayleaton.

docs/recipes/exl3.md

@@ -195,6 +195,7 @@ row's bits match it:
   for 1 to 128 rows, every codebook and width (the `mul1` half steps included), with the K split planned from the
   layer's shape alone. `exl3.Dense` runs prompt rows in 128-row blocks of the same kernel, so a prompt row's bits
   equal its decode-row bits wherever its chunk starts; `qlinear.Weight` has an `.exl3` arm for both.
+  The split counters are zeroed on the calls' stream when they grow, and each block's last program resets its own.
 - `zig/kernels/cuda/exl3_experts.cu`, `zig/src/cuda/exl3_experts.zig`: grouped routed experts (member lists,
   each expert's rotation, gate and up, GLM's SwiGLU roundings, down and the weighted combine), widths mixed within
   a layer; an expert's member count never changes a row's bits.

zig/kernels/cuda/exl3_linear.cu

@@ -20,7 +20,7 @@ struct LinearArgs {
     const half* bias;        // fp16 [n] or null
     void* y;                 // output
     float* z;                // fp32 split partials [sk, m, n] or null
-    int* counters;           // per (pass, block) arrival counts, zeroed before the launch
+    int* counters;           // per (pass, block) arrival counts: zero at launch, reset by the block's last program
     int y_dtype;
     int n, sk, m, k;
 };

zig/src/cuda/exl3.zig

@@ -204,7 +204,7 @@ pub const Dense = struct {
     candidate_mod: ?Module = null,
     candidate_k: ?candidate.Kernels = null,
     implementation: enum { reference, tiles } = .reference,
-    counters: ?memory.DeviceBuffer = null, // per pass and block, zeroed before every call
+    counters: ?memory.DeviceBuffer = null, // per pass and block: zeroed on the calls' stream, reset by the kernel
 
     pub fn load(d: *const Driver) !Dense {
         return .{ .d = d, .lin = try Linear.load(d) };
@@ -230,13 +230,21 @@ pub const Dense = struct {
         return slot.*.?.ptr;
     }
 
-    /// Grow every scratch buffer to what 128 rows of `w` need, so a call inside a graph capture allocates nothing.
-    pub fn reserve(de: *Dense, w: Weight) !void {
+    /// Grow every scratch buffer to what 128 rows of `w` need on `s`, so a call in a graph capture allocates nothing.
+    pub fn reserve(de: *Dense, s: Stream, w: Weight) !void {
         if (plan(w.k, w.n, 192, 8).sk > 1) _ = try de.scratch(&de.z, Linear.zBytes(w, 128));
-        _ = try de.scratch(&de.counters, Linear.counterBytes(w));
+        _ = try de.grownCounters(s, Linear.counterBytes(w));
         _ = try de.scratch(&de.xh, 128 * w.k * 2);
     }
 
+    /// The counters, zeroed on `s` when they grow: each block's last arriving program resets its own.
+    fn grownCounters(de: *Dense, s: Stream, bytes: usize) !u64 {
+        if (de.counters != null and de.counters.?.len >= bytes) return de.counters.?.ptr;
+        const ptr = try de.scratch(&de.counters, bytes);
+        try de.counters.?.fill8(0, s.handle);
+        return ptr;
+    }
+
     /// Decode rows: one call (the kernel takes 1..128), any count above in 128-row blocks.
     pub fn decode(de: *Dense, s: Stream, w: Weight, x: u64, x_dtype: DType, rows: usize, y: u64, y_dtype: DType) !void {
         if (de.implementation == .tiles) return de.tiles(s, w, x, x_dtype, rows, y, y_dtype, false);
@@ -278,14 +286,13 @@ pub const Dense = struct {
     fn blocks(de: *Dense, s: Stream, w: Weight, x: u64, x_dtype: DType, rows: usize, y: u64, y_dtype: DType) !void {
         const p = plan(w.k, w.n, 192, 8);
         const z: u64 = if (p.sk > 1) try de.scratch(&de.z, Linear.zBytes(w, 128)) else 0;
-        const counters = try de.scratch(&de.counters, Linear.counterBytes(w));
+        const counters = try de.grownCounters(s, Linear.counterBytes(w));
         const xe: u64 = if (x_dtype == .f32) 4 else 2;
         const ye: u64 = if (y_dtype == .f32) 4 else 2;
         var r: usize = 0;
         while (r < rows) : (r += 128) {
             const m = @min(rows - r, 128);
             const xh = try de.scratch(&de.xh, m * w.k * 2);
-            try de.counters.?.fill8(0, s.handle);
             try de.lin.matmul(s, w, x + @as(u64, r) * w.k * xe, x_dtype, m, y + @as(u64, r) * w.n * ye, y_dtype, xh, z, counters);
         }
     }
@@ -364,6 +371,23 @@ test "a column view keeps the parent's strides and moves to its first 128-column
     }
 }
 
+test "the counters are zeroed on the calls' stream only as they grow" {
+    const mock = @import("exl3_mock.zig");
+    var d = mock.init();
+    var de: Dense = .{ .d = &d, .lin = undefined };
+    defer if (de.counters) |*c| c.free();
+    const s: Stream = .{ .d = &d, .handle = @ptrFromInt(0x5100) };
+    _ = try de.grownCounters(s, 64);
+    try std.testing.expectEqual(@as(usize, 1), mock.memsets);
+    try std.testing.expectEqual(s.handle, mock.memset_stream);
+    _ = try de.grownCounters(s, 64);
+    _ = try de.grownCounters(s, 32);
+    try std.testing.expectEqual(@as(usize, 1), mock.memsets);
+    _ = try de.grownCounters(s, 128);
+    try std.testing.expectEqual(@as(usize, 2), mock.memsets);
+    try std.testing.expectEqual(@as(usize, 0), mock.legacy_memsets);
+}
+
 test "candidate selection refuses layouts it cannot represent" {
     const st = strides(128, 256, 8, false);
     var w: Weight = .{ .trellis = 1, .suh = 2, .svh = 3, .k = 128, .n = 256, .halves = 8, .codebook = .mcg, .stride_k = st.k, .stride_nb = st.nb };

zig/src/cuda/exl3_mock.zig

@@ -15,6 +15,9 @@ pub var fail_copy: usize = 0; // the copy (1 for the first) whose waits fail, 0
 pub var fail_waits: usize = 0; // how many of its waits fail before one completes
 pub var pending: bool = false; // a copy is accepted and no wait has completed it
 pub var bad_frees: usize = 0; // device frees while a copy was pending
+pub var memsets: usize = 0; // memsets on a stream the caller named
+pub var legacy_memsets: usize = 0; // memsets on the legacy default stream
+pub var memset_stream: abi.Stream = null;
 pub fn init() Driver {
     shared_calls = 0;
     shared_bytes = 0;
@@ -27,13 +30,18 @@ pub fn init() Driver {
     fail_waits = 0;
     pending = false;
     bad_frees = 0;
+    memsets = 0;
+    legacy_memsets = 0;
+    memset_stream = null;
     var d: Driver = undefined;
     d.api.cuMemAlloc_v2 = alloc;
     d.api.cuMemFree_v2 = free;
     d.api.cuMemcpyHtoD_v2 = copy;
     d.api.cuStreamSynchronize = sync;
     d.api.cuFuncSetAttribute = attribute;
     d.api.cuCtxSynchronize = context;
+    d.api.cuMemsetD8Async = memsetAsync;
+    d.api.cuMemsetD8_v2 = memsetLegacy;
     d.api.cuGetErrorName = name;
     d.api.cuGetErrorString = name;
     return d;
@@ -59,6 +67,15 @@ fn copy(_: abi.DevicePtr, src: ?*const anyopaque, n: usize) callconv(.c) abi.Res
     pending = true;
     return 0;
 }
+fn memsetAsync(_: abi.DevicePtr, _: u8, _: usize, s: abi.Stream) callconv(.c) abi.Result {
+    memsets += 1;
+    memset_stream = s;
+    return 0;
+}
+fn memsetLegacy(_: abi.DevicePtr, _: u8, _: usize) callconv(.c) abi.Result {
+    legacy_memsets += 1;
+    return 0;
+}
 fn sync(_: abi.Stream) callconv(.c) abi.Result {
     if (pending and copy_count == fail_copy and fail_waits > 0) {
         fail_waits -= 1;

zig/src/families/qwen3_5/cuda_exl3_weights.zig

@@ -142,7 +142,7 @@ fn dense(s: *Sink, kit: ?*Kit, ck: *core.Checkpoint, prefix: []const u8, mod: []
         .stride_k = st.k,
         .stride_nb = st.nb,
     };
-    if (kit) |x| try x.dense.reserve(w);
+    if (kit) |x| try x.dense.reserve(s.e.stream, w);
     return .{ .x3 = w };
 }