chore(train): engine-10 continued: the PQ2_0 tensor-core launches take their tiles from a counter

TORAD.md

@@ -99,6 +99,11 @@ which pins a commit of this branch as a submodule.
 | `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 prefill through the live-tile path (the mask-prefix hint covers at most 16 rows), so their bits move too; their decode takes the mask-prefix path, whose split was rounded already, and its rate does not move (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` |
 | `e75e318e5` | `llama-bench` puts an NVTX range `gen` (domain `llama-bench`, a registered string) around each repetition's generation, so `nsys profile -c nvtx -p gen@llama-bench --capture-range-end=stop` records the tokens and none of the depth's prefill before them; a build without the NVTX headers compiles it out. rig's roofline probe traced the whole run, and at a depth of 245,760 on an RTX 5080 (Ternary Bonsai 2 27B, graphs off, `GGML_CUDA_NVTX=1`) the prefill's ~1.19M kernels kept nsys past its 30-minute limit; started at the range the capture takes 179 s, a 2.3 MB report of the 16 tokens' 18,816 kernels. At 16,384 it reads what the full capture reads: kernels 9,553.6 against 9,559.8 us a token, the same 23 groups, the largest group apart by 6.3 us | (a profiler that does not ask for the range ignores it) |
 | `fb03d294b` | `rms_norm_fwht_cuda` and `fwht_cuda_block` request their weights (the norm's `w` and the signs, which no kernel writes) before the PDL dependency wait, as the PQ2_0 matmul does its own. The matmul after a rotation starts under it (PDL), and its ring fill and the L2 prefetch of its head (`16ad036c4`) take DRAM for the microseconds the rotation runs; read after the wait, the rotation's weights (DRAM misses each token) queued behind them. RTX 5080, Ternary Bonsai 2 27B, one build both ways: nsys over 15 tokens (graphs off, the capture started at `e75e318e5`'s range), the FFN-input norm after an out-projection ends 6.1-6.2 -> 4.7-4.8 us past it at 16,384 and 5.5-6.1 -> 4.5-4.8 at 245,760, the layer-input norm 2.1 -> 1.7 us; llama-bench with graphs on, legs N-L-L-N: tg128 at 16,384 +0.81 % (105.48 / 104.21 / 104.40 / 104.82 tok/s), tg64 at 245,760 +1.63 % (69.35 / 67.93 / 67.48 / 68.27), all four pairs faster; `test-backend-ops MUL_MAT_HADAMARD` 42/42 both ways; llama-server at one slot and at 4 slots with `--kv-unified`, 128 greedy tokens with top-5 log-probabilities: bit-identical to the published engine-32e695e both ways | `GGML_CUDA_FWHT_PREWAIT_LEGACY=1` |
+| `8c78ba9ca` | The PQ2_0 tensor-core launches (`mmvq-pq2-mma.cu`) keep a block's first tiles, as many as its ring holds before the dependency wait, and take the rest from a per-stream counter after the wait (`ggml_cuda_pq2_tile_counters`): each block's last ticket is past the tiles, and the block that takes the launch's last sets the counter back to 0. Owned tiles left a launch's last microseconds to the blocks the DRAM served later; now the blocks end within a tile of each other. RTX 5080, Ternary Bonsai 2 27B, one build both ways, graphs on, legs N-L-L-N: tg128 +1.46 % at 0 and +1.34 % at 16,384, tg64 +1.47 % at 245,760, every pair faster; pp4 `-rs 3` (the MTP verify), medians of 30 samples, +0.62 % and +0.81 %; llama-server with MTP drafts of 3, 12 greedy agentic requests: the same text, +1.14 % (same-text geomean +1.8 %) against the undisturbed legacy leg; `test-backend-ops` MUL_MAT (pq2_0) 99/99, MUL_MAT_GROUP 24/24, MUL_MAT_VEC_FUSION 1010/1010 both ways, and without the counter's reset 29, 7 and 26 of them fail; llama-server at one slot and at 4 with `--kv-unified`: bit-identical to the published engine-32e695e both ways | `GGML_CUDA_PQ2_MMA_TILES_LEGACY=1` |
+| `833d067e0` | glm5next leaves the `-sm tensor` refusal list and splits dense experts, routed experts, shared experts (as DeepSeek4) and the head, while every KDA and DSA weight, all ssm and indexer tensors and all four caches stay mirrored: one all-reduce per layer past the expert adds. `test-llama-archs -a glm5next` on 5080 + 5070 Ti: Meta row OK at 2.82e-06 NMSE vs CPU (1e-4 bar). The ffn_down_exps-mirrored mutant aborts the Meta row in the matmul split-state handler, so the green row really splits. Per-card weights from the GGUF tensor table: 66.15 GiB plus mirrored latent and indexer caches. | — |
+| `881c823c5` | `ggml_cuda_op_top_k` with a k of 64 or more past 1,024 columns selects by radix (`topk_radix`): four 8-bit histogram passes find the k-th largest key, then one ordered pass writes every column above it and the lowest columns equal to it, the tiled path's set and tie-break. Up to 16,384 columns a block keeps its row in registers; a wider row is cut into 8,192-column tiles selected in parallel, then one block a row selects among their candidates; one wide row stays with CUB's top-k. The GLM-5.3 DSA indexer (512 pools a layer) paid the tiled path's k serial block reductions: RTX 5080, k 512, 3,520 columns 194.5 -> 6.1 us, 109,020 x 3 66.3 -> 21.8, a 512-row ubatch 13.8-18x. TOP_K 523/523; two mutants fail it (the last equal key unwritten; stage 2 writing a candidate's position). | `GGML_CUDA_TOPK_RADIX_LEGACY` |
+| `33eb70bc0` | `build_attn_mha` takes `mask_is_prefix` (true by default) and the two sparse-attention callers (the DSA layers' top-k mask, and the kpool path's) pass false: every one-sequence causal mask was tagged as a prefix of the cells, the hint under which the CUDA flash attention skips its range scan and live tiles and applies the mask over the whole range, so a GLM-5.3 DSA decode read every cell's K and V to keep the indexer's 2,048. The dense layers keep the hint. Same attention by another split of the cells, bits can move at float rounding; the gain sits at long caches (436k) and is measured on the served model. | `LLAMA_ATTN_SPARSE_MASK_PREFIX_LEGACY` |
+| `78e0fadf7` | `topk_radix` also takes a row of up to 1,024 columns at any k: with CUB's top-k available such a row had no path but CUB one row after another, four launches a row (a DSA indexer's ubatch under 4,096 cells: 2,048 launches a layer). RTX 5080: 3.3-3.9 us against 8.2-12.7 (1 row) and 129-171 (16 rows). | `GGML_CUDA_TOPK_RADIX_LEGACY` |
 
 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

