diff --git a/TORAD.md b/TORAD.md index dbdce4c00e21..fc9f51a2e9d2 100644 --- a/TORAD.md +++ b/TORAD.md @@ -153,7 +153,7 @@ which pins a commit of this branch as a submodule. | `98a6656ae` | `4bce9401c`'s comb beside the stream was measured on node traces only (each card's kernel time a token -1.03 / -0.69 % under `-sm tensor`); GLM-5.3-Flash on 2 RTX PRO 6000 under `-sm tensor` ran tg64 3.7 % faster with `GGML_CUDA_HC_COMB_SIDE_LEGACY=1` in 4 of 4 interleaved pairs (rig-glm, `2e699cceb`), and on the 44-layer proxy the stream is only +0.70 % over it (95 % CI -0.15 to +1.55, tg64, 4 rotated rounds). By default the front's second kernel now makes the comb on the chain, as under the old switch, which is gone; `GGML_CUDA_HC_COMB_SIDE=1` makes it beside the stream as before (`4bce9401c`, `c6b4ecf1b`). PPL bit for bit each way (387892.2627, 386825.5220, 388534.0622); a tg node trace under `-sm tensor` launches no `dsv4_hc_comb_side` by default and 88 an evaluation with the switch, one a front (declared 100, from `c6b4ecf1b`'s row: the proxy's graph has 88 fronts, 2 a layer); `test-backend-ops` `DSV4_HC_POST`, `DSV4_HC_PRE_FUSED`, `DSV4_HC_PRE_POST`, `DSV4_HC_PRE_Q8_1` pass on both cards each way | `GGML_CUDA_HC_COMB_SIDE=1` (opt in: the comb beside the stream) | | `304cc84cc` | `98a6656ae` left `4bce9401c`'s comb beside the stream opt-in and no default ran it, while it kept a stream, two events, a list of pending weights every node was checked against, a join in every `DSV4_HC_POST` and at each evaluation's end, and `c6b4ecf1b`'s check that the weights outlive the front. They are gone with the side kernel and `GGML_CUDA_HC_COMB_SIDE`: `dsv4-hc.cu` and `dsv4-hc.cuh` are as before `4bce9401c`, and the front's second kernel makes the comb, as by default since `98a6656ae`. Kept: `6588e1d6a`'s order in `ggml_backend_compare_graph_backend` (the reference first, so work a backend leaves running past its return reads as a wrong result) and `DSV4_HC_PRE_POST`'s 1,000-iteration cases, now catching a read of the comb before the front has written it. The CUDA library holds no `dsv4_hc_comb_side` (strings: 0 lines, 8 in `5b9b8e563`'s); PPL bit for bit from the landed build (387892.2627 under `-sm tensor`, 386825.5220 at ub 1, 388534.0622 at ub 3); the scratch gate's `test-backend-ops` sets (`MUL_MAT` 1475/1475, `MUL_MAT_ID` 1076/1076, the `DSV4_HC_*` sets, `TOPK_MOE` 320/320), `test-llama-archs` and the meta tests pass | none (removed; `GGML_CUDA_HC_COMB_SIDE` is no longer read) | | `0b7316af7` | `llama_kv_cache_set_input_kpool` made about eight passes over the cache's cells each token (the pool range, the positions, each cell's pool, its completeness, the two mask rows) to rebuild the k-pool maps and masks of GLM-5.3-Flash's indexer; at 64K cached tokens it was the largest host stack of the server's main thread (rig-glm, `2e699cceb`, gdb, 14 % of the active samples were host work). `llama_kv_cells` now logs the cells its mutators change (`llama_kv_cells_log`, one consumer; a reset, a resize, an assignment or a list past 1/16 of the cells makes it rebuild) and counts each sequence's cells; `llama_kpool_views`, kept by the cache, holds each sequence's positions, position -> cell table, pool fill and complete pools, follows the log, and writes the maps from the tables, byte for byte as before. Not served from a view (the maps come from the cells as before): two cells at a position, positions spanning more pools than a map holds, used cells past n_kv, several sequences of one stream in a ubatch, more than 8 sequences, a view whose cell count differs from the cells' own. `llama_kpool_set_input` at 1 token, kpool 4, per call, views against cells: 2048 cached tokens 1.1 against 12.9 us, 8192 3.8 against 49.1, 32768 16.5 against 223.7, 65536 37.7 against 458.4; 3 tokens at 65536 77.2 against 449.1. The 44-layer proxy under `-sm tensor` at depth 32768, 3 interleaved pairs, tg64: 85.37, 84.22, 84.93 against 83.69, 83.12, 84.79 tok/s (+2.0, +1.3, +0.2 %; the microbenchmark's 207 us of an 11.8 ms token predicts 1.7 %, computed after the result); pp3 at `-ub 3` 111.2, 108.4, 122.2 against 103.1, 116.1, 98.8, too spread to read (another job held both cards at 98 % load). `LLAMA_KPOOL_INPUT_CHECK=1` on the 44-layer proxy under `llama-server -c 4096`: 6 requests with a cached prompt reused by LCP, a tail trimmed, a slot reset and a context shift, 1536 checked calls, all from views, 23 rebuilds, no difference; with `--spec-type draft-mtp --spec-draft-n-max 3` 8192 checked calls, no difference. `test-kpool-input`: 420 random rounds (append, rollback, mid-range removal, `seq_cp`, prepare's save and restore, shift and divide, duplicate and sparse positions, clear, copy-over; unified and per-stream; f16 and f32; with and without the key cache) build the 7 input tensors from one long-lived view and from the cells and compare every byte; the log is checked against a snapshot of the cells; a decode loop with rollbacks, and one that moves the newest token into a lower hole, rebuilds the view once. Mutants that fail it: a mutator not logging (`seq_rm`, `pos_set`, `set`, `seq_keep`, `pos_add`), `reset` not telling the consumer, a removal leaving the pool's mark, an insertion not seeing a taken position, the count guard, a view served with no list, the padding row, the run's base, a pool's rep, removals and insertions interleaved. `test-kv-mask` and `test-kpool-can-reuse` pass | `LLAMA_KPOOL_INPUT_LEGACY=1` (the maps from the cells every call); `LLAMA_KPOOL_INPUT_CHECK=1` builds both each call and aborts on a differing byte | -| `0b0e97b5d` | GLM-5.3-Flash's DSA layers mask the whole cache down to the cells their indexer picked (2048 and a tail of at most 3 on the proxy), and the MMA kernel read every cell under the mask: at 32K cached tokens the flash attention node took 595 / 667 us a decode token on the RTX 5080 / 5070 Ti under `-sm tensor`. Ported by hand from ggml-org/llama.cpp `8e93a9773` (#27970, Aman Gupta): FLASH_ATTN_EXT takes a bound `n_kv_max` on each mask row's finite cells (`ggml_flash_attn_ext_set_n_kv_max`, `op_params[5]`; `[4]` is this fork's mask prefix), a kernel compacts each row's finite cells into an index list, and the MMA kernel for `<512,512,1,8>` and `<576,512,1,16>` loads K, V and the mask by index over `n_kv_max` cells, one query a tile. `build_attn_sparse` passes `GGML_PAD(top_k + kpool - 1, 32)` (2080; the pad keeps a tail made longer by two cells at one position, rig-glm), the inherited DSA path `top_k`. Taken when the cache is at least 4096 cells and `n_gather = min(n_tokens, 64/ncols2) * n_kv_max`: a `test-backend-ops perf` grid of GLM's shape (32 heads on the latent, 8448 / 16640 / 33280 cells, batch 1 to 512, both cards) puts the gather at 0.25-0.99x the dense time wherever K >= `n_gather` and 1.13-1.27x under it (upstream asks 2 `n_gather`, which gave up a 3-token verify at 8K, 0.73x, and prefills at 16K, 0.63-0.72x). The 44-layer proxy, both cards, `-sm tensor`: the node a decode token 260/259 against 407/446 us at `-d 8192`, 337/335 against 595/667 at `-d 32768` (the kernel is 17 us a layer at both, the rest is the compaction, 3.6 and 11.2 us a layer), unchanged at `-d 2048` (under the gate); tg32 @ d32768 reps 2-6 of a process 89.3 against 88.2 t/s (+1.2 %; the first rep is slower each way, more so here, 73.9 against 78.2, and 3 pairs of `-r 3` means read 79.3/75.6/82.2 against 81.1/82.1/80.3); pp512 +26.7 % at `-d 32768` (1235.6 against 975.0), dense each way at `-d 8192`. PPL at `-c 8192` with the sparse kernel on every batch 348974.1045 against 349039.3242 (0.019 %), KLD 0.0016. `test-backend-ops` `FLASH_ATTN_EXT` 3229/3229 and the sparse cases pass on both cards each way; mutants that fail them: the gather reading the dense row (7 of 7), the compaction dropping each row's last index (upstream's 4) | `GGML_CUDA_FATTN_SPARSE_LEGACY=1` (the whole cache under the mask) | +| `0b0e97b5d` | GLM-5.3-Flash's DSA layers mask the whole cache down to the cells their indexer picked (2048 and a tail of at most 3 on the proxy), and the MMA kernel read every cell under the mask: at 32K cached tokens the flash attention node took 595 / 667 us a decode token on the RTX 5080 / 5070 Ti under `-sm tensor`. Ported by hand from ggml-org/llama.cpp `8e93a9773` (#27970, Aman Gupta): FLASH_ATTN_EXT takes a bound `n_kv_max` on each mask row's finite cells (`ggml_flash_attn_ext_set_n_kv_max`, `op_params[5]`; `[4]` is this fork's mask prefix), a kernel compacts each row's finite cells into an index list, and the MMA kernel for `<512,512,1,8>` and `<576,512,1,16>` loads K, V and the mask by index over `n_kv_max` cells, one query a tile. `build_attn_sparse` passes `GGML_PAD(top_k + kpool - 1, 32)` (2080; the pad keeps a tail made longer by two cells at one position, rig-glm), the inherited DSA path `top_k`. Taken when the cache is at least 4096 cells and `n_gather = min(n_tokens, 64/ncols2) * n_kv_max`: a `test-backend-ops perf` grid of GLM's shape (32 heads on the latent, 8448 / 16640 / 33280 cells, batch 1 to 512, both cards) puts the gather at 0.25-0.99x the dense time wherever K >= `n_gather` and 1.13-1.27x under it (upstream asks 2 `n_gather`, which gave up a 3-token verify at 8K, 0.73x, and prefills at 16K, 0.63-0.72x). The 44-layer proxy, both cards, `-sm tensor`: the node a decode token 260/259 against 407/446 us at `-d 8192`, 337/335 against 595/667 at `-d 32768` (the kernel is 17 us a layer at both, the rest is the compaction, 3.6 and 11.2 us a layer), unchanged at `-d 2048` (under the gate); tg32 @ d32768 reps 2-6 of a process 89.3 against 88.2 t/s (+1.2 %; the first rep is slower each way, more so here, 73.9 against 78.2, and 3 pairs of `-r 3` means read 79.3/75.6/82.2 against 81.1/82.1/80.3); pp512 +26.7 % at `-d 32768` (1235.6 against 975.0), dense each way at `-d 8192`. PPL at `-c 8192` with the sparse kernel on every batch 348974.1045 against 349039.3242 (0.019 %), KLD 0.0016. `test-backend-ops` `FLASH_ATTN_EXT` 3229/3229 and the sparse cases pass on both cards each way; mutants that fail them: the gather reading the dense row (7 of 7), the compaction dropping each row's last index (upstream's 4) | `GGML_CUDA_FATTN_SPARSE_LEGACY=1` (the dense kernel, which reads the live tiles of `2f0ca785a`; with `GGML_CUDA_FATTN_LIVE_TILES_LEGACY=1` as well for the whole cache under the mask) | | `4a88ac993` | `ggml_compute_forward_set_rows` splits a scatter's slots over threads, so two slots naming one row are a data race on the CPU backend (the same value, but a race), and GLM-5.3-Flash's indexer had two such scatters. The sparse mask: a pool the top-k took with no finite score (fewer than `select_k` pools live) carried cell 0 x r for a padded pool, or cells overlapping the tail; `build_indexer` now returns each slot's liveness (exp of the pool bias at the selected pools) and `build_attn_sparse` sends a filler slot to its own dump column past `n_kv`, in a mask widened by those columns and viewed back to `n_kv` before the cand and kq masks. The pooled-key write (`llama_kpool_set_input`): an unused slot of the fixed-size write repeated a complete pool's row; it now names a spare cell, one each (empty, or not the last of its block, and not a cell a real slot writes). Red on `0b7316af7`: `test-llama-archs -a glm5next` under a CPU ThreadSanitizer build reports 3 races in `ggml_compute_forward_set_rows`, all in the indexer key cache, and `test-kpool-input`'s new check (the rows of the write differ, and a slot naming a block's last cell computes that cell's pool) fails 20 rounds; green here. No value the model reads changes: a filler wrote 0 where the mask held 0 already or where `cand_mask` masks it again. On the 44-layer proxy, one card, every row of the final ubatch at 512 and 2048 tokens hashed per tensor: the masks, the indexer's scores and every layer's attention output are bit for bit the parent's, and with `GGML_CUDA_DISABLE_FUSION=1` on both builds so is everything to the logits. With fusion on the logits part at `-ub 512` (`-c 8192`, one chunk: mean KLD 5.63e-4, same top token 94.07 %; the parent against itself 0) and nowhere at `-ub 3` or `-ub 1`, first at layer 7's `ffn_out` in 112 of 512 rows: the larger compute buffer moves the tensors a CUDA fusion's address check (`ggml_cuda_check_fusion_memory_ranges`) reads, and that MoE fuses in one build and not the other (with the weighted-sum fusion alone off they still part there). At decode under `-sm tensor`: tg32 at `-d 32768`, two interleaved pairs, medians of reps 2-6, 85.79 / 80.11 and 84.92 / 84.13 t/s against the parent (-3.6 % pooled; the pairs disagree); `nsys`: 128 more launches and about 146 us of kernel time a token a card, and a second meta graph a token, 30-46 us after the first: `build_indexer` reshapes `pool_bias` in every DSA layer and each view is a split input of its own (`GGML_SCHED_DEBUG`: the parent's meta split takes 24 inputs, 11 of them the per-layer `pool_cells` reshapes; this one 30 and 5, the cap `GGML_SCHED_MAX_SPLIT_INPUTS`) | — | | `69b9332d2` | The meta backend (`-sm tensor`) derives a node's split from its sources', and `handle_generic` returned `UNKNOWN` for a node with no sources, so allocating the graph aborted on `GGML_ASSERT(ret.axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN)`. ARANGE routes there and the sparse mask's dump columns (`4a88ac993`, `build_attn_sparse`) are an arange, so every GLM-5.3 decode under `-sm tensor` aborted in `ggml_gallocr_alloc_graph`. Every device computes the same values for such a node: it is mirrored. Upstream master carries the same `UNKNOWN`. `test-backend-meta-sourceless` computes `y = x + arange(n)` through the meta backend over one GPU twice: red on the parent's `libggml-base` (the assert, exit 134), green here (OK, exit 0) | — | | `9c506a40d` | `process_tile`'s combine for `np > 1` held one `__syncthreads()` inside `threadIdx.y % np == 0` and another in the `else` branch — a barrier the warps reached at different instructions. Every warp now reaches one barrier, with the combine and the write-back in the `np == 0` warps around it. Ported by hand from ggml-org/llama.cpp `b74f590ea` (#27870, Siavash Norouzi); no arithmetic changes. Against `c986f0c1c`'s build: compute-sanitizer synccheck on the RTX 5070 Ti gave 3680 errors before and 0 after on upstream's repro (hsk=192 hsv=128 nh=4 [8,1] kv=512 nb=3), and 512 before and 0 after on GLM-5.3's cases (512/512 [64,1] 8192, sparse, nb 1 and 3); `FLASH_ATTN_EXT` 3229/3229 on both cards; PPL bit for bit (386825.5220 at `-c 256 -ub 1`, 388534.0622 at `-ub 3`, 387892.2627 under `-sm tensor`, 349039.3239 at `-c 8192 -sm tensor`); `flash_attn_ext_f16` by ncu at base clocks with L2 flushed over GLM's grid (3 cache lengths × 8 batches, sparse and dense, 48 cells) within 1.79 % of base, against the instrument's own 1.44 % base-against-base spread. The gate on this sha, built in a detached scratch worktree and run from its own bin: 13 `test-backend-ops` suites green (`MUL_MAT` 1475, `MUL_MAT_ID` 1076, `MUL_MAT_VEC_FUSION` 1032, `GATED_DELTA_NET` 71, `GATED_DELTA_NET_CACHE_FUSION` 77, `DSV4_HC_POST` 4, `DSV4_HC_PRE_FUSED` 37, `LORA_RANK1` 3, `DSV4_HC_PRE_POST` 5, `DSV4_HC_PRE_Q8_1` 16, `MUL_MAT_PAIR` 120, `TOPK_MOE` 320, `FLASH_ATTN_EXT` 3229), the cache-fusion persist leg at `pdl=1` and `pdl=0`, `test-llama-archs`, all four meta tests, and the 44-layer proxy under `-sm tensor` with CUDA graphs on and off (4 of 4 rows, 0 asserts) | — | @@ -181,6 +181,18 @@ which pins a commit of this branch as a submodule. | `12569c3b0` | `GGML_CUDA_COUNT_INFLIGHT=1` prints, before each replay launch of a captured `-sm tensor` evaluation, how many of that device's replays are launched and not finished (an event recorded after each launch, polled with `cudaEventQuery`). A diagnostic, off unless set: it read 0 at every launch in the deadlock `e09841eff` fixes, which ruled out the host running ahead. | unset | | `e09841eff` | Under `-sm tensor` one host thread launched a replay on device 0 and then on device 1. A replay's graph can hold more work than a device's launch queue, and `cudaGraphLaunch` then blocks inside the call until the device drains; device 0 drained only once device 1 reached the graph's first all-reduce, and device 1's launch was behind the blocked thread, so the host deadlocked: the stall in `737ee4450`'s row, which NCCL documents for any one thread launching graphs on several GPUs ("Using NCCL with CUDA Graphs"). Replays now launch device 0 from the calling thread and every other device from its own thread, and return once every launch returned, so what the caller queues next follows the replay on every device. 2x RTX PRO 6000 Blackwell, GLM-5.3-Flash IQ3_XXS, `-sm tensor -ts 1,1 -fa on -b 9 -ub 9`, q8_0 K/V/idx, internal all-reduce, one perplexity chunk, with `GGML_CUDA_COUNT_INFLIGHT=1`: before, 8k and 16k stop with device 0 blocked in its 411th replay launch after device 1's 410th, nothing in flight on either device at any launch, the same under `CUDA_MODULE_LOADING=EAGER`, while 2k (217 launches a device) and 2k over 3 chunks (651) pass and `CUDA_SCALE_LAUNCH_QUEUES=4x` lets 16k finish (1754 launches): the stall is a graph at that depth overflowing the queue. With this change 16k finishes in 80 s (PPL 2.6523, the 4x run's), 8k in 58 s, 16k at `-ub 16` in 74 s; with the switch, 16k stops at the same launch again. `test-backend-meta-capture` passes by default, with the switch and under `GGML_META_CAPTURE_LEGACY=1`. KLD against `0090733` on the real pack, 2k x 32 and 16k x 4 chunks at `-ub 1024` and at `-ub 9` (the 16k `-ub 9` base under `GGML_META_CAPTURE_LEGACY=1`, since `0090733` cannot replay there): 0.000000 in all four, same top p 99.997 / 100.000 / 99.997 / 100.000 %, the instrument's own offset. `test-backend-ops` on its libraries against the CPU, card 0: MUL_MAT 1475, MUL_MAT_ID 1096, FLASH_ATTN_EXT 3260, LIGHTNING_INDEXER 183, GET_ROWS 119, CPY 277, RMS_NORM 57, SOFT_MAX 212, ROPE 466, all passed. Against the switch on the same library, `-sm tensor -ts 1/1 -fa on`, back-to-back pairs alternating: pp512 at `-b 9 -ub 9` +0.01 % (3 pairs, 95 % CI -0.91 to +0.93 %), pp4096 at `-ub 1024` +0.02 % (CI -0.34 to +0.38 %); decode costs the launch thread's wake on every token: tg64 at depth 0 -0.77 % (172.8 -> 171.4 t/s, 4 pairs, medians of reps 2-4, CI -0.81 to -0.72 %), at depth 16384 -0.88 % (155.9 -> 154.4, CI -1.01 to -0.75 %). | `GGML_META_CAPTURE_THREADS=0` | | `471feb45d`, `1b2c4f07a` | `e09841eff`'s launch threads woke on every replay, and decode paid the wake on every token (tg64 -0.77 % at depth 0, -0.88 % at depth 16384). The calling thread launches a replay's devices in turn again and wakes the devices' threads at the replay's start; a device's thread launches its device only if the calling thread has not taken it yet (an atomic compare-and-swap a replay), which is the case when the calling thread is blocked in an earlier device's launch, and it reads the replay's graphs only once it took its device, since a thread waking after the replay could otherwise read graphs the next evaluation freed. Same box, pack and flags as `e09841eff`'s row: 16k at `-b 9 -ub 9` finishes in 80 s, PPL 2.6523, 1754 launches a device, and with the switch it stops at the same launch; `test-backend-meta-capture` passes by default, with the switch and under `GGML_META_CAPTURE_LEGACY=1`, and split, views and sourceless pass; tg64 against the switch on the same library, 4 pairs, medians of reps 2-4: depth 0 -0.03 % (95 % CI -0.12 to +0.06 %), depth 16384 -0.02 % (CI -0.14 to +0.11 %). | `GGML_META_CAPTURE_THREADS=0` | +| `3472162a7` | `b915c5e6f`'s `utf8_replace_malformed` checked a sequence's bit pattern only, so a surrogate (ED A0 80), an overlong form (C0 80, E0 80 80, F0 80 80 80), a code point past U+10FFFF (F4 90 80 80) and a lead that leads nothing (F5-FF) passed as well-formed, and the chat parser's strict JSON dump threw on them inside a tool call's string argument: the same 500, for the bytes it did not cover. The second byte of a lead is now checked against its range (Unicode Table 3-7: E0 A0-BF, ED 80-9F, F0 90-BF, F4 80-8F, the rest 80-BF; C0, C1 and F5-FF lead nothing), a malformed sequence is one U+FFFD for its maximal subpart, and an incomplete tail is held only while every byte of it is right so far. `test-utf8-replace` (18 cases) fails 10 on the function before this change and passes on it | — | +| `466d8d4f1` | The routed-expert ring's consumers (`mmvq_moe_quantize_y`) loop to a bound rounded up to a whole warp: their shuffles name all 32 lanes, and a bound that split a warp would hang it. No graph reaches that today, since `ggml_cuda_mmvq_moe_y_bytes` admits only whole 16-byte copies, which holds a quantized row to a multiple of 128 columns, so the bound is the same for every shape the ring takes. `test-backend-ops` gains `MOE_FFN_CHAIN` in Q8_0 at 2 and 3 tokens, where the ring takes Q8_0 and the down quantizes its own vectors. R3 on an RTX 5090 at `ddaaf00a7`, the ring's q8_1 block scale doubled in `mmvq_moe_quantize_y`: those two cases fail at ERR 0.25 against a tolerance of 0.005, 6 of 17 IQ3_XXS cases fail with them, and the restored build passes all 14. The 8-token cases pass under the same mutation, since 64 vectors of 4,096 columns leave the shared plan no room for y and `ggml_cuda_mmvq_moe_keeps_y` is false there (`ce77b5fd8` says so at the cases) | — | +| `458b5c1c2`, `c7f958027` | `5c03ed96f` let the fused MoE top-k's weights and ids overlap any src at 2-4 rows, but its block barrier orders only the logits' reads, and the selection bias is read past it while a sibling warp writes. At 2-4 rows the outputs may now overlap the logits themselves, or gpt-oss's reshape of them, and every other src is checked; one row keeps upstream's exception. `c7f958027` made that test identity (same buffer, data pointer and byte count) where `458b5c1c2` wrote containment, which was wider than its own ground: the kernel reads `logits->ne[1]` rows while the rule is handed `ggml_nrows`, so the tail of a logits tensor with `ne[2] > 1` passed it unread, as a bias that was ever a view of the logits would have while being read past the barrier. Neither is reachable here, so that is the rule stated as what it rests on. It also added the first test case to reach the rule at all: every registered top-k shape had 1 row or 22, and the one with 4 is refused for having 160 experts, so `{128, 4, 1, 1}` at 8 used experts is the one that gets there. No pass or fail can show that, since refusing a fusion is always correct, so it was measured: with the rule printing the row count it runs at, `TOPK_MOE` on an RTX 5090 at `700c03ac4` reached it 26 times and only ever at 4 rows, and 0 times under `GGML_CUDA_TOPK_MOE_ALIAS_LEGACY=1`. No shipped model changes, since GLM-5.3's bias is a weight, which the check never reads. Also: a PQ2_0 MUL_MAT group takes computed, non-empty nodes only (a node `ggml_build_forward_select` leaves out was a member the node loop skips); `ggml_cuda_graph_get_key` no longer reads `nodes[0]` of an empty graph; `GGML_CUDA_TIME_LAUNCH` and `GGML_CUDA_COUNT_INFLIGHT` parse as `ggml_env_switch` does; `test-backend-meta-capture` reads `GGML_META_CAPTURE_LEGACY` as the backend does (`=0` reported FAIL) | `GGML_CUDA_TOPK_MOE_ALIAS_LEGACY=1` (one row only, as before) | +| `00e9132dd` | `kpool_dirty` was one flag for the whole cache, so a context shift of ONE sequence marked every sequence's pooled keys stale and the next ubatch of ANY sequence cleared it for all of them: the others then read keys pooled around their old positions, never re-emitted (`lies`). Now `llama_kpool_stale`, a `bitset`: `seq_add` and `seq_div` mark the sequence they moved, the graph asks `get_kpool_dirty(ubatch)` over that ubatch's `seq_id_unq` alone, and `set_input` clears only those. `llama_kpool_stale::shift_regroups` keeps the mark off a shift that moves whole pools by a whole pool, which regroups nothing. `test-kpool-input`'s `stale_after_shift` emulates the indexer cache as `set_rows` writes it: a rollback and a `pos_add` on one sequence, then that sequence's decode, then the other's, leaves 0 pools holding a stale key with per-sequence marks and 1 with one flag for the cache, the control; `run_shift_regroups` covers the four shift shapes | none (a correctness fix) | +| `47fb5e66d` | `llama_kpool_set_input`'s `emitted` bitmap, a byte per cell of the indexer cache (512 KB at 524,288 cells), was allocated before the views fast path that does not read it, so every decode step paid it for nothing. Moved into the direct path, which is the only reader. `test-kpool-input` is unchanged and still byte-identical between the two paths | none (a move) | +| `caaad776e` | A rebuild re-emits every complete pool, so `n_new_max` is `n_pools` (131,074 at 524,288 cells) instead of a decode's few, and `glm5next`'s inline compression built all of their members, gates and softmax at once: about 1.07 GB of compute buffer at `d_idx` 128 that no reserve sized (`breaks`). `llama_kpool_compress` now takes the pools `LLAMA_KPOOL_COMPRESS_CHUNK` (4,096) at a time, each chunk's write expanded into the graph before the next is built, so ggml-alloc gives the next chunk the memory of the last; a decode's pools are one chunk, the graph this compression always was. Not reached from rig's serve: the only callers of `seq_add` are context shift (`server-context.cpp:3446`) and cache reuse (`:3773`), both off by default (`common.h:600`, `:657`) and neither set by rig's head or reachable as a request field. `test-kpool-input`'s `run_compress_chunks`: 4,096 pools at `d_idx` 32 take 8,388,608 bytes at once against 1,048,576 in chunks of 512, and the chunked pooled keys are byte-identical to the single pass | `LLAMA_KPOOL_COMPRESS_CHUNK` (the chunk size; one chunk over `n_pools` is the old graph) | +| `0ac4349ea` | A sparse attention row is live at its top-k cells and at its tail, the cells of its sequence at positions `[(q + 1)/r*r, q]`, and `n_kv_max` counted `r - 1` of them padded to 32 (2080 for GLM-5-Next's 2048 + 3). A sequence holding two cells at one position has a longer tail, and `fattn.cu`'s compaction kept the row's first `n_kv_max` cells and padded with -1, so the kernel attended to part of the row with no signal (`lies`; `llama-graph.cpp` named the case and assumed the pad absorbed it, which it does not have to). `llama_kpool_tail_max` counts the longest tail of the ubatch from the cells, `r - 1` unless a sequence of it holds two cells at one position (`seq_cell_count != seq_pos_distinct`, so decode pays one comparison a sequence); `build_attn_sparse` sizes `n_kv_max = GGML_PAD(top_k + n_tail, 32)`, `can_reuse` refuses a graph built on another `n_tail`, and `set_input` returns the most tail cells it granted a row so `llm_graph_input_kpool::set_input` can assert it is within the bound. Both compaction kernels now `__trap()` on a row with more finite cells than `n_kv_max` rather than drop cells, which `ggml.h` documents. `test-kpool-input`: a sequence with 11 extra cells at each of its last three positions has a tail of 36 and 2084 live cells a row, over the 2080 an `r - 1` bound gives and under the new 2112; across the random rounds 733 builds held a tail over `r - 1` cells and every one stayed within the bound, with `set_input`'s count equal to `sel_mask`'s finite columns on the view and the cell paths. R3: with `llama_kpool_tail_max` returning `r - 1` always, the fixed case and the rounds both fail; `test-kpool-can-reuse` refuses reuse on a longer tail. On an RTX 5090 at `700c03ac4`, with mask row 0 of each of the 40 `FLASH_ATTN_EXT` cases with an `n_kv_max` given `n_kv_max + 1` finite cells: the trap aborts the first case (`mask row 0 of sequence 0 has 513 finite cells, over n_kv_max 512`, exit 134). With the trap removed, which is the old kernel, the dropped cell changes the output, and the CPU comparison sees it in 2 of the 40 (NMSE 2.1e-3 and 2.4e-3 against 5e-4, both at `n_kv_max` 512 and one query); in the other 38, every one at 2,051 or more among them, one cell of a row is under the tolerance, so the trap is the signal and the comparison is not. A row overflows only when a sequence holds two cells at one position, which the batch check admits: `llama_batch_allocr` keeps a sequence's positions in a `std::set` for its continuity check and refuses only a decrease (`llama-batch.cpp:316`, `:381`). So a caller of `llama_decode` can reach it; whether a server request does is not measured | none (a correctness fix; the trap is the contract `ggml.h` states) | +| `433d88eae` | `dsv4_hc_gram_pre` took `pre' G pre` as a sum of squares, but it is summed from G's entries, so streams that cancel give a value whose rounding can land under 0 and `rsqrtf` then returns NaN for the token's whole row (`lies`). One `fmaxf(sumsq, 0)`, in a per-token path, not a per-element one. The other two RMS sites sum squares directly and cannot go negative. `700c03ac4` adds the case that reaches it, `DSV4_HC_PRE_FUSED(...,cancel)`: 12 of 16 tokens hold streams cancelling in pairs to within 1e-7 at a scale of 100 under equal pre weights, and an f32 replay of the kernel's formula, G summed in order, rounds the form under `-n_embd * eps_norm` for 15 of 40 such tokens. On an RTX 5090 at `700c03ac4` it passes with the clamp; with the clamp removed it fails, NaN from the first cancelling token (`[MUL] NaN at index 16384`, CPU -0.004938), while the other 43 `DSV4_HC_PRE_FUSED` cases still pass, so before this case no test could fail on the clamp | none (one clamp) | +| `598d83508` | The chunked KDA kernels stage `exp` of the gate's prefix sums in fp16 (stage 3's q and k operands, and the per-channel decay), which holds only for a log-decay at or under 0: a positive gate saturates them to inf, and the only guard was `cgdr_to_fp16`'s `assert`, compiled out of Release. The shipped model's gate is `lower_bound * sigmoid` with `lower_bound` asserted negative (`glm5next.cpp:82`), so this is the public op's contract, not the model's (`lies`). `ggml_gated_delta_net` now documents `g` as a log-decay, KDA's at or under 0; a raw KDA gate whose `lower_bound` is above 0 is no longer chunk-eligible and runs the recurrent kernel; an activated gate above 0 traps in `cgdr_kda_fwdsub_intra_kernel`, inside the barrier Step 0 already had (`__syncthreads_or` over the block's own loads), so the shipped path gains no sync and no extra read. `test-backend-ops` gains four KDA cases on the chunked path: strided q/k/v through the qwen35 views with activated and with raw gates, and a chunked pass of one chunk alone (128 tokens at `K` 113, so 16 chunked and 112 on the recurrent tail) over one and two sequences. On an RTX 5090 at `700c03ac4`: with the test's activated gates made positive the trap fires (`a log-decay above 0 in chunk 4 of head 3`, exit 134). With its raw gates at `lower_bound` +0.5, the four raw KDA cases (128 to 512 tokens) that turn NaN without the guard run recurrent with it, where one passes and three fail only where the CPU reference is itself NaN; without the guard all four take the chunked kernel and turn NaN at index 8,192 or 65,536, where the reference is finite (51 to 171). Read directly, `ggml_cuda_gdn_chunked_shape_eligible` takes `lower_bound` -5 and -0.001 and refuses +0.001 and +5; without the guard it takes all four | `GGML_CUDA_KDA_CHUNKED_LEGACY=1` (KDA on the recurrent kernel at every size) | +| `60ee43721` | `GGML_CUDA_FATTN_SPARSE_LEGACY=1` gates `shall_use_sparse` alone, so it restores the dense kernel as it was before the gather, which reads only the KV steps a row of its Q tile sees (`2f0ca785a`, five days earlier, with its own switch) and not the whole cache under the mask that `0b0e97b5d`'s row claimed. The row and the comment now name both switches; the gate is unchanged, since a switch per measured optimization beats one switch that undoes two | the switches themselves | +| `cdd2dff19`, `ddaaf00a7` | `4ec3d003b` gave the MTP head's kept h row the position it is for, and a position is not an identity: a slot that ends conversation Z with a row kept for position P, then takes conversation Y out of the prompt cache (`--cache-ram`) resuming at P + 1, paired Y's first row with Z's h. The drafter implemented neither `get_state` nor `set_state`, so the stash a context checkpoint carries was always empty for it, and nothing dropped the row when a slot's cells were replaced or removed. The row and its position now travel with the state they belong to: the MTP drafter has both (as the dflash one does), a sequence with no kept row has no state and an empty state clears it, `server_prompt_data` carries that blob beside the target's and the draft's, `prompt_save` fills it, the cache's `load` restores it, and `prompt_clear` clears it. The cost was draft acceptance, never the text. llama-server builds; the prompt cache's own server tests and the slot, context-shift, speculative and basic suites pass, 36 in all. The MTP path itself has no automated check here, since it needs the private pack: the change is by construction. `ddaaf00a7` clears a slot's drafter pointer when the drafter is destroyed, which `cdd2dff19` left dangling: the slots outlive it, and sleeping destroys the drafter while `prompt_clear` would still read it | `LLAMA_MTP_PENDING_H_LEGACY=1` still pairs with the kept row whatever its position | +| `c56099b8d` | `647703131` sizes the MUL_MAT_ID tile list at `min(n_expert*ceil(n_tokens/J), (ncols_dst + min(n_expert, ncols_dst)*(J - 1))/J)`. The first term counts an expert at most once for a token, which holds for every graph built here (ids is a top-k of distinct experts) and for nothing else: ids naming an expert twice for one token gives it more than n_tokens columns, and four experts holding 5, 5, 5 and 1 of 16 columns at J = 4 need 7 tiles where the term allows 4. `mmq_moe_tiles_kernel` wrote what fit and dropped the rest, so those columns' rows of dst were never written and the result was wrong with no sign. The kernel now traps on a list past the bound, and the bound states the precondition it rests on. The shipped path is unchanged: one comparison an expert, in the list kernel, which runs once a launch on one block. R3 on an RTX 5090 at `793ed1e75`, with `ntiles_max` halved at the launch: the trap fires on `n_mats=64,n_used=8` (`expert 896 takes 0 column tiles from 64, over the launch's 41`, exit 134) and the restored build passes that case 4/4. That run is also what corrected the message, which printed `first` and `first + n - 1` and so read `tiles 64 to 63` for an expert taking none past an offset already over the bound | — | 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`, diff --git a/common/speculative.cpp b/common/speculative.cpp index 003a74948a9f..45517ccd7781 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1915,6 +1915,42 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { pending_pos[seq_id] = verify_pos[seq_id] + i_h; } + // the kept row and the position it is for, so a restored checkpoint pairs the token after it with the row the + // generation it belongs to produced. A sequence with no kept row has no state, and an empty one clears it: a slot + // whose cells are removed or replaced by another conversation's keeps a row that is no longer the one before + // anything (position alone does not tell the two apart, h_before) + bool get_state(llama_seq_id seq_id, std::vector & data) const override { + if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq || pending_pos[seq_id] < 0) { + return false; + } + + const llama_pos pos = pending_pos[seq_id]; + + data.resize(sizeof(llama_pos) + (size_t) n_embd * sizeof(float)); + std::memcpy(data.data(), &pos, sizeof(llama_pos)); + std::memcpy(data.data() + sizeof(llama_pos), pending_h[seq_id].data(), (size_t) n_embd * sizeof(float)); + return true; + } + + void set_state(llama_seq_id seq_id, const std::vector & data) override { + if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) { + return; + } + if (data.empty()) { + pending_pos[seq_id] = -1; + return; + } + if (data.size() != sizeof(llama_pos) + (size_t) n_embd * sizeof(float)) { + return; + } + + llama_pos pos = -1; + std::memcpy(&pos, data.data(), sizeof(llama_pos)); + + pending_pos[seq_id] = pos; + std::memcpy(pending_h[seq_id].data(), data.data() + sizeof(llama_pos), (size_t) n_embd * sizeof(float)); + } + void end(llama_seq_id seq_id) override { if (seq_id >= 0 && seq_id < (llama_seq_id) n_seq) { generating[seq_id] = 0; diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 2a93a3747e7d..ebf3a83536ee 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2468,7 +2468,7 @@ extern "C" { const struct ggml_tensor * a); // Use finite mask entries as a sparse K/V set. Set 0 to disable. - // n_kv_max must bound the number of finite entries in every mask row. + // n_kv_max must bound the number of finite entries in every mask row: CUDA traps on a row over it rather than drop cells. GGML_API void ggml_flash_attn_ext_set_n_kv_max( struct ggml_tensor * a, int32_t n_kv_max); @@ -2606,7 +2606,9 @@ extern "C" { // tensor shapes (S_k == S_v, H_v % H_k == 0): // q, k : [S_k, H_k, n_tokens, n_seqs] // v : [S_v, H_v, n_tokens, n_seqs] - // g : [1, H_v, n_tokens, n_seqs] (scalar gate) or [S_v, H_v, n_tokens, n_seqs] (KDA) + // g : [1, H_v, n_tokens, n_seqs] (scalar gate) or [S_v, H_v, n_tokens, n_seqs] (KDA), a log-decay: KDA's must be + // <= 0 (a decay exp(g) of at most 1), since CUDA's chunked path stages exp of its sums in fp16 and traps on + // one above 0; a raw KDA gate (ggml_gated_delta_net_set_raw_kda_gates) with lower_bound > 0 runs recurrent // beta : [1, H_v, n_tokens, n_seqs] // state : [S_v, S_v, H_v, n_seqs] -- initial recurrent state s0 // diff --git a/ggml/src/ggml-cuda/chunk_gated_delta_net.cu b/ggml/src/ggml-cuda/chunk_gated_delta_net.cu index f3a8fc70afe1..a7c085fabe7d 100644 --- a/ggml/src/ggml-cuda/chunk_gated_delta_net.cu +++ b/ggml/src/ggml-cuda/chunk_gated_delta_net.cu @@ -311,6 +311,7 @@ __launch_bounds__(128, 4) __global__ void cgdr_kda_fwdsub_intra_kernel( // Step 0: q (scaled), k, the activated log-decay and beta into SMEM; tokens past seq_len are zeros, which neither // decay the state nor update it + bool g_above_0 = false; for (int i = tid; i < CS * BK; i += 128) { const int t = i / BK, c = i % BK; float qv = 0.f, kv = 0.f, gv = 0.f; @@ -320,6 +321,8 @@ __launch_bounds__(128, 4) __global__ void cgdr_kda_fwdsub_intra_kernel( gv = g_chunk[t * sb2 * BK + c]; if constexpr (RAW) { gv = raw_lb * (1.0f / (1.0f + expf(gv * raw_a[h]))); + } else { + g_above_0 |= gv > 0.0f; } } s_q[t * sk + c] = qv; @@ -336,7 +339,15 @@ __launch_bounds__(128, 4) __global__ void cgdr_kda_fwdsub_intra_kernel( } s_beta[i] = beta_val; } - __syncthreads(); + // stage 3's fp16 operands hold exp of the gate's sums, at most 1 for a log-decay at or under 0 (ggml_gated_delta_net); + // a gate above it saturates them to inf, so the launch traps. A raw gate is under 0 by eligibility (raw_lb <= 0) + if constexpr (RAW) { + __syncthreads(); + } else if (__syncthreads_or(g_above_0) && tid == 0) { + printf("ERROR: chunked KDA: a log-decay above 0 in chunk %d of head %d; ggml_gated_delta_net's g must be <= 0 " + "(GGML_CUDA_KDA_CHUNKED_LEGACY=1 takes the recurrent kernel)\n", pid_chunk, h); + __trap(); + } // Step 1: G, each channel's prefix sum along the chunk { diff --git a/ggml/src/ggml-cuda/dsv4-hc.cu b/ggml/src/ggml-cuda/dsv4-hc.cu index f56646c3501c..322e72a79d1d 100644 --- a/ggml/src/ggml-cuda/dsv4-hc.cu +++ b/ggml/src/ggml-cuda/dsv4-hc.cu @@ -938,7 +938,9 @@ static __device__ __forceinline__ void dsv4_hc_gram_pre(const float * mix, const } sumsq += pre[a]*row; } - rms = rsqrtf(sumsq/pr.n_embd + pr.eps_norm); + // pre' G pre is a sum of squares, but rounded it can fall under 0 when the streams cancel, and rsqrtf of a negative + // makes the whole token NaN + rms = rsqrtf(fmaxf(sumsq, 0.0f)/pr.n_embd + pr.eps_norm); } // token it's pre, post and comb weights into weights_out, by one warp: the comb in registers (dsv4_hc_comb_regs), lane diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index 8d6d08154e68..d5f55817762e 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -6,8 +6,19 @@ #include "fattn.cuh" #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +// A row with more finite cells than n_kv_max breaks ggml_flash_attn_ext_set_n_kv_max's contract, and its first n_kv_max +// would attend to part of the row without a sign: the launch traps instead, and the next sync fails +static __device__ __forceinline__ void flash_attn_sparse_check_count( + const int count, const int n_kv_max, const int query, const int sequence) { + if (count > n_kv_max && threadIdx.x == 0) { + printf("ERROR: sparse flash attention: mask row %d of sequence %d has %d finite cells, over n_kv_max %d\n", + query, sequence, count, n_kv_max); + __trap(); + } +} + // A block a mask row: its finite entries' cells, ascending, into the row's n_kv_max indices, -1 past its count (a count -// over n_kv_max keeps the first n_kv_max: the bound is the graph's to keep). +// over n_kv_max writes the first n_kv_max and traps: the bound is the graph's to keep). // Upstream's scan (8e93a9773): 256 threads, 2048 columns a round of scalar loads. A decode's one row is one block, so // a round is a memory round trip: 11 us at 32K columns. GGML_CUDA_FATTN_SPARSE_SCAN_LEGACY=1, or a mask row not // 16-byte aligned. @@ -84,6 +95,7 @@ static __global__ void flash_attn_mask_to_sparse_indices_legacy( } const int count = row_count; + flash_attn_sparse_check_count(count, n_kv_max, query, sequence); for (int i = count + tid; i < n_kv_max; i += blockDim.x) { indices[i] = -1; } @@ -173,6 +185,7 @@ static __global__ void flash_attn_mask_to_sparse_indices( __syncthreads(); // the next chunk rewrites warp_sums } + flash_attn_sparse_check_count(row_count, n_kv_max, query, sequence); for (int i = row_count + tid; i < n_kv_max; i += sparse_scan_threads) { indices[i] = -1; } @@ -255,7 +268,9 @@ void ggml_cuda_flash_attn_ext_gather_k_q8_0( #endif // defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) } -// GGML_CUDA_FATTN_SPARSE_LEGACY=1: the hint is ignored, the kernel reads the whole cache under the mask. +// GGML_CUDA_FATTN_SPARSE_LEGACY=1: the hint is ignored, and the dense kernel runs as it did before the gather: over the +// KV steps a row of its Q tile sees (launch_fattn's kv_live, 2f0ca785a), the whole cache under the mask only with +// GGML_CUDA_FATTN_LIVE_TILES_LEGACY=1 as well. // The gather reads n_kv_max cells a query; the dense kernel reads the cache once for up to 64/ncols2 queries. The gather // wins where the cache is at least those queries' cells (GLM-5.3's DSA shape, 32 heads on the latent, RTX 5080 and // 5070 Ti: at 0.25-0.99x the dense time wherever K >= n_gather, at 1.13-1.27x wherever it is under; upstream's 2x margin diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu index a0d6f3189ca4..71f495052a30 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/ggml/src/ggml-cuda/gated_delta_net.cu @@ -428,7 +428,10 @@ bool ggml_cuda_gdn_chunked_shape_eligible(const ggml_tensor * dst) { // - q/k/v rows contiguous with any head/token/seq stride (the views qwen35 takes of the conv output are // read in place; q and k share strides), g/beta/state contiguous // - n_tokens >= 128; with K > 1 snapshot slots the last K-1 tokens go to the recurrent kernel - return (!kda || kda_chunked) && dst->src[6] == nullptr + // - KDA's raw gate, lower_bound * sigmoid, at or under 0: the chunked kernels stage exp of its sums in fp16, which a + // gate above 0 saturates (an activated gate above 0 traps there instead, cgdr_kda_fwdsub_intra_kernel) + const bool kda_raw_ok = !kda || ggml_get_op_params_i32(dst, 1) != 2 || ggml_get_op_params_f32(dst, 2) <= 0.0f; + return (!kda || kda_chunked) && kda_raw_ok && dst->src[6] == nullptr && dst->type == GGML_TYPE_F32 && src_q->type == GGML_TYPE_F32 && src_k->type == GGML_TYPE_F32 && src_v->type == GGML_TYPE_F32 && src_g->type == GGML_TYPE_F32 && src_beta->type == GGML_TYPE_F32 && src_state->type == GGML_TYPE_F32 diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index d75f9a090b85..4721b6ae01fd 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -142,10 +142,7 @@ void ggml_cuda_set_device(int device) { // attribution taken from a trace is unreliable at this scale -- including the one that put 66 % of a decode token's // 444.6 us device-idle gap on this call. Off (the default) this costs one test of a static flag. static bool ggml_cuda_time_launch() { - static const bool enabled = [] { - const char * env = getenv("GGML_CUDA_TIME_LAUNCH"); - return env != nullptr && atoi(env) == 1; - }(); + static const bool enabled = ggml_env_switch("GGML_CUDA_TIME_LAUNCH"); return enabled; } @@ -168,10 +165,7 @@ static void ggml_cuda_time_launch_record(int device, double us) { // GGML_CUDA_COUNT_INFLIGHT=1 prints, before each replay launch, the device's earlier replays not yet finished, so the last // line before a hang tells a host that ran far ahead (many) from one launch that cannot enqueue (0 or 1). static bool ggml_cuda_count_inflight() { - static const bool enabled = [] { - const char * env = getenv("GGML_CUDA_COUNT_INFLIGHT"); - return env != nullptr && atoi(env) == 1; - }(); + static const bool enabled = ggml_env_switch("GGML_CUDA_COUNT_INFLIGHT"); return enabled; } @@ -2073,7 +2067,10 @@ static bool ggml_cuda_pq2_mma_group_enabled() { // A MUL_MAT that ggml_cuda_mul_mat sends to ggml_cuda_mul_mat_vec_q and that sends it, unfused, to the tensor-core kernel static bool ggml_cuda_pq2_mma_group_member(ggml_backend_cuda_context & ctx, const ggml_tensor * mm) { - return mm->op == GGML_OP_MUL_MAT && ggml_get_op_params_i32(mm, 1) != GGML_HINT_SRC0_IS_HADAMARD && + // a node the graph does not compute (GGML_TENSOR_FLAG_COMPUTE unset, ggml_build_forward_select) is no member: the node + // loop skips it, and a group would compute it + return mm->op == GGML_OP_MUL_MAT && (mm->flags & GGML_TENSOR_FLAG_COMPUTE) && !ggml_is_empty(mm) && + ggml_get_op_params_i32(mm, 1) != GGML_HINT_SRC0_IS_HADAMARD && mm->src[1]->nb[0] == sizeof(float) && mm->nb[0] == sizeof(float) && ggml_cuda_should_use_mmvq(GGML_TYPE_PQ2_0, ggml_cuda_info().devices[ctx.device].cc, mm->src[1]->ne[1]) && ggml_cuda_pq2_mma_fuses(ctx, mm, nullptr, nullptr); @@ -3600,7 +3597,7 @@ static ggml_cuda_graph_key ggml_cuda_graph_get_key(ggml_cgraph * cgraph) { static const bool legacy = ggml_env_switch("GGML_CUDA_GRAPH_KEY_LEGACY"); ggml_cuda_graph_key key; - key.first_node = cgraph->nodes[0]; + key.first_node = cgraph->n_nodes > 0 ? cgraph->nodes[0] : nullptr; // an empty graph's slot 0 is no node of it if (legacy) { return key; } @@ -4590,13 +4587,26 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, }; bool is_ok = true; - // exception for topk-moe, which reads all its rows before it writes when they fit one block - // (ggml_cuda_topk_moe_reads_before_writes). GGML_CUDA_TOPK_MOE_ALIAS_LEGACY=1: at one row only - static const bool topk_moe_alias_legacy = ggml_env_switch("GGML_CUDA_TOPK_MOE_ALIAS_LEGACY"); - const int64_t topk_moe_rows = ggml_nrows(cgraph->nodes[node_idx]); - if (is_topk_moe && (topk_moe_alias_legacy ? topk_moe_rows == 1 : ggml_cuda_topk_moe_reads_before_writes(topk_moe_rows))) { + // exception for topk-moe at one row: its one warp reads every src before it writes + const int64_t topk_moe_rows = ggml_nrows(cgraph->nodes[node_idx]); + if (is_topk_moe && topk_moe_rows == 1) { return true; } + // at the rows that fit one block, its warps meet at a barrier once the logits (the first node's src) are read, so the + // outputs may overlap the logits themselves, or gpt-oss's reshape of them; the bias is read past the barrier, so it + // is checked (ggml_cuda_topk_moe_reads_before_writes). GGML_CUDA_TOPK_MOE_ALIAS_LEGACY=1: at one row only + static const bool topk_moe_alias_legacy = ggml_env_switch("GGML_CUDA_TOPK_MOE_ALIAS_LEGACY"); + const ggml_tensor * read_first = is_topk_moe && !topk_moe_alias_legacy && + ggml_cuda_topk_moe_reads_before_writes(topk_moe_rows) ? cgraph->nodes[node_idx]->src[0] : nullptr; + // the logits and a reshape of them, which share both, and nothing else: a src merely CONTAINED in the logits is a + // strict sub-view, and the two that can be one are both read past the barrier, not before it. The bias is one (its + // every instance here is a model weight, so GGML_OP_NONE above already skips it, and a weight shares no buffer with + // the logits); the other is the tail of a logits tensor with ne[2] > 1, which ggml_nrows counts and the launch's + // logits->ne[1] does not read. + auto reads_first_whole = [&](const ggml_tensor * src) { + return read_first != nullptr && src->buffer == read_first->buffer && src->data == read_first->data && + ggml_nbytes(src) == ggml_nbytes(read_first); + }; for (int i = 0; i < out_count; ++i) { const ggml_tensor * dst = cgraph->nodes[out_nodes[i]]; @@ -4609,7 +4619,7 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) { const ggml_tensor * src = cgraph->nodes[j]->src[src_idx]; - if (!src || src->op == GGML_OP_NONE || src == staged_src) { + if (!src || src->op == GGML_OP_NONE || src == staged_src || reads_first_whole(src)) { continue; } diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu index 208dedb4a6a2..42b0d9022ff0 100644 --- a/ggml/src/ggml-cuda/mmq.cu +++ b/ggml/src/ggml-cuda/mmq.cu @@ -320,7 +320,9 @@ int64_t ggml_cuda_mmq_moe_ncols() { } // One block: each pass scans the column tiles of blockDim.x experts, ceil(columns/J) each, on top of the earlier -// passes' total, and every expert writes its own entries at its exclusive prefix. +// passes' total, and every expert writes its own entries at its exclusive prefix. A list longer than the launch's +// column tiles (ntiles_max, whose first term counts an expert at most once a token) would leave an expert's last +// columns unlaunched and their rows of dst unwritten, with no sign: the launch traps instead, and the next sync fails. static __global__ void mmq_moe_tiles_kernel(const int32_t * __restrict__ expert_bounds, const int n_expert, const int J, const int ntiles_max, int32_t * __restrict__ moe_tiles) { __shared__ int warp_incl[WARP_SIZE]; @@ -365,7 +367,14 @@ static __global__ void mmq_moe_tiles_kernel(const int32_t * __restrict__ expert_ __syncthreads(); const int first = done + (warp > 0 ? warp_incl[warp - 1] : 0) + incl - n; - for (int k = 0; k < n && first + k < ntiles_max; ++k) { + if (first + n > ntiles_max) { + // first and n, not a range: an expert that takes none still trips this once the running offset is past the + // bound, and "tiles 64 to 63" names no tile + printf("ERROR: MUL_MAT_ID: expert %d takes %d column tiles from %d, over the launch's %d (an expert named " + "twice for one token?)\n", e, n, first, ntiles_max); + __trap(); + } + for (int k = 0; k < n; ++k) { moe_tiles[first + k] = (k << 16) | e; } __syncthreads(); diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh index bfae67ed9b97..e0e2d01889d7 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh @@ -1607,7 +1607,9 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a // MUL_MAT_ID: the grid's column tiles are the experts' non-empty ones back to back, one channel. An expert of c // columns fills ceil(c/J) <= ntx of them and an empty one none, so with n columns over n_expert experts there are - // at most (n + min(n_expert, n)*(J - 1))/J. + // at most (n + min(n_expert, n)*(J - 1))/J. The first term holds while ids names an expert at most once for a + // token, which every graph here builds (a top-k of distinct experts); mmq_moe_tiles_kernel traps on a list past + // this bound rather than drop an expert's last tiles. int64_t nchannels_grid = args.nchannels_y; ggml_cuda_pool_alloc moe_tiles(ctx.pool(id)); if (args.ids_dst && !ggml_cuda_mmq_moe_tiles_legacy()) { diff --git a/ggml/src/ggml-cuda/mmvq-moe.cu b/ggml/src/ggml-cuda/mmvq-moe.cu index 05ddc2402ba8..dac65f2b184b 100644 --- a/ggml/src/ggml-cuda/mmvq-moe.cu +++ b/ggml/src/ggml-cuda/mmvq-moe.cu @@ -201,7 +201,12 @@ static __device__ __forceinline__ void mmvq_moe_quantize_y(const mmvq_moe_dev_ar constexpr int NT = MMVQ_MOE_NW*32; const int nf4 = ncols / 4; const int total = a.ntokens*a.n_used*nf4; // a whole number of q8_1 blocks: 8 lanes are all in or all out - for (int j0 = threadIdx.x; j0 < total; j0 += MMVQ_MOE_QY_CHUNK*NT) { + // the bound is rounded up to a whole warp, so a warp's lanes are all in or all out: the shuffles below name all 32 + // lanes, and a warp the bound split would wait on lanes that never reach them. Today total is a multiple of 32 + // only because ggml_cuda_mmvq_moe_y_bytes admits a copy of a whole number of 16-byte units, which holds the rows to + // 128 columns; the j < total guards cover the load and the store + const int bound = (total + 31)/32*32; + for (int j0 = threadIdx.x; j0 < bound; j0 += MMVQ_MOE_QY_CHUNK*NT) { float4 v[MMVQ_MOE_QY_CHUNK]; #pragma unroll for (int c = 0; c < MMVQ_MOE_QY_CHUNK; ++c) { diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 460f99f15574..e0edd733ef52 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -602,8 +602,11 @@ llm_graph_input_kpool_dims llm_graph_input_kpool::current_dims( dims.n_tps = ubatch.n_tokens/dims.n_stream; dims.n_ps = ubatch.n_seqs_unq/dims.n_stream; dims.n_pools = llama_kpool_n_pools(dims.n_kv, kpool, dims.n_ps); - dims.rebuild = mctx_attn->get_kv()->get_kpool_dirty(); + dims.rebuild = mctx_attn->get_kv()->get_kpool_dirty(ubatch); dims.n_new_max = dims.rebuild ? dims.n_pools : dims.n_tps/kpool + dims.n_ps; + dims.n_tail = llama_kpool_tail_max([kv = mctx_attn->get_kv()](llama_seq_id s) -> const llama_kv_cells & { + return kv->get_cells(s); + }, ubatch, kpool); // same gate as the builder (glm5next_n_select == indexer_top_k + kpool - 1) dims.scoring = cparams.n_ctx > hparams.indexer_top_k + kpool - 1; dims.n_dump = dims.scoring ? llama_kpool_select_k(dims.n_pools, hparams.indexer_top_k, kpool) : 0; @@ -657,6 +660,9 @@ bool llm_graph_input_kpool::shapes_match(const llm_graph_input_kpool_dims & dims res &= inp.rebuild == dims.rebuild; res &= (int64_t) inp.n_new_max == dims.n_new_max; + // the sparse attention's n_kv_max is sized on it + res &= inp.n_tail == dims.n_tail; + res &= inp.new_pool_cells->ne[0] == (int64_t) inp.kpool*dims.n_new_max; res &= inp.new_pool_cells->ne[1] == dims.n_stream; @@ -4051,12 +4057,16 @@ llm_graph_input_kpool * llm_graph_context::build_inp_kpool( // n_new_max is an exact bound (a contiguous run of L tokens closes at most L/kpool + 1 // pools), FIXED for the decode phase so the graph shape does not track pools-closed-this- - // step; after a position mutation every cached key is stale, so it must re-emit all - const bool rebuild = mctx_attn->get_kv()->get_kpool_dirty(); + // step; after a position mutation of a sequence of this ubatch every cached key of its sequences is stale, so it + // must re-emit all of theirs + const bool rebuild = mctx_attn->get_kv()->get_kpool_dirty(ubatch); const int64_t n_new_max = rebuild ? n_pools : n_tps/kpool + n_ps; inp->n_new_max = (uint32_t) n_new_max; inp->rebuild = rebuild; + inp->n_tail = llama_kpool_tail_max([kv = mctx_attn->get_kv()](llama_seq_id s) -> const llama_kv_cells & { + return kv->get_cells(s); + }, ubatch, kpool); inp->pool_reps = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_pools, n_stream); ggml_set_input(inp->pool_reps); @@ -4088,6 +4098,7 @@ ggml_tensor * llm_graph_context::build_attn_sparse( ggml_tensor * top_k, ggml_tensor * sel_mask, ggml_tensor * cand_mask, + int64_t n_tail, float kq_scale, int il) const { ggml_build_forward_expand(gf, q_cur); @@ -4154,9 +4165,10 @@ ggml_tensor * llm_graph_context::build_attn_sparse( // a row's live cells: its top_k pools' cells (top_k->ne[0] = r*select_k) and its tail, positions [(q + 1)/r*r, q] // (kpool_mask_row), at most r - 1 cells unless a sequence holds two cells at one position (an rm + add before the first - // cell is gone), whose tail then has more. The kernel may read those alone; the bound rounds up to 32 cells, a CUDA - // kernel step, so that case keeps every cell at no cost (2051 and 2080 cells are 65 steps either way) - const int64_t n_kv_max = GGML_PAD(top_k->ne[0] + hparams.indexer_kpool - 1, 32); + // cell is gone), whose tail then has more. n_tail counts the longest tail of this ubatch in the cells, and set_input + // asserts no row is granted more, so the bound holds every row's live cells; the compaction traps on one it does not + GGML_ASSERT(n_tail >= (int64_t) hparams.indexer_kpool - 1); + const int64_t n_kv_max = llama_kpool_n_kv_max(top_k->ne[0], n_tail); ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, mask_top_k, sinks, v_mla, kq_scale, il, /*mask_is_prefix =*/ false, n_kv_max); cb(cur, "kqv_out", il); diff --git a/src/llama-graph.h b/src/llama-graph.h index b138c71d846b..9fd81e167a63 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1497,6 +1497,7 @@ struct llm_graph_context { ggml_tensor * top_k, // I32 [n_select, n_tokens/n_stream, n_stream], no cell twice in a row ggml_tensor * sel_mask, // F16/F32 [n_kv + dump columns, n_batch, 1, n_stream] ggml_tensor * cand_mask, // F16/F32 [n_kv, n_batch, 1, n_stream] + int64_t n_tail, // the most cells sel_mask grants a row (llm_graph_input_kpool::n_tail) float kq_scale, int il) const; diff --git a/src/llama-kv-cache-kpool.cpp b/src/llama-kv-cache-kpool.cpp index a41b2764cd15..3ad9aed76a29 100644 --- a/src/llama-kv-cache-kpool.cpp +++ b/src/llama-kv-cache-kpool.cpp @@ -27,6 +27,39 @@ uint32_t llama_kpool_select_k(uint32_t n_pools, uint32_t indexer_top_k, uint32_t return std::min(n_pools, indexer_top_k/kpool); } +int64_t llama_kpool_tail_max(const llama_kpool_cells_fn & cells_of, const llama_ubatch & ubatch, uint32_t kpool) { + const llama_pos r = (llama_pos) kpool; + + int64_t n_tail = r - 1; + + for (uint32_t k = 0; k < ubatch.n_seqs_unq; ++k) { + const llama_seq_id s = ubatch.seq_id_unq[k]; + const llama_kv_cells & cells = cells_of(s); + + // one cell a position: a tail spans at most r - 1 positions + if (cells.seq_cell_count(s) == cells.seq_pos_distinct(s)) { + continue; + } + + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + // a reserve ubatch (ubatch_reserve) names no sequence for its tokens and places none + if (ubatch.seq_id[i] == nullptr || ubatch.seq_id[i][0] != s) { + continue; + } + + const llama_pos q = ubatch.pos[i]; + + n_tail = std::max(n_tail, cells.seq_pos_count(s, (q + 1)/r*r, q + 1, UINT32_MAX)); + } + } + + return n_tail; +} + +int64_t llama_kpool_n_kv_max(int64_t n_top_k, int64_t n_tail) { + return GGML_PAD(n_top_k + n_tail, 32); +} + // sel_mask and cand_mask hold only 0.0f and -INFINITY, so f16 is exact here template struct kpool_mask_of; @@ -51,9 +84,10 @@ static void kpool_mask_fill(char * dst, int64_t n, bool f16) { } } -// pool_of[j] - b_base is the pool of cell j among the run, or negative when its pool is not complete +// pool_of[j] - b_base is the pool of cell j among the run, or negative when its pool is not complete. Returns the cells the +// tail grants template -static void kpool_mask_row( +static int64_t kpool_mask_row( T * cur_sel, T * cur_cand, const llama_pos * pos_at, @@ -66,6 +100,8 @@ static void kpool_mask_row( const T v_sel = kpool_mask_of::from(0.0f); const T v_mask = kpool_mask_of::from(-INFINITY); + int64_t n_tail = 0; + for (int64_t j = 0; j < n_kv; ++j) { const bool vis = (uint32_t) pos_at [j] <= (uint32_t) q; const bool pooled = (uint32_t) (pool_of[j] - b_base) < (uint32_t) bo_vis; @@ -73,7 +109,11 @@ static void kpool_mask_row( cur_sel [j] = vis && tail ? v_sel : v_mask; cur_cand[j] = vis && (pooled || tail) ? v_sel : v_mask; + + n_tail += vis && tail; } + + return n_tail; } // the cell an unused write slot names: empty, or not the last of its block, and not a cell the real slots write, since a pooled key @@ -131,8 +171,8 @@ struct kpool_stream_out { // one sequence alone in its stream, from its view: the bytes the loop over the cells of llama_kpool_set_input writes. // That loop's caller has zeroed pool_cells, pool_reps, new_cells and the padded mask rows, set pool_bias to -INFINITY, and written -// the dump pools and sel_mask's dump columns -static void kpool_stream_from_view( +// the dump pools and sel_mask's dump columns. Returns the most cells a row's tail was granted +static int64_t kpool_stream_from_view( const kpool_stream_out & o, const llama_kpool_views::view & v, const llama_kv_cells & cells, @@ -223,7 +263,8 @@ static void kpool_stream_from_view( } } - int64_t n_done = 0; + int64_t n_done = 0; + int64_t n_tail_max = 0; for (int64_t ii = 0; ii < o.n_tps; ++ii) { const int64_t i = s*o.n_tps + ii; @@ -246,13 +287,13 @@ static void kpool_stream_from_view( char * cur_sel = o.sel_mask + ii*o.n_sel*o.mask_ts; char * cur_cand = o.cand_mask + ii*o.n_kv *o.mask_ts; - if (o.mask_f16) { + const int64_t n_tail = o.mask_f16 ? kpool_mask_row((ggml_fp16_t *) cur_sel, (ggml_fp16_t *) cur_cand, - v.pos_at.data(), v.pblk.data(), (int32_t) b_base, o.n_kv, q, tail_start, bo_vis); - } else { + v.pos_at.data(), v.pblk.data(), (int32_t) b_base, o.n_kv, q, tail_start, bo_vis) : kpool_mask_row((float *) cur_sel, (float *) cur_cand, v.pos_at.data(), v.pblk.data(), (int32_t) b_base, o.n_kv, q, tail_start, bo_vis); - } + + n_tail_max = std::max(n_tail_max, n_tail); float * q_pool_bias = o.pool_bias + ii*o.n_pools; @@ -279,9 +320,11 @@ static void kpool_stream_from_view( } } } + + return n_tail_max; } -void llama_kpool_set_input( +int64_t llama_kpool_set_input( const llama_kpool_cells_fn & cells_of, llama_kpool_views * views, ggml_tensor * cell_pool, @@ -413,6 +456,8 @@ void llama_kpool_set_input( std::vector pos_at; std::vector run_off(n_ps); + + int64_t n_tail_max = 0; std::vector run_len(n_ps); auto seq_of = [&](int64_t s, int64_t ps) { @@ -446,9 +491,6 @@ void llama_kpool_set_input( // the fallback of an unused slot when no spare cell is left: recomputing a complete pool is idempotent const int32_t * any_rep_src = nullptr; - // the cells named as a rep by the new slots - std::vector emitted(kcache ? kv_size : 0, 0); - if (kcache) { // a pool with no rep gathers row 0; such a pool is -INFINITY in pool_bias, so discarded std::fill(cur_pool_reps, cur_pool_reps + n_pools, 0); @@ -487,7 +529,7 @@ void llama_kpool_set_input( o.mask_f16 = mask_f16; o.mask_ts = mask_ts; - kpool_stream_from_view(o, *v, cells_of(seq), seq, ubatch, s, rebuild); + n_tail_max = std::max(n_tail_max, kpool_stream_from_view(o, *v, cells_of(seq), seq, ubatch, s, rebuild)); continue; } @@ -495,6 +537,9 @@ void llama_kpool_set_input( views->get_stats().n_direct++; } + // the cells named as a rep by the new slots: a byte per cell of the cache, so only on this path, which reads it + std::vector emitted(kcache ? kv_size : 0, 0); + // [TAG_KPOOL_PACK] one packed run per sequence, NOT one full-width table: the indexer // scores every slot against every query, which would multiply the score tensor by n_seq_max { @@ -687,13 +732,13 @@ void llama_kpool_set_input( char * cur_sel = cur_sel_mask + ii*n_sel*mask_ts; char * cur_cand = cur_cand_mask + ii*n_kv *mask_ts; - if (mask_f16) { + const int64_t n_tail = mask_f16 ? kpool_mask_row((ggml_fp16_t *) cur_sel, (ggml_fp16_t *) cur_cand, - pos_at.data(), pool_of.data(), 0, n_kv, q, tail_start, bo_vis); - } else { + pos_at.data(), pool_of.data(), 0, n_kv, q, tail_start, bo_vis) : kpool_mask_row((float *) cur_sel, (float *) cur_cand, pos_at.data(), pool_of.data(), 0, n_kv, q, tail_start, bo_vis); - } + + n_tail_max = std::max(n_tail_max, n_tail); if (cur_bias) { for (int64_t j = 0; j < n_kv; ++j) { @@ -734,6 +779,80 @@ void llama_kpool_set_input( } } } + + return n_tail_max; +} + +void llama_kpool_compress( + ggml_context * ctx, + ggml_cgraph * gf, + ggml_tensor * kg_rows, + ggml_tensor * new_pool_cells, + ggml_tensor * new_pool_reps, + ggml_tensor * ape, + int64_t r, + int64_t chunk, + const llama_kpool_write_fn & write, + const llama_kpool_name_fn & name) { + GGML_ASSERT(r > 0 && chunk > 0); + GGML_ASSERT(kg_rows->ne[0] % 2 == 0); + GGML_ASSERT(new_pool_cells->ne[0] % r == 0); + + const int64_t d_idx = kg_rows->ne[0]/2; + const int64_t n_stream = new_pool_cells->ne[1]; + const int64_t n_new_max = new_pool_cells->ne[0]/r; + + GGML_ASSERT(kg_rows->ne[2] == n_stream); + GGML_ASSERT(new_pool_reps->ne[0] == n_new_max*n_stream); + GGML_ASSERT(ape->ne[0] == r && ape->ne[1] == d_idx); + + for (int64_t p0 = 0; p0 < n_new_max; p0 += chunk) { + const int64_t n = std::min(chunk, n_new_max - p0); + + // the chunk's members and rep rows, stream by stream; the whole tensors when it is every pool + ggml_tensor * cells = new_pool_cells; + ggml_tensor * reps = new_pool_reps; + + if (n != n_new_max) { + cells = ggml_view_2d(ctx, new_pool_cells, r*n, n_stream, new_pool_cells->nb[1], + ggml_row_size(new_pool_cells->type, r*p0)); + reps = ggml_view_2d(ctx, new_pool_reps, n, n_stream, ggml_row_size(new_pool_reps->type, n_new_max), + ggml_row_size(new_pool_reps->type, p0)); + + if (n_stream > 1) { + cells = ggml_cont(ctx, cells); + reps = ggml_cont(ctx, reps); + } + reps = ggml_reshape_1d(ctx, reps, n*n_stream); + } + + ggml_tensor * members = ggml_get_rows(ctx, kg_rows, cells); + name(members, "indexer_pool_members"); + + const size_t nb_mem = members->nb[1]; + + ggml_tensor * mem_k = ggml_view_4d(ctx, members, d_idx, r, n, n_stream, + nb_mem, nb_mem*r, members->nb[2], 0); + ggml_tensor * mem_g = ggml_view_4d(ctx, members, d_idx, r, n, n_stream, + nb_mem, nb_mem*r, members->nb[2], d_idx*members->nb[0]); + + // r-way softmaxes over the SLOT axis, so it must be dim 0; ape is added PRE-softmax + ggml_tensor * keys_t = ggml_cont(ctx, ggml_permute(ctx, mem_k, 1, 0, 2, 3)); + ggml_tensor * gate_t = ggml_cont(ctx, ggml_permute(ctx, mem_g, 1, 0, 2, 3)); + + gate_t = ggml_add(ctx, gate_t, ggml_reshape_4d(ctx, ape, r, d_idx, 1, 1)); + + // soft_max maps ne2/ne3 to gridDim.y/z (65535 cap); fold into ne1 + ggml_tensor * probs = ggml_soft_max(ctx, ggml_reshape_2d(ctx, gate_t, r, d_idx*n*n_stream)); + probs = ggml_reshape_4d(ctx, probs, r, d_idx, n, n_stream); + name(probs, "indexer_pool_probs"); + + ggml_tensor * pool_new = ggml_sum_rows(ctx, ggml_mul(ctx, keys_t, probs)); + pool_new = ggml_reshape_2d(ctx, pool_new, d_idx, n*n_stream); + name(pool_new, "indexer_pool_new"); + + ggml_build_forward_expand(gf, write(pool_new, reps)); + } } void llama_kpool_views::view::rebuild(const llama_kv_cells & c, llama_seq_id seq) { @@ -956,7 +1075,7 @@ const llama_kpool_views::view * llama_kpool_views::serve( return &v; } -void llama_kv_cache_set_input_kpool( +int64_t llama_kv_cache_set_input_kpool( const llama_kv_cache * kv, ggml_tensor * cell_pool, ggml_tensor * pool_cells, @@ -983,12 +1102,12 @@ void llama_kv_cache_set_input_kpool( llama_kpool_views * views = legacy ? nullptr : &kv->get_kpool_views(); - llama_kpool_set_input(cells_of, views, + const int64_t n_tail = llama_kpool_set_input(cells_of, views, cell_pool, pool_cells, bias, pool_bias, sel_mask, cand_mask, pool_reps, new_pool_cells, new_pool_reps, strm_of, kv_size, rebuild, ubatch, kpool); if (!check || !views) { - return; + return n_tail; } // the same maps from the cells; the views must have written the same bytes @@ -1033,6 +1152,8 @@ void llama_kv_cache_set_input_kpool( LLAMA_LOG_WARN("%s: k-pool input checked %" PRIu64 " calls: %" PRIu64 " streams from views, %" PRIu64 " views rebuilt, %" PRIu64 " from the cells\n", __func__, n_checked, st.n_served, st.n_rebuilt, st.n_direct); } + + return n_tail; } void llm_graph_input_kpool::set_input(const llama_ubatch * ubatch) { @@ -1057,7 +1178,7 @@ void llm_graph_input_kpool::set_input(const llama_ubatch * ubatch) { GGML_ASSERT(ggml_backend_buffer_is_host(pool_dump->buffer) && pool_dump->type == GGML_TYPE_F32 && ggml_is_contiguous(pool_dump)); std::fill_n((float *) pool_dump->data, ggml_nelements(pool_dump), -FLT_MAX); - llama_kv_cache_set_input_kpool( + const int64_t n_tail_set = llama_kv_cache_set_input_kpool( mctx_attn->get_kv(), /* cell_pool */ nullptr, pool_cells, /* bias */ nullptr, pool_bias, sel_mask, cand_mask, @@ -1067,8 +1188,12 @@ void llm_graph_input_kpool::set_input(const llama_ubatch * ubatch) { rebuild, ubatch, kpool); - // cleared here, not in build_inp_kpool: a graph built but not evaluated must not clear it + // the sparse attention's n_kv_max counts n_tail tail cells a row: a row granted more would lose cells to its compaction + GGML_ASSERT(n_tail_set <= n_tail && "k-pool: a row's tail holds more cells than the graph's sparse bound"); + + // cleared here, not in build_inp_kpool: a graph built but not evaluated must not clear it. Only this ubatch's + // sequences were re-emitted, so only theirs if (rebuild) { - mctx_attn->get_kv()->clear_kpool_dirty(); + mctx_attn->get_kv()->clear_kpool_dirty(*ubatch); } } diff --git a/src/llama-kv-cache-kpool.h b/src/llama-kv-cache-kpool.h index 0e563d504d69..52e7e2ad48f4 100644 --- a/src/llama-kv-cache-kpool.h +++ b/src/llama-kv-cache-kpool.h @@ -92,9 +92,45 @@ class llama_kpool_views { using llama_kpool_cells_fn = std::function; +// the pools a compression makes at once: a rebuild re-emits every pool (n_new_max = n_pools, 131,074 at 524,288 cells), and +// their members and softmax at once took about 1.07 GB of compute buffer at d_idx 128 that no reserve sized; a chunk takes +// about 34 MB (test-kpool-input: 8,388,608 bytes for 4,096 pools at d_idx 32), and a decode's few pools are one chunk +constexpr int64_t LLAMA_KPOOL_COMPRESS_CHUNK = 4096; + +// the write of a chunk's pooled keys (src: [d_idx, rows]) at its rep rows (I64 [rows]); and the name of a tensor +using llama_kpool_write_fn = std::function; +using llama_kpool_name_fn = std::function; + +// The pooled keys of the new pools: each pool's r member keys mixed by a softmax over their gates (ape added before it, over +// the slot axis), written at the pool's rep row. kg_rows: [2*d_idx, n_kv, n_stream], each cell's key then gate; new_pool_cells +// (I32 [r*n_new_max, n_stream]) and new_pool_reps (I64 [n_new_max*n_stream]) as llm_graph_input_kpool holds them; ape: [r, d_idx]. +// The pools go `chunk` at a time, each chunk's write expanded into gf before the next chunk is made, so the allocator gives the +// next chunk the memory of the last. One chunk is the graph this compression always was +void llama_kpool_compress( + ggml_context * ctx, + ggml_cgraph * gf, + ggml_tensor * kg_rows, + ggml_tensor * new_pool_cells, + ggml_tensor * new_pool_reps, + ggml_tensor * ape, + int64_t r, + int64_t chunk, + const llama_kpool_write_fn & write, + const llama_kpool_name_fn & name); + +// A row of the sparse attention's mask is live at its top-k pools' cells and at its tail, the cells of its sequence at +// positions [(q + 1)/r*r, q] (sel_mask's finite columns): at most r - 1 unless the sequence holds two cells at one position. +// The most cells a tail of this ubatch holds, at least r - 1, from the cells as the ubatch left them +int64_t llama_kpool_tail_max(const llama_kpool_cells_fn & cells_of, const llama_ubatch & ubatch, uint32_t kpool); + +// the live cells a mask row may hold (ggml_flash_attn_ext_set_n_kv_max): n_top_k top-k cells and an n_tail tail, padded to 32 +// cells, a step of the CUDA kernel +int64_t llama_kpool_n_kv_max(int64_t n_top_k, int64_t n_tail); + // the maps of llama_kv_cache_set_input_kpool from a cells accessor. views == nullptr builds them from the cells -// every call; else a sequence alone in its stream is served from its view, byte for byte the same maps -void llama_kpool_set_input( +// every call; else a sequence alone in its stream is served from its view, byte for byte the same maps. +// Returns the most cells a row's tail was granted (sel_mask's finite columns in a row) +int64_t llama_kpool_set_input( const llama_kpool_cells_fn & cells_of, llama_kpool_views * views, ggml_tensor * cell_pool, @@ -120,8 +156,9 @@ void llama_kpool_set_input( // sel_mask n_kv columns, then kpool*n_dump dump columns of -inf (n_kv is cand_mask's) // cand_mask bounds top-k spills a partial seq_rm would let escape // pool_reps / new_pool_cells / new_pool_reps are nullptr when the cache is off, and an entry -// is emitted only for filled == kpool: cell 0 is real, so writing its 0 slot would clobber -void llama_kv_cache_set_input_kpool( +// is emitted only for filled == kpool: cell 0 is real, so writing its 0 slot would clobber. +// Returns llama_kpool_set_input's +int64_t llama_kv_cache_set_input_kpool( const llama_kv_cache * kv, ggml_tensor * cell_pool, ggml_tensor * pool_cells, @@ -151,6 +188,7 @@ struct llm_graph_input_kpool_dims { int64_t n_pools = 0; int64_t n_dump = 0; int64_t n_new_max = 0; + int64_t n_tail = 0; bool rebuild = false; bool scoring = false; }; @@ -204,6 +242,10 @@ class llm_graph_input_kpool : public llm_graph_input_i { // set at build time: re-emit every pool after a position mutation bool rebuild = false; + // set at build time: the most cells a row's tail holds (llama_kpool_tail_max), which sizes the sparse attention's + // n_kv_max; set_input asserts no row is granted more + int64_t n_tail = 0; + // exact, since pool_bias only holds 0.0f or -INFINITY. nullptr if the fused path is off ggml_tensor * pool_bias_f16 = nullptr; // F16 [n_pools, n_tps, 1, n_stream] diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 7897ea6fadda..6f28cf71bdcf 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1285,16 +1285,54 @@ bool llama_kv_cache::get_has_shift() const { return result; } -void llama_kv_cache::set_kpool_dirty() { - kpool_dirty = true; +bool llama_kpool_stale::shift_regroups(uint32_t kpool, llama_pos p0, llama_pos p1, llama_pos shift) { + GGML_ASSERT(kpool > 0); + + const llama_pos r = (llama_pos) kpool; + + const bool whole_pools = shift % r == 0 && + (p0 <= 0 || p0 % r == 0) && + (p1 < 0 || p1 % r == 0); + + return !whole_pools; +} + +void llama_kpool_stale::mark(llama_seq_id seq_id) { + if (seq_id < 0) { + seqs.set(); + return; + } + + GGML_ASSERT(seq_id < LLAMA_MAX_SEQ); + seqs.set(seq_id); +} + +bool llama_kpool_stale::any(const llama_ubatch & ubatch) const { + for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) { + if (seqs.test(ubatch.seq_id_unq[s])) { + return true; + } + } + + return false; +} + +void llama_kpool_stale::clear(const llama_ubatch & ubatch) { + for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) { + seqs.reset(ubatch.seq_id_unq[s]); + } +} + +void llama_kv_cache::set_kpool_dirty(llama_seq_id seq_id) { + kpool_stale.mark(seq_id); } -bool llama_kv_cache::get_kpool_dirty() const { - return kpool_dirty; +bool llama_kv_cache::get_kpool_dirty(const llama_ubatch & ubatch) const { + return kpool_stale.any(ubatch); } -void llama_kv_cache::clear_kpool_dirty() const { - kpool_dirty = false; +void llama_kv_cache::clear_kpool_dirty(const llama_ubatch & ubatch) const { + kpool_stale.clear(ubatch); } llama_kpool_views & llama_kv_cache::get_kpool_views() const { diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index c89858d3ac01..cf29ccb9e1e9 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -16,6 +16,25 @@ struct llama_context; class llama_kpool_views; +// GLM-5-Next pooled-key cache: the sequences whose cached pooled keys a position mutation left stale. Pools group cells by +// absolute position, so seq_add or seq_div regroups them while every cached key still looks complete. A ubatch holding a +// marked sequence re-emits every pool of its sequences and clears only their marks; a marked sequence the ubatch does not +// hold stays marked until a ubatch of its own (one flag for the cache, cleared by another sequence's ubatch, left it stale). +// Sticky by design: a mark that clears too early is silently wrong +struct llama_kpool_stale { + // whether seq_add's shift of [p0, p1) by `shift` regroups pools of `kpool` cells: a cached key survives only if WHOLE pools + // move, a shift by a multiple of kpool over a range that starts and ends on a pool boundary (p0 < 0: from the start, + // p1 < 0: to the end, neither can be straddled) + static bool shift_regroups(uint32_t kpool, llama_pos p0, llama_pos p1, llama_pos shift); + + void mark(llama_seq_id seq_id); // seq_id < 0: every sequence + bool any (const llama_ubatch & ubatch) const; + void clear(const llama_ubatch & ubatch); + +private: + std::bitset seqs; +}; + // // llama_kv_cache // @@ -161,11 +180,11 @@ class llama_kv_cache : public llama_memory_i { bool get_has_shift() const; - // GLM-5-Next pooled-key cache: seq_add/seq_div regroup pools while every cached key still - // looks complete. sticky by design, a flag that clears too early is silently wrong - void set_kpool_dirty(); - bool get_kpool_dirty() const; - void clear_kpool_dirty() const; + // GLM-5-Next pooled-key cache, per sequence (llama_kpool_stale): seq_add/seq_div mark a sequence, a ubatch holding a + // marked one rebuilds, and its set_input clears the ubatch's sequences + void set_kpool_dirty(llama_seq_id seq_id); + bool get_kpool_dirty(const llama_ubatch & ubatch) const; + void clear_kpool_dirty(const llama_ubatch & ubatch) const; // the pooled-key input state kept from one ubatch to the next (llama-kv-cache-kpool.h); made at the first call llama_kpool_views & get_kpool_views() const; @@ -272,7 +291,7 @@ class llama_kv_cache : public llama_memory_i { bool v_trans = true; // the value tensor is transposed // see set_kpool_dirty. mutable: its only consumer runs from set_input, holding a const cache - mutable bool kpool_dirty = false; + mutable llama_kpool_stale kpool_stale; mutable std::shared_ptr kpool_views; diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 2c7e9c8c162a..a9d0459c621d 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -196,20 +196,9 @@ void llama_memory_hybrid::seq_keep(llama_seq_id seq_id) { } void llama_memory_hybrid::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { - // pools group cells by absolute position, so a cached pooled key survives a shift only if - // WHOLE pools move: the shift must be a multiple of kpool AND the range must start and - // end on a pool boundary. both callers pass an arbitrary bound. - if (mem_idx && hparams.indexer_kpool > 0) { - const llama_pos r = (llama_pos) hparams.indexer_kpool; - - // p0 < 0 means "from the start", p1 < 0 "to the end": neither can be straddled - const bool whole_pools = shift % r == 0 && - (p0 <= 0 || p0 % r == 0) && - (p1 < 0 || p1 % r == 0); - - if (!whole_pools) { - mem_attn->set_kpool_dirty(); - } + // both callers pass an arbitrary bound: a shift that regroups pools marks the sequence's pooled keys stale + if (mem_idx && hparams.indexer_kpool > 0 && llama_kpool_stale::shift_regroups(hparams.indexer_kpool, p0, p1, shift)) { + mem_attn->set_kpool_dirty(seq_id); } mem_attn->seq_add(seq_id, p0, p1, shift); if (mem_idx) mem_idx->seq_add(seq_id, p0, p1, shift); @@ -218,7 +207,7 @@ void llama_memory_hybrid::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p void llama_memory_hybrid::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { if (mem_idx && hparams.indexer_kpool > 0 && d != 1) { - mem_attn->set_kpool_dirty(); + mem_attn->set_kpool_dirty(seq_id); } mem_attn->seq_div(seq_id, p0, p1, d); if (mem_idx) mem_idx->seq_div(seq_id, p0, p1, d); diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index 2b9827f254a3..378bc3e67d66 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -422,38 +422,18 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( ggml_tensor * kg_rows = ggml_view_3d(ctx0, kbuf, 2*d_idx, n_kv, n_stream, kbuf->nb[2], kbuf->nb[3], 0); - const int64_t n_new_max = inp_kp->new_pool_cells->ne[0]/r; - - ggml_tensor * members = ggml_get_rows(ctx0, kg_rows, inp_kp->new_pool_cells); - cb(members, "indexer_pool_members", il); - - const size_t nb_mem = members->nb[1]; - - ggml_tensor * mem_k = ggml_view_4d(ctx0, members, d_idx, r, n_new_max, n_stream, - nb_mem, nb_mem*r, members->nb[2], 0); - ggml_tensor * mem_g = ggml_view_4d(ctx0, members, d_idx, r, n_new_max, n_stream, - nb_mem, nb_mem*r, members->nb[2], d_idx*members->nb[0]); - - // r-way softmaxes over the SLOT axis, so it must be dim 0; ape is added PRE-softmax - ggml_tensor * keys_t = ggml_cont(ctx0, ggml_permute(ctx0, mem_k, 1, 0, 2, 3)); - ggml_tensor * gate_t = ggml_cont(ctx0, ggml_permute(ctx0, mem_g, 1, 0, 2, 3)); - ggml_tensor * ape = ggml_cont(ctx0, ggml_transpose(ctx0, layer.indexer_comp_ape)); - gate_t = ggml_add(ctx0, gate_t, ggml_reshape_4d(ctx0, ape, r, d_idx, 1, 1)); - - // soft_max maps ne2/ne3 to gridDim.y/z (65535 cap); fold into ne1 - ggml_tensor * probs = ggml_soft_max(ctx0, ggml_reshape_2d(ctx0, gate_t, r, d_idx*n_new_max*n_stream)); - probs = ggml_reshape_4d(ctx0, probs, r, d_idx, n_new_max, n_stream); - cb(probs, "indexer_pool_probs", il); - - ggml_tensor * pool_new = ggml_sum_rows(ctx0, ggml_mul(ctx0, keys_t, probs)); - pool_new = ggml_reshape_2d(ctx0, pool_new, d_idx, n_new_max*n_stream); - cb(pool_new, "indexer_pool_new", il); // the write is by GLOBAL row while pool_reps is stream-local, so the read must not chain - // off the write tensor; expand first and let build order sequence them - ggml_build_forward_expand(gf, - mctx_idx->cpy_k_part(ctx0, pool_new, inp_kp->new_pool_reps, il, d_idx, 2*d_idx)); + // off the write tensor; every chunk's write is expanded first and build order sequences them + llama_kpool_compress(ctx0, gf, kg_rows, inp_kp->new_pool_cells, inp_kp->new_pool_reps, ape, r, + LLAMA_KPOOL_COMPRESS_CHUNK, + [&](ggml_tensor * src, ggml_tensor * rows) { + return mctx_idx->cpy_k_part(ctx0, src, rows, il, d_idx, 2*d_idx); + }, + [&](ggml_tensor * t, const char * name) { + cb(t, name, il); + }); // every pool's key, gathered into f32 rows (the unfused path, and the fused one under LLAMA_INDEXER_GATHER_LEGACY=1) const auto gather_pool_k = [&]() { @@ -584,7 +564,7 @@ ggml_tensor * llama_model_glm5next::graph::build_dsa_layer( cur = build_attn_sparse(inp_attn, layer.wo, nullptr, nullptr, q, k, k, nullptr, nullptr, layer.wv_b, - top_k, inp_kp->sel_mask, inp_kp->cand_mask, kq_scale, il); + top_k, inp_kp->sel_mask, inp_kp->cand_mask, inp_kp->n_tail, kq_scale, il); } else { cur = build_attn(inp_attn, layer.wo, nullptr, nullptr, diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index fbe9cba0278d..ad3b63eaa8bc 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -166,6 +166,9 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) llama_build_and_test(test-chat.cpp WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}) target_include_directories(test-chat PRIVATE ${PROJECT_SOURCE_DIR}/tools/server) target_link_libraries(test-chat PRIVATE server-context) + llama_build_and_test(test-utf8-replace.cpp) + target_include_directories(test-utf8-replace PRIVATE ${PROJECT_SOURCE_DIR}/tools/server) + target_link_libraries(test-utf8-replace PRIVATE server-context) # TODO: disabled on loongarch64 because the ggml-ci node lacks Python 3.8 if (NOT ${CMAKE_SYSTEM_PROCESSOR} MATCHES "loongarch64") llama_build_and_test(test-json-schema-to-grammar.cpp WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}) diff --git a/tests/test-backend-meta-capture.cpp b/tests/test-backend-meta-capture.cpp index a77eab851347..208a1d32b875 100644 --- a/tests/test-backend-meta-capture.cpp +++ b/tests/test-backend-meta-capture.cpp @@ -230,7 +230,7 @@ int main() { run(0); } - const bool legacy = getenv("GGML_META_CAPTURE_LEGACY") != nullptr; + const bool legacy = ggml_env_switch("GGML_META_CAPTURE_LEGACY"); // as the backend reads it const int executables = legacy ? 0 : (n_graphs + 1)*2; const bool ok_captures = counts.executables == executables && counts.not_captured == 0; printf("%s %d executables made (%d expected), %d graphs not captured\n", ok_captures ? "ok " : "FAIL", diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index ad7565f3a979..7da255d0fcc5 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4514,6 +4514,53 @@ struct test_dsv4_hc_pre_fused : public test_dsv4_hc { } }; +// 16 tokens whose last 12 have streams that cancel in pairs to within 1e-7, at a scale of 100, under equal pre weights +// (scale[0] 0, base[0..3] equal): the mix is ~0. CUDA's Gram front sums pre' G pre from G's entries, not the mix's +// squares, and that rounds under 0 for about a third of such tokens; unclamped, rsqrtf makes their rows NaN. The +// cancelling tokens' outputs are ~1e-3, so the first 4 tokens carry the NMSE and the Gram's rounding of the tiny RMS +// does not count. +struct test_dsv4_hc_pre_cancel : public test_dsv4_hc_pre_fused { + static constexpr int64_t n_plain = 4; + + test_dsv4_hc_pre_cancel() : test_dsv4_hc_pre_fused(GGML_TYPE_BF16, 4096, 16) {} + + std::string vars() override { + return test_dsv4_hc_pre_fused::vars() + ",cancel"; + } + + void initialize_tensors(ggml_context * ctx) override { + test_dsv4_hc_pre_fused::initialize_tensors(ctx); + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + const std::string name = ggml_get_name(t); + if (name == "scale") { + const float zero = 0.0f; + ggml_backend_tensor_set(t, &zero, 0, sizeof(float)); + } else if (name == "base") { + const float base_pre[hc] = { 0.08f, 0.08f, 0.08f, 0.08f }; + ggml_backend_tensor_set(t, base_pre, 0, sizeof(base_pre)); + } else if (name == "x") { + GGML_ASSERT(ggml_is_contiguous(t)); + std::vector x(ggml_nelements(t)); + ggml_backend_tensor_get(t, x.data(), 0, ggml_nbytes(t)); + std::mt19937 rng(tensor_seed(t)); + std::uniform_real_distribution dist(-100.0f, 100.0f); + for (int64_t it = n_plain; it < n_tokens; ++it) { + float * xt = x.data() + it*hc*n_embd; + for (int64_t i = 0; i < n_embd; ++i) { + const float a = dist(rng); + const float b = dist(rng); + xt[i + 0*n_embd] = a; + xt[i + 1*n_embd] = -a*(1.0f + 1e-7f*(i % 7 - 3)); + xt[i + 2*n_embd] = b; + xt[i + 3*n_embd] = -b*(1.0f + 1e-7f*(i % 5 - 2)); + } + } + ggml_backend_tensor_set(t, x.data(), 0, ggml_nbytes(t)); + } + } + } +}; + // The front's normed mix read by n_readers quantized MUL_MATs of 256 rows (the sublayer's projections), their outputs // summed. CUDA's front writes the mix's q8_1 copy beside it where n_embd is a multiple of 512 and a reader runs on // mul_mat_vec_q (ggml_cuda_dsv4_hc_writes_q8_1), for one reader or more, and the readers read the copy. @@ -10281,6 +10328,8 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_dsv4_hc_pre_fused(GGML_TYPE_BF16, 4096, 1, 20, true)); test_cases.emplace_back(new test_dsv4_hc_pre_fused(GGML_TYPE_BF16, 4096, 3, 20, true)); test_cases.emplace_back(new test_dsv4_hc_pre_fused(GGML_TYPE_BF16, 7168, 2, 20)); + // streams whose mix cancels: the Gram front's pre' G pre rounds under 0 and must not make a token NaN (433d88eae) + test_cases.emplace_back(new test_dsv4_hc_pre_cancel()); test_cases.emplace_back(new test_dsv4_hc_pre_fused(GGML_TYPE_BF16, 64, 5, 4)); test_cases.emplace_back(new test_dsv4_hc_pre_fused(GGML_TYPE_F32, 64, 2, 1)); test_cases.emplace_back(new test_dsv4_hc_pre_fused(GGML_TYPE_BF16, 31, 3, 4)); @@ -11550,6 +11599,17 @@ static std::vector> make_test_cases_eval() { for (int n : { 3, 5 }) { test_cases.emplace_back(new test_moe_ffn_chain(GGML_TYPE_IQ3_XXS, 8, 8, 4096, 2048, n, 7.0f)); } + // and the same FFN in Q8_0, which the ring takes past one token (an MTP verify), the down quantizing its vectors + for (int n : { 2, 3 }) { + test_cases.emplace_back(new test_moe_ffn_chain(GGML_TYPE_Q8_0, 16, 8, 4096, 2048, n, 7.0f)); + } + // and 8 tokens on each of 8 experts, where the ring's pair tables are exactly full (64 = MMVQ_MOE_MAX_PAIRS, bit 63 + // of an expert's pair mask). These reach the ring, not its own quantize of y: 64 vectors of 4,096 columns leave the + // shared-memory plan no room for them, so ggml_cuda_mmvq_moe_keeps_y is false and y comes from the shared path. The + // 2- and 3-token cases above are the ones that quantize in the ring, and a doubled block scale there fails those two + // and leaves these passing (RTX 5090, 2026-10-06) + test_cases.emplace_back(new test_moe_ffn_chain(GGML_TYPE_IQ3_XXS, 8, 8, 4096, 2048, 8, 7.0f)); + test_cases.emplace_back(new test_moe_ffn_chain(GGML_TYPE_Q8_0, 8, 8, 4096, 2048, 8, 0.0f)); // and its routed experts at a prefill batch, where MMQ tiles each expert's columns: 64 experts with 8 used, so 64 // and 256 tokens give 8 and 32 columns an expert on average, the tile width follows them and only the experts' // non-empty column tiles are launched @@ -12529,6 +12589,10 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_topk_moe({128, 1, 1, 1}, 128, with_norm, bias_probs, gate, scale_w)); test_cases.emplace_back(new test_topk_moe({129, 1, 1, 1}, 128, with_norm, bias_probs, gate, scale_w)); test_cases.emplace_back(new test_topk_moe({160, 4, 1, 1}, 160, with_norm, bias_probs, gate, scale_w)); + // the rows that fit one block of topk_moe_cuda, at an expert count the fusion takes: every other + // shape here has 1 row or 22, so the fused path's own aliasing rule was never reached (160 is no + // power of two, so ggml_cuda_should_use_topk_moe refuses the shape above) + test_cases.emplace_back(new test_topk_moe({128, 4, 1, 1}, 8, with_norm, bias_probs, gate, scale_w)); test_cases.emplace_back(new test_topk_moe({256, 22, 1, 1}, 6, with_norm, bias_probs, gate, scale_w)); // Used by DeepSeek-V4 test_cases.emplace_back(new test_topk_moe({288, 22, 1, 1}, 8, with_norm, bias_probs, gate, scale_w)); // Used by StepFun 3.7 } @@ -12626,6 +12690,12 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 512, 1, 1, false, true, 1, false, -1, true)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 128, 200, 2, 1, false, true, 4, false, -1, true)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 128, 128, 1, 1, false, true, 9, false, -1, true)); + // KDA with strided q/k/v (the qwen35 views, no cont anywhere), and one chunk of 16 tokens alone: 128 tokens with + // K = 113 snapshot slots leaves the chunked pass a single chunk and the last 112 tokens to the recurrent kernel + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 16, 128, 256, 1, 2, false, true, 1, false, -1, false, /*qkv_view=*/true)); + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 16, 128, 256, 1, 2, false, true, 1, false, -1, true, /*qkv_view=*/true)); + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 128, 128, 1, 1, false, true, /*K=*/113)); + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 128, 128, 2, 1, false, true, 113, false, -1, true)); // gated_delta_net -> cpy into the cache (fused), recurrent and chunked shapes test_cases.emplace_back(new test_gated_delta_net_cache_fusion(GGML_TYPE_F32, 4, 32, 2, 1, 2)); test_cases.emplace_back(new test_gated_delta_net_cache_fusion(GGML_TYPE_F32, 4, 32, 8, 1, 4)); diff --git a/tests/test-kpool-can-reuse.cpp b/tests/test-kpool-can-reuse.cpp index 052c21e6a4b7..bcdcb62e6369 100644 --- a/tests/test-kpool-can-reuse.cpp +++ b/tests/test-kpool-can-reuse.cpp @@ -82,6 +82,7 @@ static llm_graph_input_kpool make_inp(const case_tensors & t, bool scoring, int6 inp.new_pool_reps = t.new_pool_reps; inp.n_new_max = (uint32_t) n_new_max; inp.rebuild = rebuild; + inp.n_tail = KPOOL - 1; (void) scoring; return inp; } @@ -96,6 +97,7 @@ static llm_graph_input_kpool_dims base_dims() { d.n_pools = N_POOLS; d.n_dump = N_DUMP; d.n_new_max = N_NEW_MAX; + d.n_tail = KPOOL - 1; d.rebuild = false; d.scoring = true; return d; @@ -170,6 +172,11 @@ int main() { d.n_dump = N_DUMP - 1; CHECK(!llm_graph_input_kpool::shapes_match(d, inp), "a dump pool count change must refuse"); + // the sparse attention's n_kv_max is sized on the longest tail + d = base_dims(); + d.n_tail = KPOOL + 8; + CHECK(!llm_graph_input_kpool::shapes_match(d, inp), "a longer tail must refuse"); + // non-scoring stored graph asked against a scoring step must refuse case_tensors ts = make_tensors(ctx, false, 0); llm_graph_input_kpool inps = make_inp(ts, false, 0, false); diff --git a/tests/test-kpool-input.cpp b/tests/test-kpool-input.cpp index f6c33feb0a79..8e83e73c7bce 100644 --- a/tests/test-kpool-input.cpp +++ b/tests/test-kpool-input.cpp @@ -15,16 +15,20 @@ // is checked against a snapshot of the cells on its own. #include "../src/llama-batch.h" +#include "../src/llama-kv-cache.h" #include "../src/llama-kv-cache-kpool.h" #include "../src/llama-kv-cells.h" #include "ggml.h" +#include "ggml-alloc.h" #include "ggml-backend.h" +#include "ggml-cpu.h" #include #include #include #include +#include #include #include #include @@ -512,6 +516,30 @@ static bool dump_safe(const maps & m, int64_t n_kv, uint32_t r, char * what, siz return true; } +// the most finite columns a row of sel_mask holds among its first n_kv: the cells its tail grants +static int64_t sel_tail_max(const maps & m, int64_t n_kv) { + const int64_t n_sel = m.sel_mask->ne[0]; + const int64_t n_rows = m.sel_mask->ne[1]*m.sel_mask->ne[3]; + + int64_t n_max = 0; + + for (int64_t row = 0; row < n_rows; ++row) { + int64_t n = 0; + for (int64_t j = 0; j < n_kv; ++j) { + const int64_t k = row*n_sel + j; + const float x = m.sel_mask->type == GGML_TYPE_F16 ? ggml_fp16_to_fp32(((const ggml_fp16_t *) m.sel_mask->data)[k]) : + ((const float *) m.sel_mask->data)[k]; + n += std::isfinite(x); + } + n_max = std::max(n_max, n); + } + + return n_max; +} + +// builds whose longest tail held more than r - 1 cells (two cells at one position) +static int64_t n_tail_long = 0; + // the inputs of one ubatch built from the views and from the cells. The ubatch is placed in the cells first, as the cache // does (apply_ubatch before the inputs are set). rebuild: re-emit every pool, as after a position mutation static bool build_both(pool & w, llama_kpool_views & views, const config & cfg, bool rebuild, std::mt19937 & rng, control ctl, char * what, size_t n_what) { @@ -604,11 +632,23 @@ static bool build_both(pool & w, llama_kpool_views & views, const config & cfg, maps inc(n_kv, n_ns, n_tokens/n_ns, n_ps, cfg.r, cfg.kcache, cfg.f16, rebuild, n_dump, 0x5A); // the views go first: the reference never takes the log - llama_kpool_set_input(cells_of, &views, nullptr, inc.pool_cells, nullptr, inc.pool_bias, inc.sel_mask, inc.cand_mask, + const int64_t tail_inc = llama_kpool_set_input(cells_of, &views, nullptr, inc.pool_cells, nullptr, inc.pool_bias, inc.sel_mask, inc.cand_mask, inc.pool_reps, inc.new_pool_cells, inc.new_pool_reps, strm_of.data(), w.v_cells[0].size(), rebuild, ub, cfg.r); - llama_kpool_set_input(cells_of, nullptr, nullptr, ref.pool_cells, nullptr, ref.pool_bias, ref.sel_mask, ref.cand_mask, + const int64_t tail_ref = llama_kpool_set_input(cells_of, nullptr, nullptr, ref.pool_cells, nullptr, ref.pool_bias, ref.sel_mask, ref.cand_mask, ref.pool_reps, ref.new_pool_cells, ref.new_pool_reps, strm_of.data(), w.v_cells[0].size(), rebuild, ub, cfg.r); + // the tail a row is granted, as set_input reports it and as sel_mask holds it, within the bound the graph sizes on + const int64_t tail_sel = sel_tail_max(ref, n_kv); + const int64_t tail_max = llama_kpool_tail_max(cells_of, *ub, cfg.r); + + n_tail_long += tail_ref > (int64_t) cfg.r - 1; + + if (tail_inc != tail_ref || tail_ref != tail_sel || tail_ref > tail_max) { + snprintf(what, n_what, "tail: views report %lld cells, the cells %lld, sel_mask holds %lld, the bound is %lld", + (long long) tail_inc, (long long) tail_ref, (long long) tail_sel, (long long) tail_max); + return false; + } + if (ctl == CONTROL_CORRUPT) { ((uint8_t *) inc.sel_mask->data)[0] ^= 0x80; } @@ -749,9 +789,289 @@ static void run_controls(std::mt19937 & rng) { CHECK(views.get_stats().n_rebuilt > n_rebuilt, "a view whose changes were taken was not rebuilt"); } +// The stale marks of the pooled keys across two sequences, as a server makes them: sequence 0 shifted by a context shift +// (seq_rm of [4, 9), then seq_add of [9, end) by -5), a decode of sequence 1, a decode of sequence 0. The pooled keys are +// kept as the indexer cache keeps them, one row per complete pool written by the new-pool slots; after the last decode +// every complete pool of both sequences must hold the key of its members. `per_seq` false is the cache's former scheme, one +// flag for the cache, which the decode of sequence 1 clears: the count of stale pools it leaves shows the check can fail. +// Returns the count of complete pools whose key is not their members' after the last decode +static int64_t stale_after_shift(std::mt19937 & rng, bool per_seq) { + const uint32_t r = 4; + const uint32_t size = 256; + + pool w(size, 2, true, rng); + + llama_kpool_stale stale; // the per-sequence marks + bool flag = false; // the former flag + + // the indexer cache: rep row -> the members the key in it was made from + std::map> key; + + const llama_kpool_cells_fn cells_of = [&](llama_seq_id s) -> const llama_kv_cells & { return w.cells_of(s); }; + const uint32_t strm_of[1] = { 0 }; + + auto decode = [&](llama_seq_id s, uint32_t n) { + ubatch_data ud; + for (const llama_pos p : w.append(s, n)) { + ud.add(s, p); + } + const llama_ubatch * ub = ud.get(); + + const bool rebuild = per_seq ? stale.any(*ub) : flag; + + const int64_t n_kv = std::max(32, (w.v_cells[0].used_max_p1() + 31)/32*32); + const int64_t n_tps = ub->n_tokens; + + maps m(n_kv, 1, n_tps, 1, r, true, true, rebuild, 0, 0); + + llama_kpool_set_input(cells_of, nullptr, nullptr, m.pool_cells, nullptr, m.pool_bias, m.sel_mask, m.cand_mask, + m.pool_reps, m.new_pool_cells, m.new_pool_reps, strm_of, size, rebuild, ub, r); + + // set_rows: every slot writes its row, a real one the key of its members + const int64_t n_new_max = m.new_pool_cells->ne[0]/r; + const int32_t * cells_in = (const int32_t *) m.new_pool_cells->data; + const int64_t * rows = (const int64_t *) m.new_pool_reps->data; + + for (int64_t k = 0; k < n_new_max; ++k) { + key[rows[k]] = std::vector(cells_in + k*r, cells_in + (k + 1)*r); + } + + if (per_seq) { + if (rebuild) { + stale.clear(*ub); + } + } else { + flag = false; + } + + return rebuild; + }; + + // the complete pools of a sequence whose row does not hold the key of their members + auto n_stale = [&](llama_seq_id s) { + const auto & cells = w.cells_of(s); + + std::map> pools; + for (uint32_t i = 0; i < cells.size(); ++i) { + if (!cells.is_empty(i) && cells.seq_has(i, s)) { + auto & m = pools[cells.pos_get(i)/r]; + m.resize(r, -1); + m[cells.pos_get(i) % r] = (int32_t) i; + } + } + + int64_t n = 0; + for (const auto & [b, members] : pools) { + if (std::find(members.begin(), members.end(), -1) != members.end()) { + continue; + } + const auto it = key.find(members[r - 1]); + n += it == key.end() || it->second != members; + } + return n; + }; + + decode(0, 16); + decode(1, 16); + CHECK(n_stale(0) == 0 && n_stale(1) == 0, "the prefills left %lld and %lld stale pools", (long long) n_stale(0), (long long) n_stale(1)); + + // the context shift of sequence 0, as llama_memory_hybrid::seq_rm and seq_add apply it + w.rollback(0, 4, 9); + auto & cells = w.cells_of(0); + for (uint32_t i = 0; i < cells.size(); ++i) { + if (cells.pos_in(i, 9, INT32_MAX) && cells.seq_has(i, 0)) { + cells.pos_add(i, -5); + } + } + cells.reset_shift(); + + CHECK(llama_kpool_stale::shift_regroups(r, 9, -1, -5), "a shift by -5 over pools of 4 must regroup them"); + stale.mark(0); + flag = true; + + const bool rebuilt_1 = decode(1, 1); + const bool rebuilt_0 = decode(0, 1); + + if (per_seq) { + CHECK(!rebuilt_1, "a decode of sequence 1 rebuilt for sequence 0's shift"); + CHECK(rebuilt_0, "the decode of the shifted sequence 0 did not rebuild"); + + // and the marks are spent: the next decode of sequence 0 is a plain one + ubatch_data ud; + ud.add(0, w.p_next(0)); + CHECK(!stale.any(*ud.get()), "sequence 0 is still marked after the decode that rebuilt it"); + } + + return n_stale(0) + n_stale(1); +} + +// The new pools' keys compressed `chunk` at a time (llama_kpool_compress) against all at once, as a rebuild re-emits every +// pool: the same bytes in every pooled row, from a compute buffer of a chunk's size. Returns the buffer's bytes +static size_t compress_pools(int64_t chunk, std::vector & pooled) { + const int64_t d_idx = 32; + const int64_t r = 4; + const int64_t n_kv = 16384; + const int64_t n_pools = n_kv/r; + + ggml_backend_t cpu = ggml_backend_cpu_init(); + + ggml_init_params dp = { 8*ggml_tensor_overhead(), nullptr, true }; + ggml_context * cd = ggml_init(dp); + + ggml_tensor * kbuf = ggml_new_tensor_4d(cd, GGML_TYPE_F32, d_idx, 3, n_kv, 1); // a cell: key, gate, pooled key + ggml_tensor * cells = ggml_new_tensor_2d(cd, GGML_TYPE_I32, r*n_pools, 1); + ggml_tensor * reps = ggml_new_tensor_1d(cd, GGML_TYPE_I64, n_pools); + ggml_tensor * ape = ggml_new_tensor_2d(cd, GGML_TYPE_F32, r, d_idx); + + ggml_backend_buffer_t data = ggml_backend_alloc_ctx_tensors(cd, cpu); + + std::mt19937 rng(20261006); + std::uniform_real_distribution u(-2.0f, 2.0f); + + std::vector kv(ggml_nelements(kbuf)); + for (int64_t j = 0; j < n_kv; ++j) { + for (int64_t h = 0; h < 3*d_idx; ++h) { + kv[j*3*d_idx + h] = h < 2*d_idx ? u(rng) : 0.0f; + } + } + ggml_backend_tensor_set(kbuf, kv.data(), 0, ggml_nbytes(kbuf)); + + // a pool's members are r cells anywhere in the cache, its key in the row of the last + std::vector perm(n_kv); + for (int64_t j = 0; j < n_kv; ++j) { + perm[j] = (int32_t) j; + } + std::shuffle(perm.begin(), perm.end(), rng); + std::vector rows(n_pools); + for (int64_t p = 0; p < n_pools; ++p) { + rows[p] = perm[p*r + r - 1]; + } + ggml_backend_tensor_set(cells, perm.data(), 0, ggml_nbytes(cells)); + ggml_backend_tensor_set(reps, rows.data(), 0, ggml_nbytes(reps)); + + std::vector a(r*d_idx); + for (auto & x : a) { + x = u(rng); + } + ggml_backend_tensor_set(ape, a.data(), 0, ggml_nbytes(ape)); + + ggml_init_params gp = { 4096*ggml_tensor_overhead() + ggml_graph_overhead_custom(4096, false), nullptr, true }; + ggml_context * cg = ggml_init(gp); + ggml_cgraph * gf = ggml_new_graph_custom(cg, 4096, false); + + ggml_tensor * kg_rows = ggml_view_3d(cg, kbuf, 2*d_idx, n_kv, 1, kbuf->nb[2], kbuf->nb[3], 0); + + llama_kpool_compress(cg, gf, kg_rows, cells, reps, ape, r, chunk, + [&](ggml_tensor * src, ggml_tensor * idx) { + ggml_tensor * k2 = ggml_reshape_2d(cg, kbuf, 3*d_idx, n_kv); + ggml_tensor * dst = ggml_view_2d(cg, k2, d_idx, n_kv, k2->nb[1], ggml_row_size(k2->type, 2*d_idx)); + return ggml_set_rows(cg, dst, src, idx); + }, + [](ggml_tensor * t, const char * name) { + ggml_set_name(t, name); + }); + + ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_cpu_buffer_type()); + CHECK(ggml_gallocr_alloc_graph(galloc, gf), "the compression graph did not allocate"); + const size_t size = ggml_gallocr_get_buffer_size(galloc, 0); + + ggml_backend_graph_compute(cpu, gf); + + ggml_backend_tensor_get(kbuf, kv.data(), 0, ggml_nbytes(kbuf)); + pooled.assign(n_pools*d_idx, 0.0f); + for (int64_t p = 0; p < n_pools; ++p) { + std::copy_n(kv.data() + rows[p]*3*d_idx + 2*d_idx, d_idx, pooled.data() + p*d_idx); + } + + ggml_gallocr_free(galloc); + ggml_free(cg); + ggml_backend_buffer_free(data); + ggml_free(cd); + ggml_backend_free(cpu); + + return size; +} + +static void run_compress_chunks() { + std::vector whole; + std::vector chunked; + + const size_t size_whole = compress_pools(4096, whole); // every pool at once: the graph before chunks + const size_t size_chunked = compress_pools(512, chunked); + + printf("compressing 4,096 pools: %zu bytes of compute buffer at once, %zu in chunks of 512\n", size_whole, size_chunked); + + CHECK(whole.size() == chunked.size() && memcmp(whole.data(), chunked.data(), whole.size()*sizeof(float)) == 0, + "the pooled keys compressed in chunks differ from all at once"); + CHECK(std::any_of(whole.begin(), whole.end(), [](float x) { return x != 0.0f; }), "no pooled key was written"); + CHECK(size_chunked*4 < size_whole, "chunks of 512 of 4,096 pools took %zu bytes against %zu at once", size_chunked, size_whole); +} + +// which shifts regroup pools, as llama_memory_hybrid::seq_add asks +static void run_shift_regroups() { + CHECK(!llama_kpool_stale::shift_regroups(4, 8, -1, -4), "whole pools moved by a whole pool"); + CHECK(!llama_kpool_stale::shift_regroups(4, -1, -1, 8), "the whole sequence moved by two pools"); + CHECK( llama_kpool_stale::shift_regroups(4, 9, -1, -5), "a shift by -5"); + CHECK( llama_kpool_stale::shift_regroups(4, 9, -1, -4), "a range that starts inside a pool"); + CHECK( llama_kpool_stale::shift_regroups(4, 8, 10, -4), "a range that ends inside a pool"); +} + +// a sequence holding 11 more cells at each of its last three positions: the row at the last has a tail of 36 cells, and with +// GLM-5-Next's 2,048 top-k cells (512 pools of 4) its live cells pass the n_kv_max an r - 1 tail sizes +static void run_tail_bound(std::mt19937 & rng) { + const uint32_t r = 4; + + pool w(256, 1, true, rng); + w.append(0, 63); + + auto & cells = w.cells_of(0); + for (llama_pos p = 60; p <= 62; ++p) { + for (int k = 0; k < 11; ++k) { + const int64_t i = w.find_empty(0); + cells.pos_set(i, p); + cells.seq_add(i, 0); + } + } + + ubatch_data ud; + ud.add(0, 62); + const llama_ubatch * ub = ud.get(); + + const llama_kpool_cells_fn cells_of = [&](llama_seq_id s) -> const llama_kv_cells & { return w.cells_of(s); }; + + maps m(256, 1, 1, 1, r, false, true, false, 0, 0); + const int64_t n_tail = llama_kpool_set_input(cells_of, nullptr, nullptr, m.pool_cells, nullptr, m.pool_bias, m.sel_mask, m.cand_mask, + nullptr, nullptr, nullptr, nullptr, 0, false, ub, r); + const int64_t n_bound = llama_kpool_tail_max(cells_of, *ub, r); + + const int64_t n_top_k = 2048; + const int64_t n_live = n_top_k + n_tail; + const int64_t n_kv_old = llama_kpool_n_kv_max(n_top_k, r - 1); + const int64_t n_kv_new = llama_kpool_n_kv_max(n_top_k, n_bound); + + printf("a tail of %lld cells: %lld live cells a row, n_kv_max %lld for an r - 1 tail, %lld for the longest tail\n", + (long long) n_tail, (long long) n_live, (long long) n_kv_old, (long long) n_kv_new); + + CHECK(n_tail == 36 && n_bound == 36, "the tail: %lld cells granted, %lld counted, 36 placed", (long long) n_tail, (long long) n_bound); + CHECK(n_kv_old < n_live, "an r - 1 tail bound holds %lld live cells: the case cannot fail", (long long) n_live); + CHECK(n_kv_new >= n_live, "n_kv_max %lld under %lld live cells", (long long) n_kv_new, (long long) n_live); +} + int main() { std::mt19937 rng(20260930); + run_shift_regroups(); + run_compress_chunks(); + run_tail_bound(rng); + + const int64_t n_stale_per_seq = stale_after_shift(rng, true); + const int64_t n_stale_flag = stale_after_shift(rng, false); + + printf("a context shift of one sequence: %lld stale pools with per-sequence marks, %lld with one flag for the cache\n", + (long long) n_stale_per_seq, (long long) n_stale_flag); + CHECK(n_stale_per_seq == 0, "%lld pools kept a stale key after the shifted sequence's own decode", (long long) n_stale_per_seq); + CHECK(n_stale_flag > 0, "one flag for the cache left no pool stale: the check cannot fail"); + llama_kpool_views::stats total; const config configs[] = { @@ -786,6 +1106,9 @@ int main() { CHECK(total.n_served > 2000, "only %llu streams were served from views", (unsigned long long) total.n_served); CHECK(total.n_direct > 100, "only %llu streams took the maps from the cells: the fallbacks are not reached", (unsigned long long) total.n_direct); + printf("random rounds: %lld builds with a tail over r - 1 cells, each within llama_kpool_tail_max\n", (long long) n_tail_long); + CHECK(n_tail_long > 0, "no build held a tail over r - 1 cells: the bound is not tested"); + if (n_fail == 0) { printf("ok\n"); } diff --git a/tests/test-utf8-replace.cpp b/tests/test-utf8-replace.cpp new file mode 100644 index 000000000000..e7060a4b86eb --- /dev/null +++ b/tests/test-utf8-replace.cpp @@ -0,0 +1,74 @@ +// utf8_replace_malformed (tools/server/server-common.cpp): the server's repair of generated text before the chat parser +// reads it. The parser's JSON dump is strict, so a sequence this lets through as well-formed answers the request 500. +// Each case is the input, the `from` it is read from, the text it must leave and where it says the incomplete tail starts. + +#include "server-common.h" + +#include +#include + +static int failures = 0; + +static std::string show(const std::string & s) { + std::string out; + char buf[8]; + for (unsigned char c : s) { + snprintf(buf, sizeof(buf), c >= 0x20 && c < 0x7F ? "%c" : "\\x%02X", c); + out += buf; + } + return out; +} + +static void check(const char * name, const std::string & in, size_t from, const std::string & want, size_t want_ret) { + std::string text = in; + const size_t ret = utf8_replace_malformed(text, from); + if (text != want || ret != want_ret) { + fprintf(stderr, "FAIL %s: \"%s\" -> \"%s\" (%zu), want \"%s\" (%zu)\n", name, show(in).c_str(), + show(text).c_str(), ret, show(want).c_str(), want_ret); + failures++; + } else { + printf("ok %s\n", name); + } +} + +int main() { + const std::string R = "\xEF\xBF\xBD"; // U+FFFD + + // well-formed text is left alone, every length and range edge included + const std::string good = "a\xC2\x80\xDF\xBF\xE0\xA0\x80\xED\x9F\xBF\xEE\x80\x80\xF0\x90\x80\x80\xF4\x8F\xBF\xBF"; + check("well-formed", good, 0, good, good.size()); + + // the shapes the pattern check already caught + check("stray continuation", "a\x80" "b", 0, "a" + R + "b", 3 + 2); + check("lead then ascii", "\xC3" "A", 0, R + "A", 4); + check("lead then lead", "\xE4\xE4\xE4", 0, R + R + "\xE4", 6); + + // what a pattern check let through as well-formed, each one U+FFFD for its maximal subpart (Unicode 3.9, U+FFFD + // substitution of maximal subparts): a byte that leads nothing, or a second byte outside its lead's range + check("surrogate", "\xED\xA0\x80", 0, R + R + R, 9); + check("overlong 2", "\xC0\x80", 0, R + R, 6); + check("overlong 3", "\xE0\x80\x80", 0, R + R + R, 9); + check("overlong 4", "\xF0\x80\x80\x80", 0, R + R + R + R, 12); + check("past U+10FFFF", "\xF4\x90\x80\x80", 0, R + R + R + R, 12); + check("lead F5", "\xF5\x80\x80\x80", 0, R + R + R + R, 12); + check("lead FF", "x\xFF" "y", 0, "x" + R + "y", 5); + + // a maximal subpart is one U+FFFD, and the byte that ended it is read afresh + check("subpart of two", "\xE4\xB8" "A", 0, R + "A", 4); + check("subpart of three", "\xF0\x9F\x98" "A", 0, R + "A", 4); + + // an incomplete tail is kept back only while every byte of it is right so far + check("tail of a 3-byte", "ok\xE4\xB8", 0, "ok\xE4\xB8", 2); + check("tail E0 A0", "\xE0\xA0", 0, "\xE0\xA0", 0); + check("not a tail: E0 80", "\xE0\x80", 0, R + R, 6); + check("not a tail: ED A0", "\xED\xA0", 0, R + R, 6); + + // `from` reads only the unsent part; what is before it is not looked at again + check("from", "\x80" "ab\x80", 3, "\x80" "ab" + R, 6); + + if (failures > 0) { + fprintf(stderr, "%d case(s) failed\n", failures); + return 1; + } + return 0; +} diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index a90f0177c703..9ad295146fc1 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -926,14 +926,36 @@ size_t utf8_replace_malformed(std::string & text, size_t from) { size_t i = std::min(from, text.size()); while (i < text.size()) { const unsigned char c = byte(i); - const size_t len = c < 0x80 ? 1 : (c & 0xE0) == 0xC0 ? 2 : (c & 0xF0) == 0xE0 ? 3 : (c & 0xF8) == 0xF0 ? 4 : 0; + // a lead's length and the range its second byte must fall in (Unicode Table 3-7): E0, ED, F0 and F4 narrow it, + // which is what keeps out an overlong form, a surrogate and a code point past U+10FFFF; 80-C1 and F5-FF lead + // nothing. A pattern-only check kept those, and the chat parser's strict JSON dump then threw on them. + size_t len = 0; + unsigned char lo = 0x80; + unsigned char hi = 0xBF; + if (c < 0x80) { + len = 1; + } else if (c >= 0xC2 && c <= 0xDF) { + len = 2; + } else if (c >= 0xE0 && c <= 0xEF) { + len = 3; + lo = c == 0xE0 ? 0xA0 : 0x80; + hi = c == 0xED ? 0x9F : 0xBF; + } else if (c >= 0xF0 && c <= 0xF4) { + len = 4; + lo = c == 0xF0 ? 0x90 : 0x80; + hi = c == 0xF4 ? 0x8F : 0xBF; + } size_t n = 1; - while (n < len && i + n < text.size() && (byte(i + n) & 0xC0) == 0x80) { + while (n < len && i + n < text.size()) { + const unsigned char b = byte(i + n); + if (b < (n == 1 ? lo : 0x80) || b > (n == 1 ? hi : 0xBF)) { + break; + } n++; } const bool whole = len > 0 && n == len; if (!whole && len > 0 && i + n == text.size()) { - break; // an incomplete tail + break; // an incomplete tail: every byte so far is right, and the next token may complete it } if (!whole && first == std::string::npos) { first = i; @@ -945,7 +967,9 @@ size_t utf8_replace_malformed(std::string & text, size_t from) { fixed.append("\xEF\xBF\xBD"); } } - i += whole ? len : 1; // a malformed lead stands alone: the byte after it is read afresh + // a malformed sequence is replaced by one U+FFFD for its maximal subpart (the lead and the continuation bytes it + // took before the one that cannot continue it), as the JSON writer does; that byte is then read afresh + i += whole ? len : (len > 0 ? n : 1); } if (first == std::string::npos) { return i; diff --git a/tools/server/server-common.h b/tools/server/server-common.h index 20b532997999..a6a6d86348e0 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -270,10 +270,12 @@ llama_tokens tokenize_mixed(const llama_vocab * vocab, const json & json_prompt, // if validate_utf8(text) == text.size(), then the whole text is valid utf8 size_t validate_utf8(const std::string& text); -// replaces each malformed byte of text from `from` on (a continuation byte with no lead, a lead byte not followed by -// its continuation bytes, a byte that leads nothing) with U+FFFD, as the JSON writer shows it, and stops before an -// incomplete tail: a lead whose continuation bytes are right so far but too few, which the next token may complete. -// Returns where that tail starts (text.size() if there is none); text[from, return) is then well-formed UTF-8. +// replaces each malformed sequence of text from `from` on with one U+FFFD for its maximal subpart, as the JSON writer +// shows it: a continuation byte with no lead, a byte that leads nothing (80-C1, F5-FF), and a lead whose next byte +// cannot continue it, judged by Unicode's Table 3-7, so an overlong form, a surrogate and a code point past U+10FFFF +// are malformed too. It stops before an incomplete tail: a lead whose continuation bytes are right so far but too +// few, which the next token may complete. Returns where that tail starts (text.size() if there is none); +// text[from, return) is then well-formed UTF-8. size_t utf8_replace_malformed(std::string & text, size_t from); // process mtmd prompt, return the server_tokens containing both text tokens and media chunks diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 8d0a28b222ff..c712665b411b 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -241,7 +241,7 @@ struct server_slot { mtmd::batch_ptr mbatch = nullptr; // speculative decoding - common_speculative * spec; + common_speculative * spec = nullptr; // the context's, once init ran llama_tokens spec_draft; llama_tokens spec_prompt; @@ -323,12 +323,14 @@ struct server_slot { if (ctx_dft) { llama_state_seq_get_data_ext(ctx_dft, cur->data.drft.data(), cur_size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE); } + // the drafter's kept row travels with the prompt, as the two states above do + common_speculative_get_state(spec, id, cur->data.spec); return true; } bool prompt_load(server_prompt_cache & prompt_cache, const server_tokens & tokens) { - bool res = prompt_cache.load(prompt, tokens, ctx_tgt, ctx_dft, id); + bool res = prompt_cache.load(prompt, tokens, ctx_tgt, ctx_dft, spec, id); if (!res) { SLT_WRN(*this, "%s", "failed to load prompt from cache\n"); } @@ -340,6 +342,7 @@ struct server_slot { SLT_TRC(*this, "clearing prompt with %zu tokens\n", prompt.tokens.size()); mem.seq_rm(id, -1, -1); + common_speculative_set_state(spec, id, {}); // the cells a kept row belongs to are gone if (!prompt.checkpoints.empty()) { checkpoint_spare(prompt.checkpoints.back()); @@ -1038,6 +1041,9 @@ struct server_context_impl { int64_t t_last_load_progress_ms = 0; void destroy() { + for (server_slot & slot : slots) { + slot.spec = nullptr; // the slots outlive it: sleeping destroys and load_model builds both again + } spec.reset(); spec_init.reset(); diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 0d3beb313cea..e20396969537 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1784,13 +1784,15 @@ server_prompt_cache_state * server_prompt_cache::alloc(const server_prompt & pro /*.data =*/ { /*.main =*/ std::move(state_data_tgt), /*.drft =*/ std::move(state_data_dft), + /*.spec =*/ {}, // the caller fills it, as it does the two above }, }); return &states.back(); } -bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot) { +bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx_tgt, llama_context * ctx_dft, + common_speculative * spec, int32_t id_slot) { const int lcp_best = prompt.tokens.get_common_prefix(tokens_new); float f_keep_best = prompt.tokens.size() > 0 ? float(lcp_best) / prompt.tokens.size() : -1.0f; // empty slot: any cache entry wins @@ -1859,6 +1861,9 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok } } + // this prompt's own drafter state, or none: a slot never drafts from the row the conversation before it left + common_speculative_set_state(spec, id_slot, it_best->data.spec); + prompt = std::move(it_best->prompt); states.erase(it_best); diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 9c99143f8e19..c0c6cbd0449d 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -11,6 +11,8 @@ // TODO: prevent including the whole server-common.h as we only use server_tokens #include "server-common.h" +struct common_speculative; + enum server_task_type { SERVER_TASK_TYPE_COMPLETION, @@ -588,9 +590,12 @@ struct server_prompt { struct server_prompt_data { std::vector main; std::vector drft; + // the drafter's own per-sequence state, which the two above do not hold: the MTP head's kept h row and the position + // it is for, so a prompt restored into a slot drafts from the row its own generation left (common_speculative_get_state) + std::vector spec; size_t size() const { - return main.size() + drft.size(); + return main.size() + drft.size() + spec.size(); } }; @@ -629,7 +634,8 @@ struct server_prompt_cache { server_prompt_cache_state * alloc(const server_prompt & prompt, size_t state_size_main, size_t state_size_drft); - bool load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot); + bool load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx_tgt, llama_context * ctx_dft, + common_speculative * spec, int32_t id_slot); void update(); };