chore(train): engine-10 continued: pre-sampling probabilities on the prefilter, the live-tile attention's one copy, one scan and aligned split

TORAD.md

@@ -93,6 +93,10 @@ which pins a commit of this branch as a submodule.
 | `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) |
+| `f1e45c29c` | a request for pre-sampling probabilities (`n_probs` without `post_sampling_probs`, and every OpenAI `logprobs` request) turned the top-k prefilter off, since they are under the whole row: every decode copied each output row's n_vocab logits to the host, the CPU chain scanned them for its top k, and `get_token_probabilities` built a 248,320-entry vector and ran 248,320 `expf` a token for a softmax whose only use of the row is its sum. On an RTX 5090 at 245,752 tokens, greedy with `n_probs` 5 decoded 93.7 tok/s against 107.4 without. `llama_sampler_init_row_probs()` gives a row its softmax on the backend and changes no logit, and a backend top-k keeps probabilities a sampler before it produced aligned with its candidates; with at most k probabilities asked, the prefilter chain is [row-probs, top-k] and `get_token_probabilities` reads the k candidates' probabilities under the whole row. RTX 5080, Ternary Bonsai 2 27B served at -c 32768, n_probs 5, medians of 9: plain 95.65 -> 105.14 tok/s (105.92 without n_probs), drafted 166.20 -> 192.51 (192.15), the same tokens and top-5 ids; log-probabilities move by up to 1.9e-3, the CPU's float running sum over the row (the backend's softmax matches a double-precision one to 4e-7). test-backend-sampler `row_probs_top_k` (Qwen3 0.6B Q4_0: the backend's k probabilities against a CPU softmax of the full row to 1e-5; without top-k's gather it fails); `test_top_k_prefilter_pre_sampling_probs` (the same tokens and top-n ids, log-probabilities within 2e-5 of every logit on the CPU, measured gap 2.9e-6; the k's own softmax fails it) | `LLAMA_TOP_K_PREFILTER_LEGACY=1` keeps every logit on the CPU; more probabilities asked than the chain's k do too |
+| `818965d72` | the mma flash attention compiled its tile code once per stream-k role (a block that ends inside a tile, one that ends a tile it did not start, one that does a whole tile) as template parameters, though only the epilogue's destination depends on the role: Ternary Bonsai 2 27B's decode kernel (head 256, q4_0 K/V read raw) was 184 KB of SASS in three copies. A per-block timeline on an RTX 5080 (252 blocks, 63 a KV head) had the 4 blocks that end a tile ending last, 8-9 us after the median block, their prologues 7.8-10.5 us against 3.7-4.2: only they ran the tile-ending copy, cold in every cache after 75-283 MB of K/V through the 64 MB L2. The roles are arguments of one copy (71 KB, 255 registers, 16 bytes of stack against 24; libggml-cuda 79.1 -> 73.2 MB). test-backend-ops perf, the served layout, 4 rounds against the parent's library: a token -2.8 / -4.2 / -2.2 / -0.9 % and a 4-row verify +0.4 / -3.1 / -1.0 / -1.1 % at 16,384 / 65,536 / 131,072 / 245,760 cells; llama-bench at 16,384 inside its noise. FLASH_ATTN_EXT 3,220/3,220; the served decode's texts identical | none: the same arithmetic and stores, only the code's layout changed |
+| `6aa14679c` | the live-tile split (serving with more than one slot and `--kv-unified`) ran a memset and a scan of the mask before every flash attention, 3.2-4.5 us a call in the graph on an RTX 5080, though the live steps depend on the mask and the split alone and every attention layer of a graph reads the same mask (Ternary Bonsai 2 27B: 16 layers). `ggml_cuda_fattn_kv_live_context` keeps a graph evaluation's first scan, keyed by the mask tensor and its data, shape and strides, the stream and the split; the other layers read it and any other key scans again; it resets where the graph evaluator resets its other per-graph records, so a CUDA graph captures one scan. test-backend-ops perf (its graph repeats the op: only the first copy scans), 4 rounds: a token -13.2 / -3.6 / -2.0 / -1.1 % and a 4-row verify -13.5 / -3.7 / -1.2 / -0.7 % at 16,384 / 65,536 / 131,072 / 245,760 cells (3.3-4.2 us a call). FLASH_ATTN_EXT 3,220/3,220; the served decode at 15,616 tokens, 4 slots: the same tokens in every leg, plain and drafted, its rates inside a loaded host's noise | `GGML_CUDA_FATTN_LIVE_SCAN_EACH_LEGACY=1` |
+| `32e695ecf` | the live-tile split launched every block the SMs hold, where the plain stream-k split rounds down to a multiple of the output tiles. 01f4f0fda's 3 blocks an SM made a decode token's split 510 blocks on an RTX 5090 (127.5 per KV head) and 210 on an RTX 5070 Ti (52.5): each head's blocks started half a block's steps from its neighbour's. A cell's K and V rows hold the heads side by side (q4_0, head 256: 144 bytes each) and the L2 fetches 64-byte units, so neighbouring heads share one, fetched once only when they read the row together: 12 units a row where 9 hold it. ncu on the 5070 Ti at 131,072 cells: DRAM read 202-207 MB against 151.0 MB of K and V; rounded down to a multiple of the Q tile's output tiles (208 blocks), 152-153 MB. test-backend-ops perf, a token, the same library with and without the switch: 5090 -15.7 / -7.8 / -17.0 / -20.9 %, 5070 Ti -2.1 / -18.3 / -21.4 / -23.1 % at 16,384 / 65,536 / 131,072 / 245,760 cells; on the 5090 131,072 and 245,760 cells take 115.72 and 196.44 us, where 4104c47 (340 blocks) took 123.69 and 204.66 and 4b61c54 145.50 and 255.48. A 4-row verify (2 blocks an SM: 340 and 140 blocks) does not move; the 5080's 252 is a multiple already. The served decode on the 5090 at 245K tokens (the gate's pool, 4 x 294,912 with `--kv-unified`, its four questions, 3 legs each): plain 98.75 -> 107.92 tok/s (+9.3 %), drafted 205.43 -> 208.90 (+1.7 %), the texts first differing 35-132 tokens in: the split sums in another order. Held to engine-10's G1 bars at that shape with top-5 log-probabilities, 384 tokens: plain 4 first differences, each at a tie, |dlogprob| p99 0.095 and max 0.128; drafted 4, each at a tie, p99 0.070, max 0.078. The gate's one-slot legs take the live-tile path too: their bits move the same way and their rate does not (101.5 against 100.7 tok/s at 245K). FLASH_ATTN_EXT 3,220/3,220 on the 5070 Ti and the 5080 | `GGML_CUDA_FATTN_LIVE_BLOCKS_ALIGN_LEGACY=1` |
 
 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-cuda/common.cuh

@@ -6,6 +6,7 @@
 
 #include <cstdint>
 #include <cstdlib>
+#include <cstring>
 #include <memory>
 #include <mutex>
 
@@ -1530,6 +1531,70 @@ struct ggml_cuda_gdn_gather_context {
     }
 };
 
+// The flash attention's live KV steps (flash_attn_mask_to_KV_live) for the graph being evaluated: every attention layer of a
+// graph reads the same mask, so the first one scans it and the others read what it found. What the scan read and how it
+// split the steps is the key; a different mask, shape or stream scans again. The memory is the context's and outlives every
+// graph evaluation: a CUDA graph captured with the scan writes it again on each replay, before the layers after it read it.
+// A buffer too small for a later scan is kept, not freed, since a graph captured earlier may still name it.
+struct ggml_cuda_fattn_kv_live_context {
+    struct key_t {
+        const void * mask;
+        const void * mask_data;
+        cudaStream_t stream;
+        int64_t      ne[4];      // the mask's
+        size_t       nb[4];      // the mask's
+        int          geometry[6]; // ncols1, nbatch_fa, iter_k, the Q tiles, the output tiles per Q tile, the sequences
+
+        bool operator==(const key_t & o) const {
+            return mask == o.mask && mask_data == o.mask_data && stream == o.stream &&
+                memcmp(ne, o.ne, sizeof(ne)) == 0 && memcmp(nb, o.nb, sizeof(nb)) == 0 &&
+                memcmp(geometry, o.geometry, sizeof(geometry)) == 0;
+        }
+    };
+
+    bool                valid = false; // filled in this graph evaluation
+    key_t               key   = {};
+    int *               ptr   = nullptr;
+    size_t              size  = 0;     // ints
+    std::vector<int *>  kept;          // smaller buffers graphs captured before may still name
+
+    void reset() {
+        valid = false;
+    }
+
+    // the steps a scan of this key found in this graph evaluation, or nullptr
+    const int * find(const key_t & k) const {
+        return valid && k == key ? ptr : nullptr;
+    }
+
+    // n ints for a scan of this key to fill, found by the next find() with the same key until reset()
+    int * fill(const key_t & k, const size_t n) {
+        if (n > size) {
+            if (ptr != nullptr) {
+                kept.push_back(ptr);
+            }
+            CUDA_CHECK(cudaMalloc(&ptr, n*sizeof(int)));
+            size = n;
+        }
+        key   = k;
+        valid = true;
+        return ptr;
+    }
+
+    void release() {
+        for (int * p : kept) {
+            CUDA_CHECK(cudaFree(p));
+        }
+        kept.clear();
+        if (ptr != nullptr) {
+            CUDA_CHECK(cudaFree(ptr));
+            ptr = nullptr;
+        }
+        size  = 0;
+        valid = false;
+    }
+};
+
 // Fused conv-state update for the GDN's causal conv (build_conv_state): GET_ROWS(conv cache, s_copy) -> RESHAPE ->
 // CONCAT(state, transposed new inputs) -> the rollback snapshots' CPYs into the cache, and the SSM_CONV reading the
 // CONCAT. The graph evaluator skips the GET_ROWS, the CONCAT and the CPYs (ggml_cuda_try_ssm_conv_state_update) and
@@ -1725,6 +1790,7 @@ struct ggml_backend_cuda_context {
     ggml_cuda_stream_context concurrent_stream_context;
     ggml_cuda_gdn_gather_context gdn_gather_context;
     ggml_cuda_ssm_conv_update_context ssm_conv_update_context;
+    ggml_cuda_fattn_kv_live_context fattn_kv_live_context;
     ggml_cuda_pq2_prefetch pq2_next; // for the node being dispatched
 
     ~ggml_backend_cuda_context();
@@ -1745,6 +1811,8 @@ struct ggml_backend_cuda_context {
 
     ggml_cuda_ssm_conv_update_context & ssm_conv_updates() { return ssm_conv_update_context; }
 
+    ggml_cuda_fattn_kv_live_context & fattn_kv_live() { return fattn_kv_live_context; }
+
     cublasHandle_t cublas_handle() {
         if (cublas_handles[device][curr_stream_no] == nullptr) {
             ggml_cuda_set_device(device);