@@ -1512,6 +1512,28 @@ struct ggml_cuda_pq2_prefetch {
     int64_t      bytes = 0;
 };
 
+// One tile counter a stream for the PQ2_0 tensor-core launches (mmvq-pq2-mma.cu): past its own first tiles a block takes
+// the next by an atomic add on it, and the block that takes the launch's last ticket sets it back to 0. A launch takes
+// tickets only after its dependency wait, when every launch before it on the stream has ended, so one counter serves
+// every launch of a stream, CUDA graph replays included. Made zeroed before a graph evaluation, never inside a capture.
+struct ggml_cuda_pq2_tile_counters {
+    int * ptr = nullptr; // GGML_CUDA_MAX_STREAMS ints
+
+    void ensure() {
+        if (ptr == nullptr) {
+            CUDA_CHECK(cudaMalloc(&ptr, GGML_CUDA_MAX_STREAMS * sizeof(int)));
+            CUDA_CHECK(cudaMemset(ptr, 0, GGML_CUDA_MAX_STREAMS * sizeof(int)));
+        }
+    }
+
+    void release() {
+        if (ptr != nullptr) {
+            CUDA_CHECK(cudaFree(ptr));
+            ptr = nullptr;
+        }
+    }
+};
+
 // Owned by the backend context that evaluates the graph: registrations are keyed by node pointer, so they only mean
 // something for the evaluation that made them. Cleared at the start of every graph evaluation/capture.
 struct ggml_cuda_gdn_gather_context {
@@ -1792,9 +1814,13 @@ struct ggml_backend_cuda_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_cuda_pq2_tile_counters pq2_tile_counters;
 
     ~ggml_backend_cuda_context();
 
+    // the current stream's PQ2_0 tile counter, or nullptr before the first graph evaluation made them
+    int * pq2_tile_counter() { return pq2_tile_counters.ptr != nullptr ? pq2_tile_counters.ptr + curr_stream_no : nullptr; }
+
     cudaStream_t stream(int device, int stream) {
         if (streams[device][stream] == nullptr) {
             ggml_cuda_set_device(device);

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -738,6 +738,10 @@ ggml_backend_cuda_context::~ggml_backend_cuda_context() {
         ggml_cuda_set_device(device);
         fattn_kv_live_context.release();
     }
+    if (pq2_tile_counters.ptr != nullptr) {
+        ggml_cuda_set_device(device);
+        pq2_tile_counters.release();
+    }
     for (int i = 0; i < GGML_CUDA_MAX_DEVICES; ++i) {
         for (int j = 0; j < GGML_CUDA_MAX_STREAMS; ++j) {
             if (streams[i][j] != nullptr) {
@@ -5494,6 +5498,7 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend,
     ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context;
 
     ggml_cuda_set_device(cuda_ctx->device);
+    cuda_ctx->pq2_tile_counters.ensure(); // before a capture can begin
 
     bool use_cuda_graph             = false;
     bool cuda_graph_update_required = false;