chore(train): engine-10 continued: the raw q4_0/q8_0 flash attention at depth, llama-bench -rs

TORAD.md

@@ -91,6 +91,8 @@ which pins a commit of this branch as a submodule.
 | `b2eb4336a` | `llama-bench` had no parameter axis for the recurrent state's cache type, so a bench of a Gated DeltaNet model always ran the context's default `type_s` (f32) instead of what the server serves — Ternary Bonsai 2 27B serves an f16 state, and a bench at f32 moved twice its state bytes (48 layers x 48 x 128 x 128). `-cts`, matching `-ctk`/`-ctv`, is now a field in every printer | (new option) |
 | `16ad036c4` | between two PQ2_0 matmuls a decode token runs a chain of small kernels (norm, FWHT, q8_1 quantize, the Gated DeltaNet layer's conv/recurrence/gated norm), 1-3 us each, and PDL lets the next matmul request its own weights only under the one kernel before it — on an RTX 5080 at depth 16,384 that chain is ~1.3 ms of a 10 ms token against ~50 us floors. A PQ2_0 launch's blocks now take every gridDim.x-th tile from their index (so what a launch reads first is its matrices' heads), and each block of the launch before it, once its boxes have landed, prefetches its share of those heads into L2 (`cp.async.bulk.prefetch.L2`), sized from `GGML_CUDA_PQ2_PREFETCH_US` (default 2, at most a quarter of the card's L2; 0 disables it) times the device's DRAM rate. Ternary Bonsai 2 27B, llama-bench tg64, served state, graphs on, medians of 20 reps: RTX 5080 +2.4 % at depth 16,384 (100.54 -> 102.97 tok/s), +2.7 % at 0 (103.75 -> 106.55); RTX 5070 Ti +1.6 % (78.98 -> 80.21), +1.1 % (79.83 -> 80.71); RTX 5090, against engine-02512a3's `libggml-cuda` under this build's binaries, +4.84 % at 16,384 (152.93 -> 160.34) and +4.51 % at 0 (157.57 -> 164.68); greedy 128 tokens byte-identical to the parent at 2 and 8 us on the 5080 and at the default on the 5070 Ti. `test-backend-ops`: the 296 PQ2_0 cases and 24 MUL_MAT_GROUP cases pass with the prefetch off, at 8 us and at 64 us | `GGML_CUDA_PQ2_PREFETCH_US=0` |
 | `4104c47d5` | `llama-bench`'s own type-name table had no f32, so `-cts f32` exited on its arguments though f32 is the state's default and a type the engine serves, and `-ctk f32 -ctv f32` did the same for a K/V pair the CUDA flash attention runs. Measured on the RTX 5080: `-cts f32,f16` runs both rows (tg16 99.26 and 99.20), `-ctk f32 -ctv f32` runs, `-cts f64` still exits with the invalid-parameter error | (fix) |
+| `01f4f0fda` | the raw q4_0/q8_0 flash attention lost time to its own shared memory beside its DRAM stream: at 65,536 cells on an RTX 5080 (Ternary Bonsai 2 27B's attention: head 256, 4 KV heads at GQA 6, a bit mask) 3.43M of its 7.36M shared-memory wavefronts were bank conflicts — the V tile's dequant stored 16 bytes from 8 threads on 2 rows (4-way), and the K scales' float tile loaded 32 rows of one block a warp (4-way) — and a decode token ran 2 blocks an SM at 43,168 bytes. The dequant's store phase now spans 8 rows (store conflicts 3.2M -> 0.09M); K*Q reads each block's f16 scale from the raw rows, dropping the float tile, its pass and its barrier (load conflicts 0.23M -> 0.03M, 26,272 bytes a block, 3 blocks an SM); the stream-k fixup loads the next 8 blocks' partials before folding (6.9 -> 5.9 us); padding columns write no fixup partial; and the next raw K tile loads beside V where 2 blocks share an SM (a verify's 4-warp tile). test-backend-ops perf, the served layout, 6 rounds against the parent's library: a token -16.7 / -3.7 / -3.8 / -1.8 % and a 4-row verify -6.7 / -3.6 / -2.9 / -2.8 % at 16,384 / 65,536 / 131,072 / 245,760 cells; llama-bench as served, 4 rounds while the host swapped under other builds: tg64 +1.2 % at 65,536 (91.19 -> 92.30 tok/s) and +0.5 % at 16,384, pp4 -0.5 % and -0.1 %, inside that noise (the attention's share of the step predicts +0.3 to +0.7 %). FLASH_ATTN_EXT 3,220/3,220; greedy 64 tokens after a 7K-token prompt byte-identical to the parent | none of its own (it keeps the raw path's arithmetic): `GGML_CUDA_FATTN_Q4_0_LEGACY=1` / `GGML_CUDA_FATTN_Q8_0_LEGACY=1` restore the stock kernels |
+| `7537f40da` | `llama-bench` left the context's `n_rs_seq` at 0, so its pp4 measured a verify that writes one snapshot of each recurrent layer's state and conv window, where a drafting server (`n_rs_seq` = the draft's n_max) writes n_max + 1. `-rs` is a parameter axis like `-cts` and a field in every printer; the bench's context holds at least `n_rs_seq` + 2 rows, since the batch is clamped to the context and `split_equal` keeps a draft's last `n_rs_seq` + 1 rows in one larger ubatch (pp4 with `-rs 3` at depth 0 aborted there). RTX 5080, Ternary Bonsai 2 27B, q4_0 K/V, f16 state, pp4 in a 512 ubatch at depth 16,384, 32 samples each: `-rs 0` 386.15, `-rs 3` 383.76 tok/s (-0.62 %) | (new option) |
 
 Every switch in the last column is read once per process and parses as an integer: a `*_LEGACY` switch set to `0`
 is the same as unset (the change stays on), and `=0` turns off `GGML_CUDA_LORA_RANK1_FUSE` and

ggml/src/ggml-backend-meta.cpp

@@ -1083,7 +1083,35 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
                         split_state.ne[j] *= tensor->ne[split_state.axis];
                         if (split_state.ne[j] != 0 || tensor->src[i]->ne[src_ss[i].axis] != 0) {
                             const int64_t div = tensor->src[i]->ne[src_ss[i].axis] * split_state.nr[0];
-                            GGML_ASSERT(split_state.ne[j] % div == 0);
+                            if (split_state.ne[j] % div != 0) {
+                                char segbuf[256];
+                                int segpos = 0;
+                                for (size_t s = 0; s < src_ss[i].n_segments && segpos < 200; s++) {
+                                    segpos += snprintf(segbuf + segpos, sizeof(segbuf) - segpos, "%s[%lld]*%u",
+                                        s > 0 ? "+" : "", (long long) src_ss[i].ne[s*n_bufs + j],
+                                        src_ss[i].nr[s]);
+                                }
+                                char chainbuf[512];
+                                int chainpos = 0;
+                                // names/ops/dims only: calling back into the split computer here
+                                // recursed without bound (its cache is being filled by this walk)
+                                const ggml_tensor * anc = tensor;
+                                for (int depth = 0; depth < 6 && anc != nullptr && chainpos < 430; depth++) {
+                                    chainpos += snprintf(chainbuf + chainpos, sizeof(chainbuf) - chainpos,
+                                        "%s%s:%s[%lld]", depth > 0 ? " <- " : "",
+                                        anc->name, ggml_op_name(anc->op), (long long) anc->ne[0]);
+                                    anc = anc->src[0] == anc ? nullptr : anc->src[0];
+                                }
+                                GGML_ABORT("tensor split ratio is not integral: node=%s op=%s dst=[%lld,%lld,%lld,%lld] dstaxis=%d src%zu=[%lld,%lld,%lld,%lld] axis=%d segs=%s num=%lld den=%lld chain=%s",
+                                    tensor->name, ggml_op_name(tensor->op),
+                                    (long long) tensor->ne[0], (long long) tensor->ne[1],
+                                    (long long) tensor->ne[2], (long long) tensor->ne[3],
+                                    split_state.axis, i,
+                                    (long long) tensor->src[i]->ne[0], (long long) tensor->src[i]->ne[1],
+                                    (long long) tensor->src[i]->ne[2], (long long) tensor->src[i]->ne[3],
+                                    src_ss[i].axis, segbuf, (long long) split_state.ne[j], (long long) div,
+                                    chainbuf);
+                            }
                             split_state.ne[j] /= div;
                         }
                     }

ggml/src/ggml-cuda/cp-async.cuh

@@ -55,3 +55,23 @@ static __device__ __forceinline__ void cp_async_wait_all() {
     NO_DEVICE_CODE;
 #endif // CP_ASYNC_AVAILABLE
 }
+
+// Closes a group of this thread's asynchronous copies: the ones issued since the last group.
+static __device__ __forceinline__ void cp_async_commit_group() {
+#ifdef CP_ASYNC_AVAILABLE
+    asm volatile("cp.async.commit_group;");
+#else
+    NO_DEVICE_CODE;
+#endif // CP_ASYNC_AVAILABLE
+}
+
+// Makes each thread wait until at most its n newest groups are still in flight, every older one done.
+// As cp_async_wait_all, no synchronization beyond the thread.
+template <int n>
+static __device__ __forceinline__ void cp_async_wait_group() {
+#ifdef CP_ASYNC_AVAILABLE
+    asm volatile("cp.async.wait_group %0;" : : "n"(n));
+#else
+    NO_DEVICE_CODE;
+#endif // CP_ASYNC_AVAILABLE
+}

ggml/src/ggml-cuda/fattn-common.cuh

@@ -1071,6 +1071,75 @@ static __global__ void flash_attn_mask_to_KV_live(
     }
 }
 
+// A stream-k tile's fold: the partial results of blocks b_last - 1 down to b_first, onto the one of block b_last (dst_val,
+// max_val, rowsum), the order and the formulas every fixup kernel used, so the bits are theirs. data: the blocks' partial
+// results (block b's column jc at b*ncols*D + jc*D); meta: their KQ max and rowsum (at b*ncols + jc); skip(b): the block did
+// no work on this tile. Each of the next FATTN_FIXUP_CHUNK blocks' loads goes out before the current ones fold, so a fold
+// waits on its fmaxf and its FMAs, not on a load per block: with one load per block in turn, a decode token's fixup at
+// 65,536 cells took 6.9 us (ncu, RTX 5080, 42 blocks a tile), folded 8 ahead 5.9 us at 63.
+#define FATTN_FIXUP_CHUNK 8
+
+template <int D, int ncols, typename skip_t>
+static __device__ __forceinline__ void flash_attn_fixup_fold(
+        float & dst_val, float & max_val, float & rowsum, const float * __restrict__ data, const float2 * __restrict__ meta,
+        const int jc, const int tid, const int b_first, const int b_last, const skip_t & skip) {
+    constexpr int C = FATTN_FIXUP_CHUNK;
+
+    float  add[2][C];
+    float2 mk[2][C];
+    const auto load = [&](const int buf, const int b0) {
+#pragma unroll
+        for (int k = 0; k < C; ++k) {
+            const int b = b0 - k;
+            if (b >= b_first) {
+                add[buf][k] = data[b*ncols*D + jc*D + tid];
+                mk[buf][k]  = meta[b*ncols + jc];
+            }
+        }
+    };
+    const auto fold = [&](const int buf, const int b0) {
+#pragma unroll
+        for (int k = 0; k < C; ++k) {
+            const int b = b0 - k;
+            if (b < b_first || skip(b)) {
+                continue;
+            }
+            const float2 tmp = mk[buf][k];
+
+            const float max_val_new = fmaxf(max_val, tmp.x);
+
+            const float diff_val = max_val - max_val_new;
+            const float diff_add = tmp.x   - max_val_new;
+
+            const float scale_val = diff_val >= SOFTMAX_FTZ_THRESHOLD ? expf(diff_val) : 0.0f;
+            const float scale_add = diff_add >= SOFTMAX_FTZ_THRESHOLD ? expf(diff_add) : 0.0f;
+
+            dst_val = scale_val*dst_val + scale_add*add[buf][k];
+            rowsum  = scale_val*rowsum  + scale_add*tmp.y;
+
+            max_val = max_val_new;
+        }
+    };
+
+    int b0 = b_last - 1;
+    load(0, b0);
+    while (b0 >= b_first) {
+        if (b0 - C >= b_first) {
+            load(1, b0 - C);
+        }
+        fold(0, b0);
+        b0 -= C;
+        if (b0 < b_first) {
+            break;
+        }
+        if (b0 - C >= b_first) {
+            load(0, b0 - C);
+        }
+        fold(1, b0);
+        b0 -= C;
+    }
+}
+
 template<int D, int ncols1, int ncols2> // D == head size
 __launch_bounds__(D, 1)
 static __global__ void flash_attn_stream_k_fixup_uniform(
@@ -1130,24 +1199,8 @@ static __global__ void flash_attn_stream_k_fixup_uniform(
     }
 
     // Combine with all previous blocks in this tile.
-    for (int bidx = b_last - 1; bidx >= b_first; --bidx) {
-        const float dst_add = dst_fixup_data[bidx*ncols*D + jc*D + tid];
-
-        const float2 tmp = dst_fixup[(nblocks_stream_k + bidx)*ncols + jc];
-
-        const float max_val_new = fmaxf(max_val, tmp.x);
-
-        const float diff_val = max_val - max_val_new;
-        const float diff_add = tmp.x   - max_val_new;
-
-        const float scale_val = diff_val >= SOFTMAX_FTZ_THRESHOLD ? expf(diff_val) : 0.0f;
-        const float scale_add = diff_add >= SOFTMAX_FTZ_THRESHOLD ? expf(diff_add) : 0.0f;
-
-        dst_val = scale_val*dst_val + scale_add*dst_add;
-        rowsum  = scale_val*rowsum  + scale_add*tmp.y;
-
-        max_val = max_val_new;
-    }
+    flash_attn_fixup_fold<D, ncols>(dst_val, max_val, rowsum, dst_fixup_data, dst_fixup + nblocks_stream_k*ncols, jc, tid,
+        b_first, b_last, [](const int) { return false; });
 
     // Write back final result:
     *dst = dst_val / rowsum;
@@ -1329,28 +1382,10 @@ static __global__ void flash_attn_stream_k_fixup_live(
     // A block can have no unit of work only when there are fewer units than blocks; the test is two 64-bit divisions ahead of
     // every load, so it runs only then.
     const bool maybe_empty = fixup_legacy || total_work < nblocks;
-    for (int bidx = b_last - 1; bidx >= b_first; --bidx) {
-        if (maybe_empty && int64_t(bidx)*total_work / nblocks == int64_t(bidx + 1)*total_work / nblocks) {
-            continue; // Did not have any data.
-        }
-
-        const float dst_add = dst_fixup_data[bidx*ncols*D + jc*D + tid];
-
-        const float2 tmp = dst_fixup[(nblocks + bidx)*ncols + jc];
-
-        const float max_val_new = fmaxf(max_val, tmp.x);
-
-        const float diff_val = max_val - max_val_new;
-        const float diff_add = tmp.x   - max_val_new;
-
-        const float scale_val = diff_val >= SOFTMAX_FTZ_THRESHOLD ? expf(diff_val) : 0.0f;
-        const float scale_add = diff_add >= SOFTMAX_FTZ_THRESHOLD ? expf(diff_add) : 0.0f;
-
-        dst_val = scale_val*dst_val + scale_add*dst_add;
-        rowsum  = scale_val*rowsum  + scale_add*tmp.y;
-
-        max_val = max_val_new;
-    }
+    flash_attn_fixup_fold<D, ncols>(dst_val, max_val, rowsum, dst_fixup_data, dst_fixup + nblocks*ncols, jc, tid,
+        b_first, b_last, [&](const int b) { // did not have any data
+            return maybe_empty && int64_t(b)*total_work / nblocks == int64_t(b + 1)*total_work / nblocks;
+        });
 
     *dst = dst_val / rowsum;
 }